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