From 6520c4e4bc39b7315f4f9416fe68b883cfdde8ec Mon Sep 17 00:00:00 2001
From: Tom Smeding <t.j.smeding@uu.nl>
Date: Wed, 3 Apr 2024 15:06:29 +0200
Subject: Fix enumShape
---
src/Data/Array/Mixed.hs | 14 +++++---------
1 file changed, 5 insertions(+), 9 deletions(-)
(limited to 'src')
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 = []
--
cgit v1.2.3-70-g09d2