From 0fd727dcb3fe05816aa9c68be5ebac84a55fcf4b Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Mon, 3 Jun 2024 21:28:58 +0200
Subject: Cleanup of shape functions

---
 src/Data/Array/Nested/Internal/Shape.hs | 46 +++++++++++++++++++++------------
 1 file changed, 29 insertions(+), 17 deletions(-)

(limited to 'src/Data')

diff --git a/src/Data/Array/Nested/Internal/Shape.hs b/src/Data/Array/Nested/Internal/Shape.hs
index bce85b0..c66b467 100644
--- a/src/Data/Array/Nested/Internal/Shape.hs
+++ b/src/Data/Array/Nested/Internal/Shape.hs
@@ -61,10 +61,10 @@ listrUncons :: ListR n1 i -> Maybe (UnconsListRRes i n1)
 listrUncons (i ::: sh') = Just (UnconsListRRes sh' i)
 listrUncons ZR = Nothing
 
-listrShow :: forall sh i. (i -> ShowS) -> ListR sh i -> ShowS
+listrShow :: forall n i. (i -> ShowS) -> ListR n i -> ShowS
 listrShow f l = showString "[" . go "" l . showString "]"
   where
-    go :: String -> ListR sh' i -> ShowS
+    go :: String -> ListR n' i -> ShowS
     go _ ZR = id
     go prefix (x ::: xs) = showString prefix . f x . go "," xs
 
@@ -156,6 +156,9 @@ ixCvtRX (n :.: (idx :: IxR m Int)) =
 ixrTail :: IxR (n + 1) i -> IxR n i
 ixrTail (IxR list) = IxR (listrTail list)
 
+ixrAppend :: forall n m i. IxR n i -> IxR m i -> IxR (n + m) i
+ixrAppend = coerce (listrAppend @_ @i)
+
 ixrLengthSNat :: IxR n i -> SNat n
 ixrLengthSNat (IxR sh) = listrLengthSNat sh
 
@@ -219,6 +222,9 @@ shrSize (n :$: sh) = n * shrSize sh
 shrTail :: ShR (n + 1) i -> ShR n i
 shrTail (ShR list) = ShR (listrTail list)
 
+shrAppend :: forall n m i. ShR n i -> ShR m i -> ShR (n + m) i
+shrAppend = coerce (listrAppend @_ @i)
+
 shrLengthSNat :: ShR n i -> SNat n
 shrLengthSNat (ShR sh) = listrLengthSNat sh
 
@@ -319,24 +325,9 @@ listsIndex p pT (SS (i :: SNat i')) ((_ :: f n) ::$ (sh :: ListS sh' f))
   = listsIndex p pT i sh
 listsIndex _ _ _ ZS = error "Index into empty shape"
 
-shsTakeLen :: Perm is -> ShS sh -> ShS (TakeLen is sh)
-shsTakeLen = coerce (listsTakeLenPerm @SNat)
-
-shsPermute :: Perm is -> ShS sh -> ShS (Permute is sh)
-shsPermute = coerce (listsPermute @SNat)
-
-shsIndex :: Proxy is -> Proxy shT -> SNat i -> ShS sh -> SNat (Index i sh)
-shsIndex pis pshT i sh = coerce (fst (listsIndex @SNat pis pshT i (coerce sh)))
-
 listsPermutePrefix :: forall f is sh. Perm is -> ListS sh f -> ListS (PermutePrefix is sh) f
 listsPermutePrefix perm sh = listsAppend (listsPermute perm (listsTakeLenPerm perm sh)) (listsDropLenPerm perm sh)
 
-ixsPermutePrefix :: forall i is sh. Perm is -> IxS sh i -> IxS (PermutePrefix is sh) i
-ixsPermutePrefix = coerce (listsPermutePrefix @(Const i))
-
-shsPermutePrefix :: forall is sh. Perm is -> ShS sh -> ShS (PermutePrefix is sh)
-shsPermutePrefix = coerce (listsPermutePrefix @SNat)
-
 
 -- | An index into a shape-typed array.
 --
@@ -386,6 +377,12 @@ ixCvtSX (n :.$ sh) = n :.% ixCvtSX sh
 ixsTail :: IxS (n : sh) i -> IxS sh i
 ixsTail (IxS list) = IxS (listsTail list)
 
+ixsAppend :: forall sh sh' i. IxS sh i -> IxS sh' i -> IxS (sh ++ sh') i
+ixsAppend = coerce (listsAppend @_ @(Const i))
+
+ixsPermutePrefix :: forall i is sh. Perm is -> IxS sh i -> IxS (PermutePrefix is sh) i
+ixsPermutePrefix = coerce (listsPermutePrefix @(Const i))
+
 
 -- | The shape of a shape-typed array given as a list of 'SNat' values.
 --
@@ -440,10 +437,25 @@ shCvtSX (n :$$ sh) = SKnown n :$% shCvtSX sh
 shsTail :: ShS (n : sh) -> ShS sh
 shsTail (ShS list) = ShS (listsTail list)
 
+shsAppend :: forall sh sh'. ShS sh -> ShS sh' -> ShS (sh ++ sh')
+shsAppend = coerce (listsAppend @_ @SNat)
+
 shsSize :: ShS sh -> Int
 shsSize ZSS = 1
 shsSize (n :$$ sh) = fromSNat' n * shsSize sh
 
+shsTakeLen :: Perm is -> ShS sh -> ShS (TakeLen is sh)
+shsTakeLen = coerce (listsTakeLenPerm @SNat)
+
+shsPermute :: Perm is -> ShS sh -> ShS (Permute is sh)
+shsPermute = coerce (listsPermute @SNat)
+
+shsIndex :: Proxy is -> Proxy shT -> SNat i -> ShS sh -> SNat (Index i sh)
+shsIndex pis pshT i sh = coerce (fst (listsIndex @SNat pis pshT i (coerce sh)))
+
+shsPermutePrefix :: forall is sh. Perm is -> ShS sh -> ShS (PermutePrefix is sh)
+shsPermutePrefix = coerce (listsPermutePrefix @SNat)
+
 -- | Evidence for the static part of a shape. This pops up only when you are
 -- polymorphic in the element type of an array.
 type KnownShS :: [Nat] -> Constraint
-- 
cgit v1.2.3-70-g09d2