diff options
Diffstat (limited to 'src/Array.hs')
| -rw-r--r-- | src/Array.hs | 15 |
1 files changed, 15 insertions, 0 deletions
diff --git a/src/Array.hs b/src/Array.hs index 9a770c4..0d585a9 100644 --- a/src/Array.hs +++ b/src/Array.hs @@ -17,11 +17,13 @@ data Shape n where ShNil :: Shape Z ShCons :: Shape n -> Int -> Shape (S n) deriving instance Show (Shape n) +deriving instance Eq (Shape n) data Index n where IxNil :: Index Z IxCons :: Index n -> Int -> Index (S n) deriving instance Show (Index n) +deriving instance Eq (Index n) shapeSize :: Shape n -> Int shapeSize ShNil = 0 @@ -38,6 +40,10 @@ toLinearIndex :: Shape n -> Index n -> Int toLinearIndex ShNil IxNil = 0 toLinearIndex (sh `ShCons` n) (idx `IxCons` i) = toLinearIndex sh idx * n + i +emptyShape :: SNat n -> Shape n +emptyShape SZ = ShNil +emptyShape (SS m) = emptyShape m `ShCons` 0 + -- | TODO: this Vector is a boxed vector, which is horrendously inefficient. data Array (n :: Nat) t = Array (Shape n) (Vector t) @@ -49,6 +55,9 @@ arrayShape (Array sh _) = sh arraySize :: Array n t -> Int arraySize (Array sh _) = shapeSize sh +emptyArray :: SNat n -> Array n t +emptyArray n = Array (emptyShape n) V.empty + arrayIndex :: Array n t -> Index n -> t arrayIndex arr@(Array sh _) idx = arrayIndexLinear arr (toLinearIndex sh idx) @@ -58,6 +67,12 @@ arrayIndexLinear (Array _ v) i = v V.! i arrayIndex1 :: Array (S n) t -> Int -> Array n t arrayIndex1 (Array (sh `ShCons` _) v) i = let sz = shapeSize sh in Array sh (V.slice (sz * i) sz v) +arrayGenerate :: Shape n -> (Index n -> t) -> Array n t +arrayGenerate sh f = arrayGenerateLin sh (f . fromLinearIndex sh) + +arrayGenerateLin :: Shape n -> (Int -> t) -> Array n t +arrayGenerateLin sh f = Array sh (V.generate (shapeSize sh) f) + arrayGenerateM :: Monad m => Shape n -> (Index n -> m t) -> m (Array n t) arrayGenerateM sh f = arrayGenerateLinM sh (f . fromLinearIndex sh) |