From 6520c4e4bc39b7315f4f9416fe68b883cfdde8ec Mon Sep 17 00:00:00 2001 From: Tom Smeding 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