diff --git a/vector/src/Data/Vector.hs b/vector/src/Data/Vector.hs index 991dae5a..9517d4fc 100644 --- a/vector/src/Data/Vector.hs +++ b/vector/src/Data/Vector.hs @@ -94,7 +94,7 @@ module Data.Vector ( indexed, -- ** Mapping - map, imap, concatMap, + map, imap, concatMap, iconcatMap, -- ** Monadic mapping mapM, imapM, mapM_, imapM_, forM, forM_, @@ -1096,6 +1096,13 @@ concatMap :: (a -> Vector b) -> Vector a -> Vector b {-# INLINE concatMap #-} concatMap = G.concatMap +-- | Map a function to every element of a vector and its index, and concatenate the results. +-- +-- @since 0.13.3.0 +iconcatMap :: (Int -> a -> Vector b) -> Vector a -> Vector b +{-# INLINE iconcatMap #-} +iconcatMap = G.iconcatMap + -- Monadic mapping -- --------------- diff --git a/vector/src/Data/Vector/Generic.hs b/vector/src/Data/Vector/Generic.hs index 35acd0fb..c2114f7a 100644 --- a/vector/src/Data/Vector/Generic.hs +++ b/vector/src/Data/Vector/Generic.hs @@ -84,7 +84,7 @@ module Data.Vector.Generic ( indexed, -- ** Mapping - map, imap, concatMap, + map, imap, concatMap, iconcatMap, -- ** Monadic mapping mapM, imapM, mapM_, imapM_, forM, forM_, @@ -1112,6 +1112,16 @@ concatMap f = unstream . Bundle.map f . stream +-- | Apply a function to every element of a vector and its index, and concatenate the results. +-- +-- @since 0.13.3.0 +iconcatMap :: (Vector v a, Vector v b) => (Int -> a -> v b) -> v a -> v b +{-# INLINE iconcatMap #-} +iconcatMap f = unstream + . Bundle.concatVectors + . Bundle.inplace (S.map (uncurry f) . S.indexed) id + . stream + -- Monadic mapping -- --------------- diff --git a/vector/src/Data/Vector/Primitive.hs b/vector/src/Data/Vector/Primitive.hs index 2d67fba7..95d0d790 100644 --- a/vector/src/Data/Vector/Primitive.hs +++ b/vector/src/Data/Vector/Primitive.hs @@ -84,7 +84,7 @@ module Data.Vector.Primitive ( -- * Elementwise operations -- ** Mapping - map, imap, concatMap, + map, imap, concatMap, iconcatMap, -- ** Monadic mapping mapM, imapM, mapM_, imapM_, forM, forM_, @@ -882,6 +882,13 @@ concatMap :: (Prim a, Prim b) => (a -> Vector b) -> Vector a -> Vector b {-# INLINE concatMap #-} concatMap = G.concatMap +-- | Apply a function to every element of a vector and its index, and concatenate the results. +-- +-- @since 0.13.3.0 +iconcatMap :: (Prim a, Prim b) => (Int -> a -> Vector b) -> Vector a -> Vector b +{-# INLINE iconcatMap #-} +iconcatMap = G.iconcatMap + -- Monadic mapping -- --------------- diff --git a/vector/src/Data/Vector/Storable.hs b/vector/src/Data/Vector/Storable.hs index 88f2332c..65565b05 100644 --- a/vector/src/Data/Vector/Storable.hs +++ b/vector/src/Data/Vector/Storable.hs @@ -81,7 +81,7 @@ module Data.Vector.Storable ( -- * Elementwise operations -- ** Mapping - map, imap, concatMap, + map, imap, concatMap, iconcatMap, -- ** Monadic mapping mapM, imapM, mapM_, imapM_, forM, forM_, @@ -893,6 +893,13 @@ concatMap :: (Storable a, Storable b) => (a -> Vector b) -> Vector a -> Vector b {-# INLINE concatMap #-} concatMap = G.concatMap +-- | Apply a function to every element of a vector and its index, and concatenate the results. +-- +-- @since 0.13.3.0 +iconcatMap :: (Storable a, Storable b) => (Int -> a -> Vector b) -> Vector a -> Vector b +{-# INLINE iconcatMap #-} +iconcatMap = G.iconcatMap + -- Monadic mapping -- --------------- diff --git a/vector/src/Data/Vector/Strict.hs b/vector/src/Data/Vector/Strict.hs index faedc2e5..01143b4b 100644 --- a/vector/src/Data/Vector/Strict.hs +++ b/vector/src/Data/Vector/Strict.hs @@ -92,7 +92,7 @@ module Data.Vector.Strict ( indexed, -- ** Mapping - map, imap, concatMap, + map, imap, concatMap, iconcatMap, -- ** Monadic mapping mapM, imapM, mapM_, imapM_, forM, forM_, @@ -1175,6 +1175,13 @@ concatMap :: (a -> Vector b) -> Vector a -> Vector b {-# INLINE concatMap #-} concatMap = G.concatMap +-- | Apply a function to every element of a vector and its index, and concatenate the results. +-- +-- @since 0.13.3.0 +iconcatMap :: (Int -> a -> Vector b) -> Vector a -> Vector b +{-# INLINE iconcatMap #-} +iconcatMap = G.iconcatMap + -- Monadic mapping -- --------------- diff --git a/vector/src/Data/Vector/Unboxed.hs b/vector/src/Data/Vector/Unboxed.hs index 673e4946..f84b38f2 100644 --- a/vector/src/Data/Vector/Unboxed.hs +++ b/vector/src/Data/Vector/Unboxed.hs @@ -132,7 +132,7 @@ module Data.Vector.Unboxed ( indexed, -- ** Mapping - map, imap, concatMap, + map, imap, concatMap, iconcatMap, -- ** Monadic mapping mapM, imapM, mapM_, imapM_, forM, forM_, @@ -931,6 +931,13 @@ concatMap :: (Unbox a, Unbox b) => (a -> Vector b) -> Vector a -> Vector b {-# INLINE concatMap #-} concatMap = G.concatMap +-- | Apply a function to every element of a vector and its index, and concatenate the results. +-- +-- @since 0.13.3.0 +iconcatMap :: (Unbox a, Unbox b) => (Int -> a -> Vector b) -> Vector a -> Vector b +{-# INLINE iconcatMap #-} +iconcatMap = G.iconcatMap + -- Monadic mapping -- --------------- diff --git a/vector/tests/Tests/Vector/Property.hs b/vector/tests/Tests/Vector/Property.hs index dc8f0b63..135818d7 100644 --- a/vector/tests/Tests/Vector/Property.hs +++ b/vector/tests/Tests/Vector/Property.hs @@ -147,7 +147,7 @@ testPolymorphicFunctions _ = $(testProperties [ {- 'prop_unsafeBackpermute, -} -- Mapping - 'prop_map, 'prop_imap, 'prop_concatMap, + 'prop_map, 'prop_imap, 'prop_concatMap, 'prop_iconcatMap, -- Monadic mapping 'prop_mapM, 'prop_mapM_, 'prop_forM, 'prop_forM_, @@ -435,6 +435,11 @@ testPolymorphicFunctions _ = $(testProperties [ where prop :: P ((a -> v a) -> v a -> v a) = V.concatMap `eq` concatMap + prop_iconcatMap = forAll arbitrary $ \xs -> + forAll (sized (\n -> resize (n `div` V.length xs) arbitrary)) $ \f -> unP prop f xs + where + prop :: P ((Int -> a -> v a) -> v a -> v a) = V.iconcatMap `eq` iconcatMap + prop_uniq :: P (v a -> v a) = V.uniq `eq` (map head . group) diff --git a/vector/tests/Utilities.hs b/vector/tests/Utilities.hs index 5aa9e8b0..77e2be86 100644 --- a/vector/tests/Utilities.hs +++ b/vector/tests/Utilities.hs @@ -286,6 +286,9 @@ imapM = withIndexFirst mapM imapM_ :: Monad m => (Int -> a -> m b) -> [a] -> m () imapM_ = withIndexFirst mapM_ +iconcatMap :: (Int -> a -> [a]) -> [a] -> [a] +iconcatMap f = concat . withIndexFirst map f + izipWith :: (Int -> a -> a -> a) -> [a] -> [a] -> [a] izipWith = withIndexFirst zipWith