From 9f47aa6a2bcd772388a5d5150ca7254e4eb95bc2 Mon Sep 17 00:00:00 2001 From: Mikolaj Konarski Date: Tue, 2 Dec 2025 13:16:33 +0100 Subject: Generalize mgeneratePrim to potentially avoid @fmap fromIntegral@ --- src/Data/Array/Nested/Mixed.hs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) (limited to 'src/Data/Array/Nested/Mixed.hs') 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 -- cgit v1.2.3-70-g09d2