Make ShS a newtype over ShX

TODO: use lemmas in place of the unsafeCoerceRefl

1 files changed, 85 insertions, 34 deletions

diff --git a/src/Data/Array/Nested/Shaped/Shape.hs b/src/Data/Array/Nested/Shaped/Shape.hs
index 6485c72..d815adf 100644
--- a/src/Data/Array/Nested/Shaped/Shape.hs
+++ b/src/Data/Array/Nested/Shaped/Shape.hs
@@ -331,21 +331,34 @@ ixsToLinear = \sh i -> go sh i 0
 -- can also retrieve the array shape from a 'KnownShS' dictionary.
 type role ShS nominal
 type ShS :: [Nat] -> Type
-newtype ShS sh = ShS (ListS sh SNat)
+newtype ShS sh = ShS (ShX (MapJust sh) Int)
   deriving (Generic)
 
 instance Eq (ShS sh) where _ == _ = True
 instance Ord (ShS sh) where compare _ _ = EQ
 
 pattern ZSS :: forall sh. () => sh ~ '[] => ShS sh
-pattern ZSS = ShS ZS
+pattern ZSS <- ShS (matchZSX -> Just Refl)
+  where ZSS = ShS ZSX
+
+matchZSX :: forall sh i. ShX (MapJust sh) i -> Maybe (sh :~: '[])
+matchZSX ZSX | Refl <- lemMapJustEmpty @sh Refl = Just Refl
+matchZSX _ = Nothing
 
 pattern (:$$)
   :: forall {sh1}.
      forall n sh. (n : sh ~ sh1)
   => SNat n -> ShS sh -> ShS sh1
-pattern i :$$ shl <- ShS (listsUncons -> Just (UnconsListSRes (ShS -> shl) i))
-  where i :$$ ShS shl = ShS (i ::$ shl)
+pattern i :$$ shl <- (shsUncons -> Just (UnconsShSRes i shl))
+  where i :$$ ShS shl = ShS (SKnown i :$% shl)
+
+data UnconsShSRes sh1 =
+  forall n sh. (n : sh ~ sh1) => UnconsShSRes (SNat n) (ShS sh)
+shsUncons :: forall sh1. ShS sh1 -> Maybe (UnconsShSRes sh1)
+shsUncons (ShS (SKnown x :$% sh'))
+  | Refl <- lemMapJustCons @sh1 Refl
+  = Just (UnconsShSRes x (ShS sh'))
+shsUncons (ShS _) = Nothing
 
 infixr 3 :$$
 
@@ -355,15 +368,16 @@ infixr 3 :$$
 deriving instance Show (ShS sh)
 #else
 instance Show (ShS sh) where
-  showsPrec _ (ShS l) = listsShow (shows . fromSNat) l
+  showsPrec d (ShS shx) = showsPrec d shx
 #endif
 
 instance NFData (ShS sh) where
-  rnf (ShS ZS) = ()
-  rnf (ShS (SNat ::$ l)) = rnf (ShS l)
+  rnf (ShS shx) = rnf shx
 
 instance TestEquality ShS where
-  testEquality (ShS l1) (ShS l2) = listsEqType l1 l2
+  testEquality (ShS shx1) (ShS shx2) = case shxEqType shx1 shx2 of
+    Nothing -> Nothing
+    Just Refl -> Just unsafeCoerceRefl
 
 -- | @'shsEqual' = 'testEquality'@. (Because 'ShS' is a singleton, types are
 -- equal if and only if values are equal.)
@@ -371,17 +385,20 @@ shsEqual :: ShS sh -> ShS sh' -> Maybe (sh :~: sh')
 shsEqual = testEquality
 
 shsLength :: ShS sh -> Int
-shsLength (ShS l) = listsLength l
+shsLength (ShS shx) = shxLength shx
 
-shsRank :: ShS sh -> SNat (Rank sh)
-shsRank (ShS l) = listsRank l
+shsRank :: forall sh. ShS sh -> SNat (Rank sh)
+shsRank (ShS shx) =
+  gcastWith (unsafeCoerceRefl
+             :: Rank (MapJust sh) :~: Rank sh) $
+  shxRank shx
 
 shsSize :: ShS sh -> Int
-shsSize ZSS = 1
-shsSize (n :$$ sh) = fromSNat' n * shsSize sh
+shsSize (ShS sh) = shxSize sh
 
 -- | This is a partial @const@ that fails when the second argument
--- doesn't match the first.
+-- doesn't match the first. It also has a better error message comparing
+-- to just coercing 'shxFromList'.
 shsFromList :: ShS sh -> [Int] -> ShS sh
 shsFromList topsh topl = go topsh topl `seq` topsh
   where
@@ -397,38 +414,72 @@ shsFromList topsh topl = go topsh topl `seq` topsh
 
 {-# INLINEABLE shsToList #-}
 shsToList :: ShS sh -> [Int]
-shsToList topsh = build (\(cons :: Int -> is -> is) (nil :: is) ->
-  let go :: ShS sh -> is
-      go ZSS = nil
-      go (sn :$$ sh) = fromSNat' sn `cons` go sh
-  in go topsh)
+shsToList = coerce shxToList
 
 shsHead :: ShS (n : sh) -> SNat n
-shsHead (ShS list) = listsHead list
+shsHead (ShS shx) = case shxHead shx of
+  SKnown SNat -> SNat
 
-shsTail :: ShS (n : sh) -> ShS sh
-shsTail (ShS list) = ShS (listsTail list)
+shsTail :: forall n sh. ShS (n : sh) -> ShS sh
+shsTail = coerce (shxTail @_ @_ @Int)
 
-shsInit :: ShS (n : sh) -> ShS (Init (n : sh))
-shsInit (ShS list) = ShS (listsInit list)
+shsInit :: forall n sh. ShS (n : sh) -> ShS (Init (n : sh))
+shsInit =
+  gcastWith (unsafeCoerceRefl
+             :: Init (Just n : MapJust sh) :~: MapJust (Init (n : sh))) $
+  coerce (shxInit @_ @_ @Int)
 
-shsLast :: ShS (n : sh) -> SNat (Last (n : sh))
-shsLast (ShS list) = listsLast list
+shsLast :: forall n sh. ShS (n : sh) -> SNat (Last (n : sh))
+shsLast (ShS shx) =
+  gcastWith (unsafeCoerceRefl
+             :: Last (Just n : MapJust sh) :~: Just (Last (n : sh))) $
+  case shxLast shx of
+    SKnown SNat -> SNat
 
 shsAppend :: forall sh sh'. ShS sh -> ShS sh' -> ShS (sh ++ sh')
-shsAppend = coerce (listsAppend @_ @SNat)
+shsAppend =
+  gcastWith (unsafeCoerceRefl
+             :: MapJust sh ++ MapJust sh' :~: MapJust (sh ++ sh')) $
+  coerce (shxAppend @_ @_ @Int)
+
+shsTakeLen :: forall is sh. Perm is -> ShS sh -> ShS (TakeLen is sh)
+shsTakeLen =
+  gcastWith (unsafeCoerceRefl
+             :: TakeLen is (MapJust sh) :~: MapJust (TakeLen is sh)) $
+  coerce shxTakeLen
 
-shsTakeLen :: Perm is -> ShS sh -> ShS (TakeLen is sh)
-shsTakeLen = coerce (listsTakeLenPerm @SNat)
+shsDropLen :: forall is sh. Perm is -> ShS sh -> ShS (DropLen is sh)
+shsDropLen =
+  gcastWith (unsafeCoerceRefl
+             :: DropLen is (MapJust sh) :~: MapJust (DropLen is sh)) $
+  coerce shxDropLen
 
-shsPermute :: Perm is -> ShS sh -> ShS (Permute is sh)
-shsPermute = coerce (listsPermute @SNat)
+shsPermute :: forall is sh. Perm is -> ShS sh -> ShS (Permute is sh)
+shsPermute =
+  gcastWith (unsafeCoerceRefl
+             :: Permute is (MapJust sh) :~: MapJust (Permute is sh)) $
+  coerce shxPermute
 
-shsIndex :: Proxy is -> Proxy shT -> SNat i -> ShS sh -> SNat (Index i sh)
-shsIndex pis pshT i sh = coerce (listsIndex @SNat pis pshT i (coerce sh))
+shsIndex :: forall is shT i sh.
+            Proxy is -> Proxy shT -> SNat i -> ShS sh -> SNat (Index i sh)
+shsIndex pis pshT i sh =
+  gcastWith (unsafeCoerceRefl
+             :: Index i (MapJust sh) :~: Just (Index i sh)) $
+  case shxIndex pis pshT i (coerce sh) of
+    SKnown SNat -> SNat
 
 shsPermutePrefix :: forall is sh. Perm is -> ShS sh -> ShS (PermutePrefix is sh)
-shsPermutePrefix = coerce (listsPermutePrefix @SNat)
+shsPermutePrefix perm (ShS shx)
+  {- TODO: here and elsewhere, solve the module dependency cycle and add this:
+  | Refl <- lemTakeLenMapJust perm sh
+  , Refl <- lemDropLenMapJust perm sh
+  , Refl <- lemPermuteMapJust perm sh
+  , Refl <- lemMapJustApp (shsPermute perm (shsTakeLen perm sh)) (shsDropLen perm sh) -}
+  = gcastWith (unsafeCoerceRefl
+               :: Permute is (TakeLen is (MapJust sh))
+                  ++ DropLen is (MapJust sh)
+                  :~: MapJust (Permute is (TakeLen is sh) ++ DropLen is sh)) $
+    ShS (shxPermutePrefix perm shx)
 
 type family Product sh where
   Product '[] = 1