Add iconcatMap, concatMapM and iconcatMapM

vector/src/Data/Vector.hs

@@ -94,11 +94,11 @@ module Data.Vector (
   indexed,
 
   -- ** Mapping
-  map, imap, concatMap,
+  map, imap, concatMap, iconcatMap,
 
   -- ** Monadic mapping
   mapM, imapM, mapM_, imapM_, forM, forM_,
-  iforM, iforM_,
+  iforM, iforM_, concatMapM, iconcatMapM,
 
   -- ** Zipping
   zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
@@ -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
 -- ---------------
 
@@ -1151,6 +1158,22 @@ iforM_ :: Monad m => Vector a -> (Int -> a -> m b) -> m ()
 {-# INLINE iforM_ #-}
 iforM_ = G.iforM_
 
+-- | Apply the monadic action to all elements of the vector, yielding a vector of results, and
+-- concatenate the results.
+--
+-- @since 0.13.3.0
+concatMapM :: (Monad m) => (a -> m (Vector b)) -> Vector a -> m (Vector b)
+{-# INLINE concatMapM #-}
+concatMapM = G.concatMapM
+
+-- | Apply the monadic action to every element of the vector and its index, yielding a vector of
+-- results, and concatenate the results.
+--
+-- @since 0.13.3.0
+iconcatMapM :: (Monad m) => (Int -> a -> m (Vector b)) -> Vector a -> m (Vector b)
+{-# INLINE iconcatMapM #-}
+iconcatMapM = G.iconcatMapM
+
 -- Zipping
 -- -------
 

vector/src/Data/Vector/Generic.hs

@@ -84,11 +84,11 @@ module Data.Vector.Generic (
   indexed,
 
   -- ** Mapping
-  map, imap, concatMap,
+  map, imap, concatMap, iconcatMap,
 
   -- ** Monadic mapping
   mapM, imapM, mapM_, imapM_, forM, forM_,
-  iforM, iforM_,
+  iforM, iforM_, concatMapM, iconcatMapM,
 
   -- ** Zipping
   zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
@@ -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
 -- ---------------
 
@@ -1167,6 +1177,29 @@ iforM_ :: (Monad m, Vector v a) => v a -> (Int -> a -> m b) -> m ()
 {-# INLINE iforM_ #-}
 iforM_ as f = imapM_ f as
 
+-- | Apply the monadic action to all elements of the vector, yielding a vector of results, and
+-- concatenate the results.
+--
+-- @since 0.13.3.0
+concatMapM :: (Monad m, Vector v a, Vector v b) => (a -> m (v b)) -> v a -> m (v b)
+{-# INLINE concatMapM #-}
+concatMapM f = unstreamM
+               . MBundle.concatVectors
+               . MBundle.mapM f
+               . MBundle.fromVector
+
+-- | Apply the monadic action to every element of the vector and its index, yielding a vector of
+-- results, and concatenate the results.
+--
+-- @since 0.13.3.0
+iconcatMapM :: (Monad m, Vector v a, Vector v b) => (Int -> a -> m (v b)) -> v a -> m (v b)
+{-# INLINE iconcatMapM #-}
+iconcatMapM f = unstreamM
+                . MBundle.concatVectors
+                . MBundle.mapM (uncurry f)
+                . MBundle.indexed
+                . MBundle.fromVector
+
 -- Zipping
 -- -------
 

vector/src/Data/Vector/Primitive.hs

@@ -84,11 +84,11 @@ module Data.Vector.Primitive (
   -- * Elementwise operations
 
   -- ** Mapping
-  map, imap, concatMap,
+  map, imap, concatMap, iconcatMap,
 
   -- ** Monadic mapping
   mapM, imapM, mapM_, imapM_, forM, forM_,
-  iforM, iforM_,
+  iforM, iforM_, concatMapM, iconcatMapM,
 
   -- ** Zipping
   zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
@@ -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
 -- ---------------
 
@@ -942,6 +949,22 @@ iforM_ :: (Monad m, Prim a) => Vector a -> (Int -> a -> m b) -> m ()
 {-# INLINE iforM_ #-}
 iforM_ = G.iforM_
 
+-- | Apply the monadic action to all elements of the vector, yielding a vector of results, and
+-- concatenate the results.
+--
+-- @since 0.13.3.0
+concatMapM :: (Monad m, Prim a, Prim b) => (a -> m (Vector b)) -> Vector a -> m (Vector b)
+{-# INLINE concatMapM #-}
+concatMapM = G.concatMapM
+
+-- | Apply the monadic action to every element of the vector and its index, yielding a vector of
+-- results, and concatenate the results.
+--
+-- @since 0.13.3.0
+iconcatMapM :: (Monad m, Prim a, Prim b) => (Int -> a -> m (Vector b)) -> Vector a -> m (Vector b)
+{-# INLINE iconcatMapM #-}
+iconcatMapM = G.iconcatMapM
+
 -- Zipping
 -- -------
 

vector/src/Data/Vector/Storable.hs

@@ -81,11 +81,11 @@ module Data.Vector.Storable (
   -- * Elementwise operations
 
   -- ** Mapping
-  map, imap, concatMap,
+  map, imap, concatMap, iconcatMap,
 
   -- ** Monadic mapping
   mapM, imapM, mapM_, imapM_, forM, forM_,
-  iforM, iforM_,
+  iforM, iforM_, concatMapM, iconcatMapM,
 
   -- ** Zipping
   zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
@@ -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
 -- ---------------
 
@@ -953,6 +960,22 @@ iforM_ :: (Monad m, Storable a) => Vector a -> (Int -> a -> m b) -> m ()
 {-# INLINE iforM_ #-}
 iforM_ = G.iforM_
 
+-- | Apply the monadic action to all elements of the vector, yielding a vector of results, and
+-- concatenate the results.
+--
+-- @since 0.13.3.0
+concatMapM :: (Monad m, Storable a, Storable b) => (a -> m (Vector b)) -> Vector a -> m (Vector b)
+{-# INLINE concatMapM #-}
+concatMapM = G.concatMapM
+
+-- | Apply the monadic action to every element of the vector and its index, yielding a vector of
+-- results, and concatenate the results.
+--
+-- @since 0.13.3.0
+iconcatMapM :: (Monad m, Storable a, Storable b) => (Int -> a -> m (Vector b)) -> Vector a -> m (Vector b)
+{-# INLINE iconcatMapM #-}
+iconcatMapM = G.iconcatMapM
+
 -- Zipping
 -- -------
 

vector/src/Data/Vector/Strict.hs

@@ -92,11 +92,11 @@ module Data.Vector.Strict (
   indexed,
 
   -- ** Mapping
-  map, imap, concatMap,
+  map, imap, concatMap, iconcatMap,
 
   -- ** Monadic mapping
   mapM, imapM, mapM_, imapM_, forM, forM_,
-  iforM, iforM_,
+  iforM, iforM_, concatMapM, iconcatMapM,
 
   -- ** Zipping
   zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
@@ -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
 -- ---------------
 
@@ -1242,6 +1249,22 @@ iforM_ :: Monad m => Vector a -> (Int -> a -> m b) -> m ()
 {-# INLINE iforM_ #-}
 iforM_ = G.iforM_
 
+-- | Apply the monadic action to all elements of the vector, yielding a vector of results, and
+-- concatenate the results.
+--
+-- @since 0.13.3.0
+concatMapM :: (Monad m) => (a -> m (Vector b)) -> Vector a -> m (Vector b)
+{-# INLINE concatMapM #-}
+concatMapM = G.concatMapM
+
+-- | Apply the monadic action to every element of the vector and its index, yielding a vector of
+-- results, and concatenate the results.
+--
+-- @since 0.13.3.0
+iconcatMapM :: (Monad m) => (Int -> a -> m (Vector b)) -> Vector a -> m (Vector b)
+{-# INLINE iconcatMapM #-}
+iconcatMapM = G.iconcatMapM
+
 -- Zipping
 -- -------
 

vector/src/Data/Vector/Unboxed.hs

@@ -132,11 +132,11 @@ module Data.Vector.Unboxed (
   indexed,
 
   -- ** Mapping
-  map, imap, concatMap,
+  map, imap, concatMap, iconcatMap,
 
   -- ** Monadic mapping
   mapM, imapM, mapM_, imapM_, forM, forM_,
-  iforM, iforM_,
+  iforM, iforM_, concatMapM, iconcatMapM,
 
   -- ** Zipping
   zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
@@ -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
 -- ---------------
 
@@ -987,6 +994,22 @@ iforM_ :: (Monad m, Unbox a) => Vector a -> (Int -> a -> m b) -> m ()
 {-# INLINE iforM_ #-}
 iforM_ = G.iforM_
 
+-- | Apply the monadic action to all elements of the vector, yielding a vector of results, and
+-- concatenate the results.
+--
+-- @since 0.13.3.0
+concatMapM :: (Monad m, Unbox a, Unbox b) => (a -> m (Vector b)) -> Vector a -> m (Vector b)
+{-# INLINE concatMapM #-}
+concatMapM = G.concatMapM
+
+-- | Apply the monadic action to every element of the vector and its index, yielding a vector of
+-- results, and concatenate the results.
+--
+-- @since 0.13.3.0
+iconcatMapM :: (Monad m, Unbox a, Unbox b) => (Int -> a -> m (Vector b)) -> Vector a -> m (Vector b)
+{-# INLINE iconcatMapM #-}
+iconcatMapM = G.iconcatMapM
+
 -- Zipping
 -- -------
 

vector/tests/Tests/Vector/Property.hs

@@ -147,11 +147,11 @@ 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_,
-        'prop_imapM, 'prop_imapM_,
+        'prop_imapM, 'prop_imapM_, 'prop_concatMapM, 'prop_iconcatMapM,
 
         -- Zipping
         'prop_zipWith, 'prop_zipWith3,
@@ -313,6 +313,17 @@ testPolymorphicFunctions _ = $(testProperties [
             = V.imapM `eq` imapM
     prop_imapM_ :: P ((Int -> a -> Writer [a] ()) -> v a -> Writer [a] ())
             = V.imapM_ `eq` imapM_
+
+    prop_concatMapM    = forAll arbitrary $ \xs ->
+                        forAll (sized (\n -> resize (n `div` V.length xs) arbitrary)) $ \f -> unP prop f xs
+      where
+        prop :: P ((a -> Identity (v a)) -> v a -> Identity (v a)) = V.concatMapM `eq` concatMapM
+
+    prop_iconcatMapM   = forAll arbitrary $ \xs ->
+                        forAll (sized (\n -> resize (n `div` V.length xs) arbitrary)) $ \f -> unP prop f xs
+      where
+        prop :: P ((Int -> a -> Identity (v a)) -> v a -> Identity (v a)) = V.iconcatMapM `eq` iconcatMapM
+
     prop_izipWith :: P ((Int -> a -> a -> a) -> v a -> v a -> v a) = V.izipWith `eq` izipWith
     prop_zipWithM :: P ((a -> a -> Identity a) -> v a -> v a -> Identity (v a))
             = V.zipWithM `eq` zipWithM
@@ -435,6 +446,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,15 @@ 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
+
+concatMapM :: Monad m => (a -> m [a]) -> [a] -> m [a]
+concatMapM f = fmap concat . mapM f
+
+iconcatMapM :: Monad m => (Int -> a -> m [a]) -> [a] -> m [a]
+iconcatMapM f = fmap concat . withIndexFirst mapM f
+
 izipWith :: (Int -> a -> a -> a) -> [a] -> [a] -> [a]
 izipWith = withIndexFirst zipWith