Generalize mgeneratePrim to potentially avoid @fmap fromIntegral@

1 files changed, 5 insertions, 3 deletions

diff --git a/src/Data/Array/Nested/Mixed.hs b/src/Data/Array/Nested/Mixed.hs
index 515e867..d658ed3 100644
--- a/src/Data/Array/Nested/Mixed.hs
+++ b/src/Data/Array/Nested/Mixed.hs
@@ -751,9 +751,11 @@ mgenerate sh f = case shxEnum sh of
 
 -- | An optimized special case of `mgenerate', where the function results
 -- are of a primitive type and so there's not need to verify the shapes
--- of them all are equal.
-mgeneratePrim :: forall sh a. PrimElt a
-              => IShX sh -> (IIxX sh -> a) -> Mixed sh a
+-- of them all are equal. This is also generalized to aribitrary @Num@ index
+-- type compared to @mgenerate@.
+{-# INLINE mgeneratePrim #-}
+mgeneratePrim :: forall sh a i. (PrimElt a, Num i)
+              => IShX sh -> (IxX sh i -> a) -> Mixed sh a
 mgeneratePrim sh f =
   let g i = f (ixxFromLinear sh i)
   in mfromVector sh $ VS.generate (shxSize sh) g