diff options
| author | Mikolaj Konarski <mikolaj.konarski@funktory.com> | 2025-12-02 13:16:33 +0100 |
|---|---|---|
| committer | Mikolaj Konarski <mikolaj.konarski@funktory.com> | 2025-12-02 13:16:33 +0100 |
| commit | 9f47aa6a2bcd772388a5d5150ca7254e4eb95bc2 (patch) | |
| tree | 4c3a1b8a7b1a734e83f161f2b1be58ce4470cfa3 /src/Data/Array/Nested/Mixed.hs | |
| parent | ba5a31c976f80421464af1af8d6ab1e2a154cd83 (diff) | |
Generalize mgeneratePrim to potentially avoid @fmap fromIntegral@
Diffstat (limited to 'src/Data/Array/Nested/Mixed.hs')
| -rw-r--r-- | src/Data/Array/Nested/Mixed.hs | 8 |
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 |