aboutsummaryrefslogtreecommitdiff
path: root/src/Data
diff options
context:
space:
mode:
authorTom Smeding <t.j.smeding@uu.nl>2024-04-03 15:06:29 +0200
committerTom Smeding <t.j.smeding@uu.nl>2024-04-03 15:06:29 +0200
commit6520c4e4bc39b7315f4f9416fe68b883cfdde8ec (patch)
tree8315b318c6be640ab46b306270aa644e7623a472 /src/Data
parent6f912510c594de511bc92b7f47d0d405b64aaa1e (diff)
Fix enumShape
Diffstat (limited to 'src/Data')
-rw-r--r--src/Data/Array/Mixed.hs14
1 files changed, 5 insertions, 9 deletions
diff --git a/src/Data/Array/Mixed.hs b/src/Data/Array/Mixed.hs
index f597bbf..2ad1d26 100644
--- a/src/Data/Array/Mixed.hs
+++ b/src/Data/Array/Mixed.hs
@@ -134,16 +134,12 @@ toLinearIdx = \sh i -> fst (go sh i)
in (sz * i + lidx, n * sz)
enumShape :: IxX sh -> [IxX sh]
-enumShape = \sh -> go 0 sh id []
+enumShape = \sh -> go sh id []
where
- go :: Int -> IxX sh -> (IxX sh -> a) -> [a] -> [a]
- go _ IZX _ = id
- go i (n ::@ sh) f
- | i < n = go (i + 1) (n ::@ sh) f . go 0 sh (f . (i ::@))
- | otherwise = id
- go i (n ::? sh) f
- | i < n = go (i + 1) (n ::? sh) f . go 0 sh (f . (i ::?))
- | otherwise = id
+ go :: IxX sh -> (IxX sh -> a) -> [a] -> [a]
+ go IZX f = (f IZX :)
+ go (n ::@ sh) f = foldr (.) id [go sh (f . (i ::@)) | i <- [0 .. n-1]]
+ go (n ::? sh) f = foldr (.) id [go sh (f . (i ::?)) | i <- [0 .. n-1]]
shapeLshape :: IxX sh -> U.ShapeL
shapeLshape IZX = []