aboutsummaryrefslogtreecommitdiff
path: root/src/Data/Array/Nested/Shaped/Shape.hs
diff options
context:
space:
mode:
authorTom Smeding <tom@tomsmeding.com>2025-11-18 12:57:17 +0100
committerTom Smeding <tom@tomsmeding.com>2025-11-18 12:57:17 +0100
commita158f928c6d617fa950f74c517e94c5062422efd (patch)
treea82cc806c9b2cfd45183e6b65f98c377575af41a /src/Data/Array/Nested/Shaped/Shape.hs
parente53b10d35bdb414520ac1d5b00307c4c62a9e1db (diff)
Generalise sh?Enum functions to arbitrary NumHEADmaster
The explicit usage at Int in the non-primed versions of the functions results in an auto-generated SPEC rule at Int, so even without any SPECIALIZE pragmas, using sh?Enum' at Int should not (need to) result in multiple specialisations.
Diffstat (limited to 'src/Data/Array/Nested/Shaped/Shape.hs')
-rw-r--r--src/Data/Array/Nested/Shaped/Shape.hs10
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.