Flesh out shaped sized lists

1 files changed, 51 insertions, 5 deletions

diff --git a/src/Data/Array/Nested/Internal.hs b/src/Data/Array/Nested/Internal.hs
index e42de12..9cabdc6 100644
--- a/src/Data/Array/Nested/Internal.hs
+++ b/src/Data/Array/Nested/Internal.hs
@@ -1043,6 +1043,14 @@ instance (KnownShape sh, Storable a, Num a) => Num (Shaped sh (Primitive a)) whe
 deriving via Shaped sh (Primitive Int) instance KnownShape sh => Num (Shaped sh Int)
 deriving via Shaped sh (Primitive Double) instance KnownShape sh => Num (Shaped sh Double)
 
+type ListS :: Type -> [Nat] -> Type
+data ListS i n where
+  ZS :: ListS i '[]
+  (::$) :: forall n sh i. i -> ListS i sh -> ListS i (n : sh)
+deriving instance Show i => Show (ListS i n)
+deriving instance Eq i => Eq (ListS i n)
+infixr 3 ::$
+
 -- | An index into a shape-typed array.
 --
 -- For convenience, this contains regular 'Int's instead of bounded integers
@@ -1050,15 +1058,53 @@ deriving via Shaped sh (Primitive Double) instance KnownShape sh => Num (Shaped
 -- shape-typed array is known statically, you can also retrieve the array shape
 -- from a 'KnownShape' dictionary.
 type IxS :: Type -> [Nat] -> Type
-data IxS i sh where
-  ZIS :: IxS i '[]
-  (:.$) :: forall n sh i. i -> IxS i sh -> IxS i (n : sh)
-deriving instance Show i => Show (IxS i n)
-deriving instance Eq i => Eq (IxS i n)
+newtype IxS i sh = IxS (ListS i sh)
+  deriving (Show, Eq)
+
+pattern ZIS :: forall sh i. () => sh ~ '[] => IxS i sh
+pattern ZIS = IxS ZS
+
+pattern (:.$)
+  :: forall {sh1} {i}.
+     forall n sh. (n : sh ~ sh1)
+  => i -> IxS i sh -> IxS i sh1
+pattern i :.$ shl <- (unconsIxS -> Just (UnconsIxSRes shl i))
+  where i :.$ (IxS shl) = IxS (i ::$ shl)
+{-# COMPLETE ZIS, (:.$) #-}
 infixr 3 :.$
 
+data UnconsIxSRes i sh1 =
+  forall n sh. (n : sh ~ sh1) => UnconsIxSRes (IxS i sh) i
+unconsIxS :: IxS i sh1 -> Maybe (UnconsIxSRes i sh1)
+unconsIxS (IxS shl) = case shl of
+  i ::$ shl' -> Just (UnconsIxSRes (IxS shl') i)
+  ZS -> Nothing
+
 type IIxS = IxS Int
 
+type StaticShapeS :: Type -> [Nat] -> Type
+newtype StaticShapeS i sh = StaticShapeS (ListS i sh)
+  deriving (Show, Eq)
+
+pattern ZSS :: forall sh i. () => sh ~ '[] => StaticShapeS i sh
+pattern ZSS = StaticShapeS ZS
+
+pattern (:$$)
+  :: forall {sh1} {i}.
+     forall n sh. (n : sh ~ sh1)
+  => i -> StaticShapeS i sh -> StaticShapeS i sh1
+pattern i :$$ shl <- (unconsStaticShapeS -> Just (UnconsStaticShapeSRes shl i))
+  where i :$$ (StaticShapeS shl) = StaticShapeS (i ::$ shl)
+{-# COMPLETE ZSS, (:$$) #-}
+infixr 3 :$$
+
+data UnconsStaticShapeSRes i sh1 =
+  forall n sh. (n : sh ~ sh1) => UnconsStaticShapeSRes (StaticShapeS i sh) i
+unconsStaticShapeS :: StaticShapeS i sh1 -> Maybe (UnconsStaticShapeSRes i sh1)
+unconsStaticShapeS (StaticShapeS shl) = case shl of
+  i ::$ shl' -> Just (UnconsStaticShapeSRes (StaticShapeS shl') i)
+  ZS -> Nothing
+
 zeroIxS :: SShape sh -> IIxS sh
 zeroIxS ShNil = ZIS
 zeroIxS (ShCons _ sh) = 0 :.$ zeroIxS sh