Add some utility functions for arrays

src/Data/NumIdr/Array.idr

@@ -1,41 +1,63 @@
 module Data.NumIdr.Array
 
 import Data.Vect
+import Data.Reify
 import Data.NumIdr.PrimArray
 
 %default total
 
 
 export
-data Array : Vect rank Nat -> Type -> Type where
-  MkArray : PrimArray a -> Array dim a
+record Array {0 rk : Nat} (s : Vect rk Nat) a where
+  constructor MkArray
+  shape : Reify s
+  contents : PrimArray a
 
+mkArray : {s : _} -> PrimArray a -> Array s a
+mkArray arr = MkArray (MkReify _) arr
 
 
 export
-size : Array dim a -> Nat
-size (MkArray arr) = length arr
+size : Array s a -> Nat
+size arr = length arr.contents
 
 export
-fromVect : {0 dim : Vect rank Nat} -> Vect (product dim) a -> Array dim a
-fromVect = MkArray . fromList . toList
+shape : {0 s : Vect rk Nat} -> Array s a -> Vect rk Nat
+shape arr = getReify arr.shape
 
 export
-reshape : {0 dim : Vect rank Nat} -> (0 dim' : Vect rank' Nat) ->
-           Array dim a -> product dim = product dim' =>
-           Array dim' a
-reshape _ (MkArray arr) = MkArray arr
+rank : Array s a -> Nat
+rank arr = length $ shape arr
 
-Vects : Vect rank Nat -> Type -> Type
+export
+fromVect : {s : Vect rk Nat} -> Vect (product s) a -> Array s a
+fromVect = mkArray . fromList . toList
+
+public export
+Vects : Vect rk Nat -> Type -> Type
 Vects [] a = a
-Vects (d :: dim) a = Vect d (Vects dim a)
+Vects (d :: s) a = Vect d (Vects s a)
+
+export
+fromVects : {s : _} -> Vects s a -> Array s a
+fromVects = mkArray . fromList . collapse
+  where
+    collapse : {s : Vect rk Nat} -> Vects s a -> List a
+    collapse {s=[]} x = [x]
+    collapse {s=_::_} v = concat $ map collapse v
+
+export
+reshape : {0 s : Vect rk Nat} -> (s' : Vect rk' Nat) ->
+           Array s a -> product s = product s' =>
+           Array s' a
+reshape _ arr = mkArray arr.contents
 
 
 
 export
-constant : a -> (dim : Vect rank Nat) -> Array dim a
-constant x dim = MkArray $ create (product dim) (const x)
+constant : a -> (s : Vect rk Nat) -> Array s a
+constant x s = mkArray $ create (product s) (const x)
 
 export
-zeros : Num a => (dim : Vect rank Nat) -> Array dim a
+zeros : Num a => (s : Vect rk Nat) -> Array s a
 zeros = constant 0