Add partition
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Fri, 11 Dec 2009 02:22:27 +0000 (02:22 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Fri, 11 Dec 2009 02:22:27 +0000 (02:22 +0000)
Data/Vector.hs
Data/Vector/Generic.hs
Data/Vector/Generic/Mutable.hs
Data/Vector/Primitive.hs
Data/Vector/Storable.hs
Data/Vector/Unboxed.hs
tests/Tests/Vector.hs

index e72279d..71a538e 100644 (file)
@@ -49,7 +49,7 @@ module Data.Vector (
 
   -- * Filtering
   filter, ifilter, takeWhile, dropWhile,
 
   -- * Filtering
   filter, ifilter, takeWhile, dropWhile,
-  unstablePartition, span, break,
+  partition, unstablePartition, span, break,
 
   -- * Searching
   elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
 
   -- * 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
 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)
 -- 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)
index e6f8f3c..ae43f6d 100644 (file)
@@ -53,7 +53,7 @@ module Data.Vector.Generic (
 
   -- * Filtering
   filter, ifilter, takeWhile, dropWhile,
 
   -- * Filtering
   filter, ifilter, takeWhile, dropWhile,
-  unstablePartition, span, break,
+  partition, unstablePartition, span, break,
 
   -- * Searching
   elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
 
   -- * 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
 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
 -- 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
 unstablePartition :: Vector v a => (a -> Bool) -> v a -> (v a, v a)
 {-# INLINE unstablePartition #-}
 unstablePartition f = unstablePartition_stream f . stream
index cd1b07c..fcd9b95 100644 (file)
@@ -28,7 +28,7 @@ module Data.Vector.Generic.Mutable (
   -- * Internal operations
   unstream, transform, unstreamR, transformR,
   unsafeAccum, accum, unsafeUpdate, update, reverse,
   -- * 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
 ) 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)
 
       (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 #-}
 partitionUnknown :: (PrimMonad m, MVector v a)
         => (a -> Bool) -> Stream a -> m (v (PrimState m) a, v (PrimState m) a)
 {-# INLINE partitionUnknown #-}
index d4cf10b..5207fb2 100644 (file)
@@ -45,7 +45,7 @@ module Data.Vector.Primitive (
 
   -- * Filtering
   filter, ifilter, takeWhile, dropWhile,
 
   -- * Filtering
   filter, ifilter, takeWhile, dropWhile,
-  unstablePartition, span, break,
+  partition, unstablePartition, span, break,
 
   -- * Searching
   elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
 
   -- * 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
 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)
 -- 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)
index bab7c27..e945b84 100644 (file)
@@ -45,7 +45,7 @@ module Data.Vector.Storable (
 
   -- * Filtering
   filter, ifilter, takeWhile, dropWhile,
 
   -- * Filtering
   filter, ifilter, takeWhile, dropWhile,
-  unstablePartition, span, break,
+  partition, unstablePartition, span, break,
 
   -- * Searching
   elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
 
   -- * 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
 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
 -- that satisfy the predicate and the second one those that don't. The order
 -- of the elements is not preserved.
 unstablePartition
index 2c6c5d2..6659f09 100644 (file)
@@ -47,7 +47,7 @@ module Data.Vector.Unboxed (
 
   -- * Filtering
   filter, ifilter, takeWhile, dropWhile,
 
   -- * Filtering
   filter, ifilter, takeWhile, dropWhile,
-  unstablePartition, span, break,
+  partition, unstablePartition, span, break,
 
   -- * Searching
   elem, notElem, find, findIndex, findIndices, elemIndex, elemIndices,
 
   -- * 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
 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)
 -- 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)
index ce46d84..512c0e7 100644 (file)
@@ -94,7 +94,7 @@ testPolymorphicFunctions _ = $(testProperties [
         'prop_imap, 'prop_izipWith, 'prop_izipWith3,
 
         'prop_filter, 'prop_ifilter, 'prop_takeWhile, 'prop_dropWhile,
         '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,
 
         '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_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
 
     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