Add iconcatMap (#531)

Author: toyboot4e

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
 -- ---------------
 

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
 -- ---------------
 

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
 -- ---------------
 

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
 -- ---------------
 

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
 -- ---------------
 

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
 -- ---------------
 

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)
 

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