diff options
| -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 = [] |