diff options
Diffstat (limited to 'src/Data/Array/Nested/Mixed/Shape.hs')
| -rw-r--r-- | src/Data/Array/Nested/Mixed/Shape.hs | 10 |
1 files changed, 7 insertions, 3 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 |