Add groupBy and group for vectors (#427)

This adds `groupBy` and `group` for generic vectors, returning a list of
slices of the input vector.

Specialized versions of groupBy and group are added for:

- Data.Vector
- Data.Vector.Unboxed
- Data.Vector.Storable
- Data.Vector.Primitive

These functions are all `O(n)`, and do not fuse, but because they
return slices of the input, they do not copy.

Co-authored-by: Alex Kalderimis <[email protected]>

Author: alexkalderimis

vector/changelog.md

@@ -44,6 +44,9 @@
    [#382](https://github.com/haskell/vector/pull/382)
 * Remove redundant `Storable` constraints on to/from `ForeignPtr` conversions
 * Add `unsafeCast` to `Primitive` vectors
+* Add `groupBy` and `group` for `Data.Vector.Generic` and the specialized
+  version in `Data.Vector`, `Data.Vector.Unboxed`, `Data.Vector.Storable` and
+  `Data.Vector.Primitive`.
 
 # Changes in version 0.12.3.1
 

vector/src/Data/Vector.hs

@@ -119,7 +119,7 @@ module Data.Vector (
   takeWhile, dropWhile,
 
   -- ** Partitioning
-  partition, unstablePartition, partitionWith, span, break,
+  partition, unstablePartition, partitionWith, span, break, groupBy, group,
 
   -- ** Searching
   elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
@@ -1410,6 +1410,46 @@ break :: (a -> Bool) -> Vector a -> (Vector a, Vector a)
 {-# INLINE break #-}
 break = G.break
 
+-- | /O(n)/ Split a vector into a list of slices, using a predicate function.
+--
+-- The concatenation of this list of slices is equal to the argument vector,
+-- and each slice contains only equal elements, as determined by the equality
+-- predicate function.
+--
+-- Does not fuse.
+--
+-- >>> import qualified Data.Vector as V
+-- >>> import           Data.Char (isUpper)
+-- >>> V.groupBy (\a b -> isUpper a == isUpper b) (V.fromList "Mississippi River")
+-- ["M","ississippi ","R","iver"]
+--
+-- See also 'Data.List.groupBy', 'group'.
+--
+-- @since 0.13.0.1
+groupBy :: (a -> a -> Bool) -> Vector a -> [Vector a]
+{-# INLINE groupBy #-}
+groupBy = G.groupBy
+
+-- | /O(n)/ Split a vector into a list of slices of the input vector.
+--
+-- The concatenation of this list of slices is equal to the argument vector,
+-- and each slice contains only equal elements.
+--
+-- Does not fuse.
+--
+-- This is the equivalent of 'groupBy (==)'.
+--
+-- >>> import qualified Data.Vector as V
+-- >>> V.group (V.fromList "Mississippi")
+-- ["M","i","ss","i","ss","i","pp","i"]
+--
+-- See also 'Data.List.group'.
+--
+-- @since 0.13.0.1
+group :: Eq a => Vector a -> [Vector a]
+{-# INLINE group #-}
+group = G.groupBy (==)
+
 -- Searching
 -- ---------
 

vector/src/Data/Vector/Generic.hs

@@ -107,7 +107,7 @@ module Data.Vector.Generic (
   takeWhile, dropWhile,
 
   -- ** Partitioning
-  partition, partitionWith, unstablePartition, span, break,
+  partition, partitionWith, unstablePartition, span, break, groupBy, group,
 
   -- ** Searching
   elem, notElem, find, findIndex, findIndexR, findIndices, elemIndex, elemIndices,
@@ -1552,6 +1552,47 @@ break f xs = case findIndex f xs of
                Just i  -> (unsafeSlice 0 i xs, unsafeSlice i (length xs - i) xs)
                Nothing -> (xs, empty)
 
+-- | /O(n)/ Split a vector into a list of slices.
+--
+-- The concatenation of this list of slices is equal to the argument vector,
+-- and each slice contains only equal elements, as determined by the equality
+-- predicate function.
+--
+-- >>> import qualified Data.Vector as V
+-- >>> import           Data.Char (isUpper)
+-- >>> V.groupBy (\a b -> isUpper a == isUpper b) (V.fromList "Mississippi River")
+-- ["M","ississippi ","R","iver"]
+--
+-- See also 'Data.List.groupBy'.
+--
+-- @since 0.13.0.1
+{-# INLINE groupBy #-}
+groupBy :: (Vector v a) => (a -> a -> Bool) -> v a -> [v a]
+groupBy _ v | null v = []
+groupBy f v =
+  let h = unsafeHead v
+      tl = unsafeTail v
+   in case findIndex (not . f h) tl of
+      Nothing -> [v]
+      Just n -> unsafeTake (n + 1) v : groupBy f (unsafeDrop (n + 1) v)
+
+-- | /O(n)/ Split a vector into a list of slices.
+--
+-- The concatenation of this list of slices is equal to the argument vector,
+-- and each slice contains only equal elements.
+--
+-- This is the equivalent of 'groupBy (==)'.
+--
+-- >>> import qualified Data.Vector as V
+-- >>> V.group (V.fromList "Mississippi")
+-- ["M","i","ss","i","ss","i","pp","i"]
+--
+-- See also 'Data.List.group'.
+--
+-- @since 0.13.0.1
+group :: (Vector v a , Eq a) => v a -> [v a]
+{-# INLINE group #-}
+group = groupBy (==)
 
 -- Searching
 -- ---------

vector/src/Data/Vector/Primitive.hs

@@ -104,7 +104,7 @@ module Data.Vector.Primitive (
   takeWhile, dropWhile,
 
   -- ** Partitioning
-  partition, unstablePartition, partitionWith, span, break,
+  partition, unstablePartition, partitionWith, span, break, groupBy, group,
 
   -- ** Searching
   elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
@@ -1169,6 +1169,46 @@ break :: Prim a => (a -> Bool) -> Vector a -> (Vector a, Vector a)
 {-# INLINE break #-}
 break = G.break
 
+-- | /O(n)/ Split a vector into a list of slices, using a predicate function.
+--
+-- The concatenation of this list of slices is equal to the argument vector,
+-- and each slice contains only equal elements, as determined by the equality
+-- predicate function.
+--
+-- Does not fuse.
+--
+-- >>> import qualified Data.Vector.Primitive as VP
+-- >>> import           Data.Char (isUpper)
+-- >>> VP.groupBy (\a b -> isUpper a == isUpper b) (VP.fromList "Mississippi River")
+-- ["M","ississippi ","R","iver"]
+--
+-- See also 'Data.List.groupBy', 'group'.
+--
+-- @since 0.13.0.1
+groupBy :: Prim a => (a -> a -> Bool) -> Vector a -> [Vector a]
+{-# INLINE groupBy #-}
+groupBy = G.groupBy
+
+-- | /O(n)/ Split a vector into a list of slices of the input vector.
+--
+-- The concatenation of this list of slices is equal to the argument vector,
+-- and each slice contains only equal elements.
+--
+-- Does not fuse.
+--
+-- This is the equivalent of 'groupBy (==)'.
+--
+-- >>> import qualified Data.Vector.Primitive as VP
+-- >>> VP.group (VP.fromList "Mississippi")
+-- ["M","i","ss","i","ss","i","pp","i"]
+--
+-- See also 'Data.List.group'.
+--
+-- @since 0.13.0.1
+group :: (Prim a, Eq a) => Vector a -> [Vector a]
+{-# INLINE group #-}
+group = G.groupBy (==)
+
 -- Searching
 -- ---------
 

vector/src/Data/Vector/Storable.hs

@@ -101,7 +101,7 @@ module Data.Vector.Storable (
   takeWhile, dropWhile,
 
   -- ** Partitioning
-  partition, unstablePartition, partitionWith, span, break,
+  partition, unstablePartition, partitionWith, span, break, groupBy, group,
 
   -- ** Searching
   elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
@@ -1190,6 +1190,46 @@ break :: Storable a => (a -> Bool) -> Vector a -> (Vector a, Vector a)
 {-# INLINE break #-}
 break = G.break
 
+-- | /O(n)/ Split a vector into a list of slices, using a predicate function.
+--
+-- The concatenation of this list of slices is equal to the argument vector,
+-- and each slice contains only equal elements, as determined by the equality
+-- predicate function.
+--
+-- Does not fuse.
+--
+-- >>> import qualified Data.Vector.Storable as VS
+-- >>> import           Data.Char (isUpper)
+-- >>> VS.groupBy (\a b -> isUpper a == isUpper b) (VS.fromList "Mississippi River")
+-- ["M","ississippi ","R","iver"]
+--
+-- See also 'Data.List.groupBy', 'group'.
+--
+-- @since 0.13.0.1
+groupBy :: Storable a => (a -> a -> Bool) -> Vector a -> [Vector a]
+{-# INLINE groupBy #-}
+groupBy = G.groupBy
+
+-- | /O(n)/ Split a vector into a list of slices of the input vector.
+--
+-- The concatenation of this list of slices is equal to the argument vector,
+-- and each slice contains only equal elements.
+--
+-- Does not fuse.
+--
+-- This is the equivalent of 'groupBy (==)'.
+--
+-- >>> import qualified Data.Vector.Storable as VS
+-- >>> VS.group (VS.fromList "Mississippi")
+-- ["M","i","ss","i","ss","i","pp","i"]
+--
+-- See also 'Data.List.group'.
+--
+-- @since 0.13.0.1
+group :: (Storable a, Eq a) => Vector a -> [Vector a]
+{-# INLINE group #-}
+group = G.groupBy (==)
+
 -- Searching
 -- ---------
 

vector/src/Data/Vector/Unboxed.hs

@@ -146,7 +146,7 @@ module Data.Vector.Unboxed (
   takeWhile, dropWhile,
 
   -- ** Partitioning
-  partition, unstablePartition, partitionWith, span, break,
+  partition, unstablePartition, partitionWith, span, break, groupBy, group,
 
   -- ** Searching
   elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
@@ -1182,6 +1182,44 @@ break :: Unbox a => (a -> Bool) -> Vector a -> (Vector a, Vector a)
 {-# INLINE break #-}
 break = G.break
 
+-- | /O(n)/ Split a vector into a list of slices, using a predicate function.
+--
+-- The concatenation of this list of slices is equal to the argument vector,
+-- and each slice contains only equal elements, as determined by the equality
+-- predicate function.
+--
+-- Does not fuse.
+--
+-- >>> import qualified Data.Vector.Unboxed as VU
+-- >>> import           Data.Char (isUpper)
+-- >>> VU.groupBy (\a b -> isUpper a == isUpper b) (VU.fromList "Mississippi River")
+-- ["M","ississippi ","R","iver"]
+--
+-- See also 'Data.List.groupBy', 'group'.
+--
+-- @since 0.13.0.1
+groupBy :: Unbox a => (a -> a -> Bool) -> Vector a -> [Vector a]
+groupBy = G.groupBy
+
+-- | /O(n)/ Split a vector into a list of slices of the input vector.
+--
+-- The concatenation of this list of slices is equal to the argument vector,
+-- and each slice contains only equal elements.
+--
+-- Does not fuse.
+--
+-- This is the equivalent of 'groupBy (==)'.
+--
+-- >>> import qualified Data.Vector.Unboxed as VU
+-- >>> VU.group (VU.fromList "Mississippi")
+-- ["M","i","ss","i","ss","i","pp","i"]
+--
+-- See also 'Data.List.group'.
+--
+-- @since 0.13.0.1
+group :: (Unbox a, Eq a) => Vector a -> [Vector a]
+group = G.groupBy (==)
+
 -- Searching
 -- ---------
 

vector/tests/Tests/Vector/Property.hs

@@ -169,6 +169,7 @@ testPolymorphicFunctions _ = $(testProperties [
         'prop_partition, {- 'prop_unstablePartition, -}
         'prop_partitionWith,
         'prop_span, 'prop_break,
+        'prop_groupBy,
 
         -- Searching
         'prop_elem, 'prop_notElem,
@@ -334,6 +335,7 @@ testPolymorphicFunctions _ = $(testProperties [
       = V.partitionWith `eq` partitionWith
     prop_span :: P ((a -> Bool) -> v a -> (v a, v a)) = V.span `eq` span
     prop_break :: P ((a -> Bool) -> v a -> (v a, v a)) = V.break `eq` break
+    prop_groupBy :: P ((a -> a -> Bool) -> v a -> [v a]) = V.groupBy `eq` groupBy
 
     prop_elem    :: P (a -> v a -> Bool) = V.elem `eq` elem
     prop_notElem :: P (a -> v a -> Bool) = V.notElem `eq` notElem