diff options
Diffstat (limited to 'src/Data/Array/Nested')
| -rw-r--r-- | src/Data/Array/Nested/Mixed/Shape.hs | 10 | ||||
| -rw-r--r-- | src/Data/Array/Nested/Ranked/Shape.hs | 10 | ||||
| -rw-r--r-- | src/Data/Array/Nested/Shaped/Shape.hs | 10 |
3 files changed, 21 insertions, 9 deletions
diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs index 8777739..0464f1f 100644 --- a/src/Data/Array/Nested/Mixed/Shape.hs +++ b/src/Data/Array/Nested/Mixed/Shape.hs @@ -472,11 +472,15 @@ shxSplitApp _ ZKX idx = (ZSX, idx) shxSplitApp p (_ :!% ssh) (i :$% idx) = first (i :$%) (shxSplitApp p ssh idx) shxEnum :: IShX sh -> [IIxX sh] -shxEnum = \sh -> go sh id [] +shxEnum = shxEnum' + +{-# INLINABLE shxEnum' #-} -- ensure this can be specialised at use site +shxEnum' :: Num i => IShX sh -> [IxX sh i] +shxEnum' = \sh -> go sh id [] where - go :: IShX sh -> (IIxX sh -> a) -> [a] -> [a] + go :: Num i => IShX sh -> (IxX sh i -> a) -> [a] -> [a] go ZSX f = (f ZIX :) - go (n :$% sh) f = foldr (.) id [go sh (f . (i :.%)) | i <- [0 .. fromSMayNat' n - 1]] + go (n :$% sh) f = foldr (.) id [go sh (f . (fromIntegral i :.%)) | i <- [0 .. fromSMayNat' n - 1]] shxCast :: StaticShX sh' -> IShX sh -> Maybe (IShX sh') shxCast ZKX ZSX = Just ZSX diff --git a/src/Data/Array/Nested/Ranked/Shape.hs b/src/Data/Array/Nested/Ranked/Shape.hs index a1fa247..200fac5 100644 --- a/src/Data/Array/Nested/Ranked/Shape.hs +++ b/src/Data/Array/Nested/Ranked/Shape.hs @@ -331,11 +331,15 @@ shrPermutePrefix :: forall n i. [Int] -> ShR n i -> ShR n i shrPermutePrefix = coerce (listrPermutePrefix @i) shrEnum :: IShR sh -> [IIxR sh] -shrEnum = \sh -> go sh id [] +shrEnum = shrEnum' + +{-# INLINABLE shrEnum' #-} -- ensure this can be specialised at use site +shrEnum' :: Num i => IShR sh -> [IxR sh i] +shrEnum' = \sh -> go sh id [] where - go :: IShR sh -> (IIxR sh -> a) -> [a] -> [a] + go :: Num i => IShR sh -> (IxR sh i -> a) -> [a] -> [a] go ZSR f = (f ZIR :) - go (n :$: sh) f = foldr (.) id [go sh (f . (i :.:)) | i <- [0 .. n - 1]] + go (n :$: sh) f = foldr (.) id [go sh (f . (fromIntegral i :.:)) | i <- [0 .. n - 1]] -- | Untyped: length is checked at runtime. diff --git a/src/Data/Array/Nested/Shaped/Shape.hs b/src/Data/Array/Nested/Shaped/Shape.hs index 6eaec10..582cc50 100644 --- a/src/Data/Array/Nested/Shaped/Shape.hs +++ b/src/Data/Array/Nested/Shaped/Shape.hs @@ -390,11 +390,15 @@ shsFromIxS :: IxS sh i -> ShS sh shsFromIxS (IxS l) = shsFromListS l shsEnum :: ShS sh -> [IIxS sh] -shsEnum = \sh -> go sh id [] +shsEnum = shsEnum' + +{-# INLINABLE shsEnum' #-} -- ensure this can be specialised at use site +shsEnum' :: Num i => ShS sh -> [IxS sh i] +shsEnum' = \sh -> go sh id [] where - go :: ShS sh -> (IIxS sh -> a) -> [a] -> [a] + go :: Num i => ShS sh -> (IxS sh i -> a) -> [a] -> [a] go ZSS f = (f ZIS :) - go (n :$$ sh) f = foldr (.) id [go sh (f . (i :.$)) | i <- [0 .. fromSNat' n - 1]] + go (n :$$ sh) f = foldr (.) id [go sh (f . (fromIntegral i :.$)) | i <- [0 .. fromSNat' n - 1]] -- | Untyped: length is checked at runtime. |
