diff options
| author | Tom Smeding <t.j.smeding@uu.nl> | 2024-04-03 15:06:29 +0200 |
|---|---|---|
| committer | Tom Smeding <t.j.smeding@uu.nl> | 2024-04-03 15:06:29 +0200 |
| commit | 6520c4e4bc39b7315f4f9416fe68b883cfdde8ec (patch) | |
| tree | 8315b318c6be640ab46b306270aa644e7623a472 /src | |
| parent | 6f912510c594de511bc92b7f47d0d405b64aaa1e (diff) | |
Fix enumShape
Diffstat (limited to 'src')
| -rw-r--r-- | src/Data/Array/Mixed.hs | 14 |
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 = [] |