{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE DataKinds #-}
module ReWire.Vectors where

import ReWire
import qualified ReWire.Finite as F
import qualified ReWire.FiniteComp as FC
import Prelude hiding ((++), zip, take, drop, map, zipWith)

{-# INLINE replicate #-}
replicate :: KnownNat n => a -> Vec n a
replicate :: forall (n :: Nat) a. KnownNat n => a -> Vec n a
replicate = a -> Vec n a
forall (n :: Nat) a. KnownNat n => a -> Vec n a
rwPrimVecReplicate

{-# INLINE reverse #-}
reverse :: Vec n a -> Vec n a
reverse :: forall (n :: Nat) a. Vec n a -> Vec n a
reverse = Vec n a -> Vec n a
forall (n :: Nat) a. Vec n a -> Vec n a
rwPrimVecReverse

{-# INLINE slice #-}
slice :: (KnownNat i, KnownNat n) => Proxy i -> Vec ((i + n) + m) a -> Vec n a
slice :: forall (i :: Nat) (n :: Nat) (m :: Nat) a.
(KnownNat i, KnownNat n) =>
Proxy i -> Vec ((i + n) + m) a -> Vec n a
slice = Proxy i -> Vec ((i + n) + m) a -> Vec n a
forall (i :: Nat) (n :: Nat) (m :: Nat) a.
(KnownNat i, KnownNat n) =>
Proxy i -> Vec ((i + n) + m) a -> Vec n a
rwPrimVecSlice

{-# INLINE rslice #-}
rslice :: (KnownNat i, KnownNat n) => Proxy i -> Vec ((i + n) + m) a -> Vec n a
rslice :: forall (i :: Nat) (n :: Nat) (m :: Nat) a.
(KnownNat i, KnownNat n) =>
Proxy i -> Vec ((i + n) + m) a -> Vec n a
rslice = Proxy i -> Vec ((i + n) + m) a -> Vec n a
forall (i :: Nat) (n :: Nat) (m :: Nat) a.
(KnownNat i, KnownNat n) =>
Proxy i -> Vec ((i + n) + m) a -> Vec n a
rwPrimVecRSlice

{-# INLINE index #-}
index :: Vec n a -> Finite n -> a
index :: forall (n :: Nat) a. Vec n a -> Finite n -> a
index = Vec n a -> Finite n -> a
forall (n :: Nat) a. Vec n a -> Finite n -> a
rwPrimVecIndex

{-# INLINE index' #-}
index' :: KnownNat n => Vec ((n + m) + 1) a -> Proxy n -> a
index' :: forall (n :: Nat) (m :: Nat) a.
KnownNat n =>
Vec ((n + m) + 1) a -> Proxy n -> a
index' = Vec ((n + m) + 1) a -> Proxy n -> a
forall (n :: Nat) (m :: Nat) a.
KnownNat n =>
Vec ((n + m) + 1) a -> Proxy n -> a
rwPrimVecIndexProxy

{-# INLINE (++) #-}
(++) :: Vec n a -> Vec m a -> Vec (n + m) a
++ :: forall (n :: Nat) a (m :: Nat). Vec n a -> Vec m a -> Vec (n + m) a
(++) = Vec n a -> Vec m a -> Vec (n + m) a
forall (n :: Nat) a (m :: Nat). Vec n a -> Vec m a -> Vec (n + m) a
rwPrimVecConcat

{-# INLINE empty #-}
empty :: Vec 0 a
empty :: forall a. Vec 0 a
empty = [a] -> Vec 0 a
forall (n :: Nat) a. KnownNat n => [a] -> Vec n a
fromList []

{-# INLINE singleton #-}
singleton :: a -> Vec 1 a
singleton :: forall a. a -> Vec 1 a
singleton a
a = [a] -> Vec 1 a
forall (n :: Nat) a. KnownNat n => [a] -> Vec n a
fromList [a
a]

{-# INLINE cons #-}
cons :: a -> Vec n a -> Vec (1 + n) a
cons :: forall a (n :: Nat). a -> Vec n a -> Vec (1 + n) a
cons a
x Vec n a
v = [a] -> Vec 1 a
forall (n :: Nat) a. KnownNat n => [a] -> Vec n a
fromList [a
x] Vec 1 a -> Vec n a -> Vector Vector (1 + n) a
forall (n :: Nat) a (m :: Nat). Vec n a -> Vec m a -> Vec (n + m) a
++ Vec n a
v

{-# INLINE snoc #-}
snoc :: Vec n a -> a -> Vec (n + 1) a
snoc :: forall (n :: Nat) a. Vec n a -> a -> Vec (n + 1) a
snoc Vec n a
v a
x = Vec n a
v Vec n a -> Vec 1 a -> Vector Vector (n + 1) a
forall (n :: Nat) a (m :: Nat). Vec n a -> Vec m a -> Vec (n + m) a
++ [a] -> Vec 1 a
forall (n :: Nat) a. KnownNat n => [a] -> Vec n a
fromList [a
x]

{-# INLINE head #-}
head :: Vec (1 + n) a -> a
head :: forall (n :: Nat) a. Vec (1 + n) a -> a
head Vec (1 + n) a
v = Vec ((0 + n) + 1) a -> Proxy 0 -> a
forall (n :: Nat) (m :: Nat) a.
KnownNat n =>
Vec ((n + m) + 1) a -> Proxy n -> a
index' Vec (1 + n) a
Vec ((0 + n) + 1) a
v (Proxy 0
forall (n :: Nat). Proxy n
Proxy :: Proxy 0)

{-# INLINE last #-}
last :: KnownNat n => Vec n a -> a
last :: forall (n :: Nat) a. KnownNat n => Vec n a -> a
last Vec n a
v = Vec n a -> Finite n -> a
forall (n :: Nat) a. Vec n a -> Finite n -> a
index Vec n a
v (Vec n a -> Finite n
forall (n :: Nat) a. KnownNat n => Vec n a -> Finite n
lastIndex Vec n a
v)

{-# INLINE lastIndex #-}
lastIndex :: KnownNat n => Vec n a -> Finite n
lastIndex :: forall (n :: Nat) a. KnownNat n => Vec n a -> Finite n
lastIndex Vec n a
_ = Finite n
forall (n :: Nat). KnownNat n => Finite n
F.maxBound

-- {-# INLINE lastIndex' #-}
-- lastIndex' :: KnownNat n => Vec (1 + n) a -> Proxy n
-- lastIndex' _ = Proxy

{-# INLINE take #-}
take :: KnownNat n => Vec (n + m) a -> Vec n a
take :: forall (n :: Nat) (m :: Nat) a.
KnownNat n =>
Vec (n + m) a -> Vec n a
take = Proxy 0 -> Vec ((0 + n) + m) a -> Vector Vector n a
forall (i :: Nat) (n :: Nat) (m :: Nat) a.
(KnownNat i, KnownNat n) =>
Proxy i -> Vec ((i + n) + m) a -> Vec n a
slice (Proxy 0
forall (n :: Nat). Proxy n
Proxy :: Proxy 0)

{-# INLINE init #-}
init :: KnownNat n => Vec (n + 1) a -> Vec n a
init :: forall (n :: Nat) a. KnownNat n => Vec (n + 1) a -> Vec n a
init = Vec (n + 1) a -> Vec n a
forall (n :: Nat) (m :: Nat) a.
KnownNat n =>
Vec (n + m) a -> Vec n a
take

{-# INLINE drop #-}
drop :: KnownNat m => Vec (n + m) a -> Vec m a
drop :: forall (m :: Nat) (n :: Nat) a.
KnownNat m =>
Vec (n + m) a -> Vec m a
drop = Proxy 0 -> Vec ((0 + m) + n) a -> Vector Vector m a
forall (i :: Nat) (n :: Nat) (m :: Nat) a.
(KnownNat i, KnownNat n) =>
Proxy i -> Vec ((i + n) + m) a -> Vec n a
rslice (Proxy 0
forall (n :: Nat). Proxy n
Proxy :: Proxy 0)

{-# INLINE tail #-}
tail :: KnownNat n => Vec (1 + n) a -> Vec n a
tail :: forall (n :: Nat) a. KnownNat n => Vec (1 + n) a -> Vec n a
tail = Vec (1 + n) a -> Vec n a
forall (m :: Nat) (n :: Nat) a.
KnownNat m =>
Vec (n + m) a -> Vec m a
drop

{-# INLINE update #-}
update :: KnownNat n => Vec n a -> Finite n -> a -> Vec n a
update :: forall (n :: Nat) a.
KnownNat n =>
Vec n a -> Finite n -> a -> Vec n a
update Vec n a
v Finite n
i a
x = (Finite n -> a) -> Vec n a
forall (n :: Nat) a. KnownNat n => (Finite n -> a) -> Vec n a
generate (\ Finite n
j -> if Finite n
j Finite n -> Finite n -> Bool
forall (n :: Nat). Finite n -> Finite n -> Bool
FC.== Finite n
i then a
x else Vec n a -> Finite n -> a
forall (n :: Nat) a. Vec n a -> Finite n -> a
index Vec n a
v Finite n
j)

-- {-# INLINE bulkUpdate #-}
-- bulkUpdate :: KnownNat n => Vec n a -> Vec m (Finite n,a) -> Vec n a
-- bulkUpdate = rwPrimVecBulkUpdate

{-# INLINE map #-}
map :: (a -> b) -> Vec n a -> Vec n b
map :: forall a b (n :: Nat). (a -> b) -> Vec n a -> Vec n b
map = (a -> b) -> Vec n a -> Vec n b
forall a b (n :: Nat). (a -> b) -> Vec n a -> Vec n b
rwPrimVecMap

{-# INLINE generate #-}
generate :: KnownNat n => (Finite n -> a) -> Vec n a
generate :: forall (n :: Nat) a. KnownNat n => (Finite n -> a) -> Vec n a
generate = (Finite n -> a) -> Vec n a
forall (n :: Nat) a. KnownNat n => (Finite n -> a) -> Vec n a
rwPrimVecGenerate

-- {-# INLINE iterate #-}
-- iterate :: KnownNat n => Proxy n -> (a -> a) -> a -> Vec n a
-- iterate = rwPrimVecIterate

-- {-# INLINE zip #-}
-- zip :: Vec n a -> Vec n b -> Vec n (a , b)
-- zip = rwPrimVecZip

-- {-# INLINE zipWith #-}
-- zipWith :: (a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
-- zipWith f vs ws = map (uncurry f) (zip vs ws)

-- {-# INLINE zipWith3 #-}
-- zipWith3 :: (a -> b -> c -> d) -> Vec n a -> Vec n b -> Vec n c -> Vec n d
-- zipWith3 f vs ws = zipWith (uncurry f) (zip vs ws)

-- | Returns evens from v concatenated with evens from w
{-# INLINE packlo #-}
packlo :: KnownNat n => Vec n a -> Vec n a -> Vec n a
packlo :: forall (n :: Nat) a. KnownNat n => Vec n a -> Vec n a -> Vec n a
packlo Vec n a
v Vec n a
w = (Finite n -> a) -> Vec n a
forall (n :: Nat) a. KnownNat n => (Finite n -> a) -> Vec n a
generate (\ Finite n
fi ->
      if Finite n
fi Finite n -> Finite n -> Bool
forall (n :: Nat). Finite n -> Finite n -> Bool
FC.< Finite n
n then Vec n a -> Finite n -> a
forall (n :: Nat) a. Vec n a -> Finite n -> a
index Vec n a
v (Finite n
fi Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
FC.* Finite n
two)
                 else Vec n a -> Finite n -> a
forall (n :: Nat) a. Vec n a -> Finite n -> a
index Vec n a
w ((Finite n
fi Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
FC.- Finite n
n) Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
FC.* Finite n
two))
  where
      one :: Finite n
one = Integer -> Finite n
forall (n :: Nat). KnownNat n => Integer -> Finite n
F.finite Integer
1
      two :: Finite n
two = Integer -> Finite n
forall (n :: Nat). KnownNat n => Integer -> Finite n
F.finite Integer
2
      n' :: Finite n
n' = Vec n a -> Finite n
forall (n :: Nat) a. KnownNat n => Vec n a -> Finite n
lastIndex Vec n a
v Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
`FC.div` Finite n
two
      n :: Finite n
n = if Finite n -> Bool
forall (n :: Nat). Finite n -> Bool
FC.even (Vec n a -> Finite n
forall (n :: Nat) a. KnownNat n => Vec n a -> Finite n
lastIndex Vec n a
v) then Finite n
n' else Finite n
n' Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
FC.+ Finite n
one

-- | Returns odds from v concatenated with odds from w
{-# INLINE packhi #-}
packhi :: KnownNat n => Vec n a -> Vec n a -> Vec n a
packhi :: forall (n :: Nat) a. KnownNat n => Vec n a -> Vec n a -> Vec n a
packhi Vec n a
v Vec n a
w = (Finite n -> a) -> Vec n a
forall (n :: Nat) a. KnownNat n => (Finite n -> a) -> Vec n a
generate (\ Finite n
fi ->
      if Finite n
fi Finite n -> Finite n -> Bool
forall (n :: Nat). Finite n -> Finite n -> Bool
FC.< Finite n
n then Vec n a -> Finite n -> a
forall (n :: Nat) a. Vec n a -> Finite n -> a
index Vec n a
v ((Finite n
fi Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
FC.* Finite n
two) Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
FC.+ Finite n
one)
                 else Vec n a -> Finite n -> a
forall (n :: Nat) a. Vec n a -> Finite n -> a
index Vec n a
w (((Finite n
fi Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
FC.- Finite n
n) Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
FC.* Finite n
two) Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
FC.+ Finite n
one))
  where
      one :: Finite n
one = Integer -> Finite n
forall (n :: Nat). KnownNat n => Integer -> Finite n
F.finite Integer
1
      two :: Finite n
two = Integer -> Finite n
forall (n :: Nat). KnownNat n => Integer -> Finite n
F.finite Integer
2
      n' :: Finite n
n' = Vec n a -> Finite n
forall (n :: Nat) a. KnownNat n => Vec n a -> Finite n
lastIndex Vec n a
v Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
`FC.div` Finite n
two
      n :: Finite n
n = if Finite n -> Bool
forall (n :: Nat). Finite n -> Bool
FC.even (Vec n a -> Finite n
forall (n :: Nat) a. KnownNat n => Vec n a -> Finite n
lastIndex Vec n a
v) then Finite n
n' else Finite n
n' Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
FC.+ Finite n
one

-- | Returns the first half of v interleaved with the first half of w
{-# INLINE unpacklo #-}
unpacklo :: KnownNat n => Vec n a -> Vec n a -> Vec n a
unpacklo :: forall (n :: Nat) a. KnownNat n => Vec n a -> Vec n a -> Vec n a
unpacklo Vec n a
v Vec n a
w = (Finite n -> a) -> Vec n a
forall (n :: Nat) a. KnownNat n => (Finite n -> a) -> Vec n a
generate (\ Finite n
fi ->
      if Finite n -> Bool
forall (n :: Nat). Finite n -> Bool
FC.even Finite n
fi then Vec n a -> Finite n -> a
forall (n :: Nat) a. Vec n a -> Finite n -> a
index Vec n a
v (Finite n
fi Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
`FC.div` Finite n
two)
                   else Vec n a -> Finite n -> a
forall (n :: Nat) a. Vec n a -> Finite n -> a
index Vec n a
w ((Finite n
fi Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
FC.- Finite n
one) Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
`FC.div` Finite n
two))
  where
      one :: Finite n
one = Integer -> Finite n
forall (n :: Nat). KnownNat n => Integer -> Finite n
F.finite Integer
1
      two :: Finite n
two = Integer -> Finite n
forall (n :: Nat). KnownNat n => Integer -> Finite n
F.finite Integer
2

-- | Returns the second half of v interleaved with the second half of w
{-# INLINE unpackhi #-}
unpackhi :: KnownNat n => Vec n a -> Vec n a -> Vec n a
unpackhi :: forall (n :: Nat) a. KnownNat n => Vec n a -> Vec n a -> Vec n a
unpackhi Vec n a
v Vec n a
w = (Finite n -> a) -> Vec n a
forall (n :: Nat) a. KnownNat n => (Finite n -> a) -> Vec n a
generate (\ Finite n
fi ->
      if Finite n -> Bool
forall (n :: Nat). Finite n -> Bool
FC.even Finite n
fi then Vec n a -> Finite n -> a
forall (n :: Nat) a. Vec n a -> Finite n -> a
index Vec n a
v (Finite n
n Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
FC.+ (Finite n
fi Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
`FC.div` Finite n
two))
                   else Vec n a -> Finite n -> a
forall (n :: Nat) a. Vec n a -> Finite n -> a
index Vec n a
w (Finite n
n Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
FC.+ ((Finite n
fi Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
FC.- Finite n
one) Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
`FC.div` Finite n
two)))
  where
      one :: Finite n
one = Integer -> Finite n
forall (n :: Nat). KnownNat n => Integer -> Finite n
F.finite Integer
1
      two :: Finite n
two = Integer -> Finite n
forall (n :: Nat). KnownNat n => Integer -> Finite n
F.finite Integer
2
      n' :: Finite n
n' = Vec n a -> Finite n
forall (n :: Nat) a. KnownNat n => Vec n a -> Finite n
lastIndex Vec n a
v Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
`FC.div` Finite n
two
      n :: Finite n
n = if Finite n -> Bool
forall (n :: Nat). Finite n -> Bool
FC.even (Vec n a -> Finite n
forall (n :: Nat) a. KnownNat n => Vec n a -> Finite n
lastIndex Vec n a
v) then Finite n
n' else Finite n
n' Finite n -> Finite n -> Finite n
forall (n :: Nat). KnownNat n => Finite n -> Finite n -> Finite n
FC.+ Finite n
one

-- | lookup value at index n in vector
{-# INLINE (!) #-}
(!) :: KnownNat n => Vec ((n + m) + 1) a -> Proxy n -> a
! :: forall (n :: Nat) (m :: Nat) a.
KnownNat n =>
Vec ((n + m) + 1) a -> Proxy n -> a
(!) = Vec ((n + m) + 1) a -> Proxy n -> a
forall (n :: Nat) (m :: Nat) a.
KnownNat n =>
Vec ((n + m) + 1) a -> Proxy n -> a
index'

infixl 9 !

-- | assign new value a to index i
{-# INLINE (!=) #-}
(!=) :: KnownNat n => Vec n a -> Finite n -> a -> Vec n a
Vec n a
v != :: forall (n :: Nat) a.
KnownNat n =>
Vec n a -> Finite n -> a -> Vec n a
!= Finite n
i = Vec n a -> Finite n -> a -> Vec n a
forall (n :: Nat) a.
KnownNat n =>
Vec n a -> Finite n -> a -> Vec n a
update Vec n a
v Finite n
i

infixr 2  !=