aboutsummaryrefslogtreecommitdiff
path: root/src/Data/Array/Nested/Ranked
diff options
context:
space:
mode:
authorTom Smeding <tom@tomsmeding.com>2025-11-18 11:39:39 +0100
committerTom Smeding <tom@tomsmeding.com>2025-11-18 11:39:39 +0100
commitaeb62e1d59b4cdc426f6ca0ef8df6a0737bb0f86 (patch)
tree3fb5c3255b19a4961d2adad8eb32b1a6ff66ebcd /src/Data/Array/Nested/Ranked
parent0ae9eeef25e72a0a69b6ed078d282f6892e1fc0b (diff)
Add shsEnum and shrEnum
Diffstat (limited to 'src/Data/Array/Nested/Ranked')
-rw-r--r--src/Data/Array/Nested/Ranked/Shape.hs7
1 files changed, 7 insertions, 0 deletions
diff --git a/src/Data/Array/Nested/Ranked/Shape.hs b/src/Data/Array/Nested/Ranked/Shape.hs
index 488b4d8..a1fa247 100644
--- a/src/Data/Array/Nested/Ranked/Shape.hs
+++ b/src/Data/Array/Nested/Ranked/Shape.hs
@@ -330,6 +330,13 @@ shrZipWith f (ShR l1) (ShR l2) = ShR $ listrZipWith f l1 l2
shrPermutePrefix :: forall n i. [Int] -> ShR n i -> ShR n i
shrPermutePrefix = coerce (listrPermutePrefix @i)
+shrEnum :: IShR sh -> [IIxR sh]
+shrEnum = \sh -> go sh id []
+ where
+ go :: IShR sh -> (IIxR sh -> a) -> [a] -> [a]
+ go ZSR f = (f ZIR :)
+ go (n :$: sh) f = foldr (.) id [go sh (f . (i :.:)) | i <- [0 .. n - 1]]
+
-- | Untyped: length is checked at runtime.
instance KnownNat n => IsList (ListR n i) where