diff options
Diffstat (limited to 'src/Data/Array/Nested/Mixed.hs')
| -rw-r--r-- | src/Data/Array/Nested/Mixed.hs | 13 |
1 files changed, 6 insertions, 7 deletions
diff --git a/src/Data/Array/Nested/Mixed.hs b/src/Data/Array/Nested/Mixed.hs index d658ed3..182943d 100644 --- a/src/Data/Array/Nested/Mixed.hs +++ b/src/Data/Array/Nested/Mixed.hs @@ -728,6 +728,8 @@ msize = shxSize . mshape -- the entire hierarchy (after distributing out tuples) must be a rectangular -- array. The type of 'mgenerate' allows this requirement to be broken very -- easily, hence the runtime check. +-- +-- If your element type @a@ is a scalar, use the faster 'mgeneratePrim'. mgenerate :: forall sh a. KnownElt a => IShX sh -> (IIxX sh -> a) -> Mixed sh a mgenerate sh f = case shxEnum sh of [] -> memptyArrayUnsafe sh @@ -739,9 +741,6 @@ mgenerate sh f = case shxEnum sh of else runST $ do vecs <- mvecsUnsafeNew sh firstelem mvecsWrite sh firstidx firstelem vecs - -- This is likely fine if @a@ is big, but if @a@ is a scalar - -- this array copying is inefficient so it's better to use - -- the @mgeneratePrim@ below. forM_ restidxs $ \idx -> do let val = f idx when (not (mshapeTreeEq (Proxy @a) (mshapeTree val) shapetree)) $ @@ -749,10 +748,10 @@ mgenerate sh f = case shxEnum sh of mvecsWrite sh idx val vecs mvecsFreeze sh vecs --- | 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. This is also generalized to aribitrary @Num@ index --- type compared to @mgenerate@. +-- | An optimized special case of 'mgenerate', where the function results +-- are of a primitive type and so there's not need to check that all shapes +-- are equal. This is also generalized to an arbitrary @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 |