From: Roman Leshchinskiy
Date: Fri, 11 Dec 2009 02:22:27 +0000 (+0000)
Subject: Add partition
X-Git-Tag: 0_5~10
X-Git-Url: http://git.haskell.org/darcs-mirrors/vector.git/commitdiff_plain/9ee6ff43275f48611453ae47eb037c11ceda7363
Add partition
---
diff --git a/Data/Vector.hs b/Data/Vector.hs
index e72279d..71a538e 100644
--- a/Data/Vector.hs
+++ b/Data/Vector.hs
@@ -49,7 +49,7 @@ module Data.Vector (
-- * Filtering
filter, ifilter, takeWhile, dropWhile,
- unstablePartition, span, break,
+ partition, unstablePartition, span, break,
-- * Searching
elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
@@ -518,6 +518,14 @@ dropWhile :: (a -> Bool) -> Vector a -> Vector a
dropWhile = G.dropWhile
-- | Split the vector in two parts, the first one containing those elements
+-- that satisfy the predicate and the second one those that don't. The
+-- relative order of the elements is preserved at the cost of a (sometimes)
+-- reduced performance compared to 'unstablePartition'.
+partition :: (a -> Bool) -> Vector a -> (Vector a, Vector a)
+{-# INLINE partition #-}
+partition = G.partition
+
+-- | Split the vector in two parts, the first one containing those elements
-- that satisfy the predicate and the second one those that don't. The order
-- of the elements is not preserved.
unstablePartition :: (a -> Bool) -> Vector a -> (Vector a, Vector a)
diff --git a/Data/Vector/Generic.hs b/Data/Vector/Generic.hs
index e6f8f3c..ae43f6d 100644
--- a/Data/Vector/Generic.hs
+++ b/Data/Vector/Generic.hs
@@ -53,7 +53,7 @@ module Data.Vector.Generic (
-- * Filtering
filter, ifilter, takeWhile, dropWhile,
- unstablePartition, span, break,
+ partition, unstablePartition, span, break,
-- * Searching
elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
@@ -863,8 +863,29 @@ dropWhile :: Vector v a => (a -> Bool) -> v a -> v a
dropWhile f = unstream . Stream.dropWhile f . stream
-- | Split the vector in two parts, the first one containing those elements
+-- that satisfy the predicate and the second one those that don't. The
+-- relative order of the elements is preserved at the cost of a (sometimes)
+-- reduced performance compared to 'unstablePartition'.
+partition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
+{-# INLINE partition #-}
+partition f = partition_stream f . stream
+
+-- FIXME: Make this inplace-fusible (look at how stable_partition is
+-- implemented in C++)
+
+partition_stream :: Vector v a => (a -> Bool) -> Stream a -> (v a, v a)
+{-# INLINE_STREAM partition_stream #-}
+partition_stream f s = s `seq` runST (
+ do
+ (mv1,mv2) <- M.partitionStream f s
+ v1 <- unsafeFreeze mv1
+ v2 <- unsafeFreeze mv2
+ return (v1,v2))
+
+-- | Split the vector in two parts, the first one containing those elements
-- that satisfy the predicate and the second one those that don't. The order
--- of the elements is not preserved.
+-- of the elements is not preserved but the operation is often faster than
+-- 'partition'.
unstablePartition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
{-# INLINE unstablePartition #-}
unstablePartition f = unstablePartition_stream f . stream
diff --git a/Data/Vector/Generic/Mutable.hs b/Data/Vector/Generic/Mutable.hs
index cd1b07c..fcd9b95 100644
--- a/Data/Vector/Generic/Mutable.hs
+++ b/Data/Vector/Generic/Mutable.hs
@@ -28,7 +28,7 @@ module Data.Vector.Generic.Mutable (
-- * Internal operations
unstream, transform, unstreamR, transformR,
unsafeAccum, accum, unsafeUpdate, update, reverse,
- unstablePartition, unstablePartitionStream
+ unstablePartition, unstablePartitionStream, partitionStream
) where
import qualified Data.Vector.Fusion.Stream as Stream
@@ -693,6 +693,41 @@ unstablePartitionMax f s n
(i,j) <- Stream.foldM' put (0, n) s
return (unsafeSlice 0 i v, unsafeSlice j (n-j) v)
+partitionStream :: (PrimMonad m, MVector v a)
+ => (a -> Bool) -> Stream a -> m (v (PrimState m) a, v (PrimState m) a)
+{-# INLINE partitionStream #-}
+partitionStream f s
+ = case upperBound (Stream.size s) of
+ Just n -> partitionMax f s n
+ Nothing -> partitionUnknown f s
+
+partitionMax :: (PrimMonad m, MVector v a)
+ => (a -> Bool) -> Stream a -> Int -> m (v (PrimState m) a, v (PrimState m) a)
+{-# INLINE partitionMax #-}
+partitionMax f s n
+ = do
+ v <- INTERNAL_CHECK(checkLength) "unstablePartitionMax" n
+ $ unsafeNew n
+
+ let {-# INLINE_INNER put #-}
+ put (i,j) x
+ | f x = do
+ unsafeWrite v i x
+ return (i+1,j)
+
+ | otherwise = let j' = j-1 in
+ do
+ unsafeWrite v j' x
+ return (i,j')
+
+ (i,j) <- Stream.foldM' put (0,n) s
+ INTERNAL_CHECK(check) "partitionMax" "invalid indices" (i <= j)
+ $ return ()
+ let l = unsafeSlice 0 i v
+ r = unsafeSlice j (n-j) v
+ reverse r
+ return (l,r)
+
partitionUnknown :: (PrimMonad m, MVector v a)
=> (a -> Bool) -> Stream a -> m (v (PrimState m) a, v (PrimState m) a)
{-# INLINE partitionUnknown #-}
diff --git a/Data/Vector/Primitive.hs b/Data/Vector/Primitive.hs
index d4cf10b..5207fb2 100644
--- a/Data/Vector/Primitive.hs
+++ b/Data/Vector/Primitive.hs
@@ -45,7 +45,7 @@ module Data.Vector.Primitive (
-- * Filtering
filter, ifilter, takeWhile, dropWhile,
- unstablePartition, span, break,
+ partition, unstablePartition, span, break,
-- * Searching
elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
@@ -469,6 +469,14 @@ dropWhile :: Prim a => (a -> Bool) -> Vector a -> Vector a
dropWhile = G.dropWhile
-- | Split the vector in two parts, the first one containing those elements
+-- that satisfy the predicate and the second one those that don't. The
+-- relative order of the elements is preserved at the cost of a (sometimes)
+-- reduced performance compared to 'unstablePartition'.
+partition :: Prim a => (a -> Bool) -> Vector a -> (Vector a, Vector a)
+{-# INLINE partition #-}
+partition = G.partition
+
+-- | Split the vector in two parts, the first one containing those elements
-- that satisfy the predicate and the second one those that don't. The order
-- of the elements is not preserved.
unstablePartition :: Prim a => (a -> Bool) -> Vector a -> (Vector a, Vector a)
diff --git a/Data/Vector/Storable.hs b/Data/Vector/Storable.hs
index bab7c27..e945b84 100644
--- a/Data/Vector/Storable.hs
+++ b/Data/Vector/Storable.hs
@@ -45,7 +45,7 @@ module Data.Vector.Storable (
-- * Filtering
filter, ifilter, takeWhile, dropWhile,
- unstablePartition, span, break,
+ partition, unstablePartition, span, break,
-- * Searching
elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
@@ -504,6 +504,14 @@ dropWhile :: Storable a => (a -> Bool) -> Vector a -> Vector a
dropWhile = G.dropWhile
-- | Split the vector in two parts, the first one containing those elements
+-- that satisfy the predicate and the second one those that don't. The
+-- relative order of the elements is preserved at the cost of a (sometimes)
+-- reduced performance compared to 'unstablePartition'.
+partition :: Storable a => (a -> Bool) -> Vector a -> (Vector a, Vector a)
+{-# INLINE partition #-}
+partition = G.partition
+
+-- | Split the vector in two parts, the first one containing those elements
-- that satisfy the predicate and the second one those that don't. The order
-- of the elements is not preserved.
unstablePartition
diff --git a/Data/Vector/Unboxed.hs b/Data/Vector/Unboxed.hs
index 2c6c5d2..6659f09 100644
--- a/Data/Vector/Unboxed.hs
+++ b/Data/Vector/Unboxed.hs
@@ -47,7 +47,7 @@ module Data.Vector.Unboxed (
-- * Filtering
filter, ifilter, takeWhile, dropWhile,
- unstablePartition, span, break,
+ partition, unstablePartition, span, break,
-- * Searching
elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
@@ -463,6 +463,14 @@ dropWhile :: Unbox a => (a -> Bool) -> Vector a -> Vector a
dropWhile = G.dropWhile
-- | Split the vector in two parts, the first one containing those elements
+-- that satisfy the predicate and the second one those that don't. The
+-- relative order of the elements is preserved at the cost of a (sometimes)
+-- reduced performance compared to 'unstablePartition'.
+partition :: Unbox a => (a -> Bool) -> Vector a -> (Vector a, Vector a)
+{-# INLINE partition #-}
+partition = G.partition
+
+-- | Split the vector in two parts, the first one containing those elements
-- that satisfy the predicate and the second one those that don't. The order
-- of the elements is not preserved.
unstablePartition :: Unbox a => (a -> Bool) -> Vector a -> (Vector a, Vector a)
diff --git a/tests/Tests/Vector.hs b/tests/Tests/Vector.hs
index ce46d84..512c0e7 100644
--- a/tests/Tests/Vector.hs
+++ b/tests/Tests/Vector.hs
@@ -94,7 +94,7 @@ testPolymorphicFunctions _ = $(testProperties [
'prop_imap, 'prop_izipWith, 'prop_izipWith3,
'prop_filter, 'prop_ifilter, 'prop_takeWhile, 'prop_dropWhile,
- 'prop_span, 'prop_break,
+ 'prop_partition, 'prop_span, 'prop_break,
'prop_elem, 'prop_notElem,
'prop_find, 'prop_findIndex, 'prop_findIndices,
@@ -198,6 +198,8 @@ testPolymorphicFunctions _ = $(testProperties [
prop_ifilter :: P ((Int -> a -> Bool) -> v a -> v a) = V.ifilter `eq` ifilter
prop_takeWhile :: P ((a -> Bool) -> v a -> v a) = V.takeWhile `eq` takeWhile
prop_dropWhile :: P ((a -> Bool) -> v a -> v a) = V.dropWhile `eq` dropWhile
+ prop_partition :: P ((a -> Bool) -> v a -> (v a, v a))
+ = V.partition `eq` partition
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