diff options
Diffstat (limited to 'src/Data/Array/Nested')
| -rw-r--r-- | src/Data/Array/Nested/Internal/Shape.hs | 62 |
1 files changed, 40 insertions, 22 deletions
diff --git a/src/Data/Array/Nested/Internal/Shape.hs b/src/Data/Array/Nested/Internal/Shape.hs index 5fb2e7f..878ea7e 100644 --- a/src/Data/Array/Nested/Internal/Shape.hs +++ b/src/Data/Array/Nested/Internal/Shape.hs @@ -90,6 +90,13 @@ listrShow f l = showString "[" . go "" l . showString "]" go _ ZR = id go prefix (x ::: xs) = showString prefix . f x . go "," xs +listrLength :: ListR n i -> Int +listrLength = length + +listrRank :: ListR n i -> SNat n +listrRank ZR = SNat +listrRank (_ ::: sh) = snatSucc (listrRank sh) + listrAppend :: ListR n i -> ListR m i -> ListR (n + m) i listrAppend ZR sh = sh listrAppend (x ::: xs) sh = x ::: listrAppend xs sh @@ -121,10 +128,6 @@ listrIndex SZ (x ::: _) = x listrIndex (SS i) (_ ::: xs) | Refl <- lemLeqSuccSucc (Proxy @k) (Proxy @n) = listrIndex i xs listrIndex _ ZR = error "k + 1 <= 0" -listrRank :: ListR n i -> SNat n -listrRank ZR = SNat -listrRank (_ ::: sh) = snatSucc (listrRank sh) - listrPermutePrefix :: forall i n. [Int] -> ListR n i -> ListR n i listrPermutePrefix = \perm sh -> listrFromList perm $ \sperm -> @@ -175,6 +178,12 @@ type IIxR n = IxR n Int instance Show i => Show (IxR n i) where showsPrec _ (IxR l) = listrShow shows l +ixrLength :: IxR sh i -> Int +ixrLength (IxR l) = listrLength l + +ixrRank :: IxR n i -> SNat n +ixrRank (IxR sh) = listrRank sh + ixrZero :: SNat n -> IIxR n ixrZero SZ = ZIR ixrZero (SS n) = 0 :.: ixrZero n @@ -204,9 +213,6 @@ ixrLast (IxR list) = listrLast list ixrAppend :: forall n m i. IxR n i -> IxR m i -> IxR (n + m) i ixrAppend = coerce (listrAppend @_ @i) -ixrRank :: IxR n i -> SNat n -ixrRank (IxR sh) = listrRank sh - ixrPermutePrefix :: forall n i. [Int] -> IxR n i -> IxR n i ixrPermutePrefix = coerce (listrPermutePrefix @i) @@ -268,6 +274,15 @@ shrEqRank (ShR sh) (ShR sh') = listrEqRank sh sh' shrEqual :: Eq i => ShR n i -> ShR n' i -> Maybe (n :~: n') shrEqual (ShR sh) (ShR sh') = listrEqual sh sh' +shrLength :: ShR sh i -> Int +shrLength (ShR l) = listrLength l + +-- | This function can also be used to conjure up a 'KnownNat' dictionary; +-- pattern matching on the returned 'SNat' with the 'pattern SNat' pattern +-- synonym yields 'KnownNat' evidence. +shrRank :: ShR n i -> SNat n +shrRank (ShR sh) = listrRank sh + -- | The number of elements in an array described by this shape. shrSize :: IShR n -> Int shrSize ZSR = 1 @@ -288,12 +303,6 @@ shrLast (ShR list) = listrLast list shrAppend :: forall n m i. ShR n i -> ShR m i -> ShR (n + m) i shrAppend = coerce (listrAppend @_ @i) --- | This function can also be used to conjure up a 'KnownNat' dictionary; --- pattern matching on the returned 'SNat' with the 'pattern SNat' pattern --- synonym yields 'KnownNat' evidence. -shrRank :: ShR n i -> SNat n -shrRank (ShR sh) = listrRank sh - shrPermutePrefix :: forall n i. [Int] -> ShR n i -> ShR n i shrPermutePrefix = coerce (listrPermutePrefix @i) @@ -381,6 +390,13 @@ listsShow f l = showString "[" . go "" l . showString "]" go _ ZS = id go prefix (x ::$ xs) = showString prefix . f x . go "," xs +listsLength :: ListS sh f -> Int +listsLength = getSum . listsFold (\_ -> Sum 1) + +listsRank :: ListS sh f -> SNat (Rank sh) +listsRank ZS = SNat +listsRank (_ ::$ sh) = snatSucc (listsRank sh) + listsToList :: ListS sh (Const i) -> [i] listsToList ZS = [] listsToList (Const i ::$ is) = i : listsToList is @@ -403,10 +419,6 @@ listsAppend :: ListS sh f -> ListS sh' f -> ListS (sh ++ sh') f listsAppend ZS idx' = idx' listsAppend (i ::$ idx) idx' = i ::$ listsAppend idx idx' -listsRank :: ListS sh f -> SNat (Rank sh) -listsRank ZS = SNat -listsRank (_ ::$ sh) = snatSucc (listsRank sh) - listsTakeLenPerm :: forall f is sh. Perm is -> ListS sh f -> ListS (TakeLen is sh) f listsTakeLenPerm PNil _ = ZS listsTakeLenPerm (_ `PCons` is) (n ::$ sh) = n ::$ listsTakeLenPerm is sh @@ -468,6 +480,12 @@ instance Functor (IxS sh) where instance Foldable (IxS sh) where foldMap f (IxS l) = listsFold (f . getConst) l +ixsLength :: IxS sh i -> Int +ixsLength (IxS l) = listsLength l + +ixsRank :: IxS sh i -> SNat (Rank sh) +ixsRank (IxS l) = listsRank l + ixsZero :: ShS sh -> IIxS sh ixsZero ZSS = ZIS ixsZero (_ :$$ sh) = 0 :.$ ixsZero sh @@ -534,11 +552,15 @@ shsEqual :: ShS sh -> ShS sh' -> Maybe (sh :~: sh') shsEqual = testEquality shsLength :: ShS sh -> Int -shsLength (ShS l) = getSum (listsFold (\_ -> Sum 1) l) +shsLength (ShS l) = listsLength l shsRank :: ShS sh -> SNat (Rank sh) shsRank (ShS l) = listsRank l +shsSize :: ShS sh -> Int +shsSize ZSS = 1 +shsSize (n :$$ sh) = fromSNat' n * shsSize sh + shsToList :: ShS sh -> [Int] shsToList ZSS = [] shsToList (sn :$$ sh) = fromSNat' sn : shsToList sh @@ -575,10 +597,6 @@ shsLast (ShS list) = listsLast 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) |