diff options
author | Tom Smeding <tom@tomsmeding.com> | 2024-06-03 21:28:58 +0200 |
---|---|---|
committer | Tom Smeding <tom@tomsmeding.com> | 2024-06-03 21:28:58 +0200 |
commit | 0fd727dcb3fe05816aa9c68be5ebac84a55fcf4b (patch) | |
tree | ac34666b460d1f02ddc482ae06ca4228d5602adc /src/Data/Array/Nested | |
parent | 3286a65fe6e4735aaadef5addecbe3c3f7ed3468 (diff) |
Cleanup of shape functions
Diffstat (limited to 'src/Data/Array/Nested')
-rw-r--r-- | src/Data/Array/Nested/Internal/Shape.hs | 46 |
1 files changed, 29 insertions, 17 deletions
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 |