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.

1 files changed, 7 insertions, 3 deletions

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.