Improve the performance of `partition` for lazy and strict bytestrings
authorSean <burton.seanr@gmail.com>
Fri, 18 Sep 2015 08:14:01 +0000 (09:14 +0100)
committerSean <burton.seanr@gmail.com>
Tue, 29 Sep 2015 10:17:24 +0000 (11:17 +0100)
Data/ByteString.hs
Data/ByteString/Internal.hs
Data/ByteString/Lazy.hs

index 42263a6..c141441 100644 (file)
@@ -231,7 +231,7 @@ import Data.Maybe               (isJust, listToMaybe)
 #ifndef __NHC__
 import Control.Exception        (finally, bracket, assert, throwIO)
 #else
-import Control.Exception       (bracket, finally)
+import Control.Exception       (bracket, finally)
 #endif
 import Control.Monad            (when)
 
@@ -1259,14 +1259,42 @@ find f p = case findIndex f p of
 -- > partition p bs == (filter p xs, filter (not . p) xs)
 --
 partition :: (Word8 -> Bool) -> ByteString -> (ByteString, ByteString)
-partition p bs = (filter p bs, filter (not . p) bs)
---TODO: use a better implementation
-
--- ---------------------------------------------------------------------
--- Searching for substrings
-
--- | /O(n)/ The 'isPrefixOf' function takes two ByteStrings and returns 'True'
--- iff the first is a prefix of the second.
+partition f s = unsafeDupablePerformIO $
+    do fp' <- mallocByteString len
+       withForeignPtr fp' $ \p ->
+           do let end = p `plusPtr` (len - 1)
+              mid <- sep 0 p end
+              rev mid end
+              let i = mid `minusPtr` p
+              return (PS fp' 0 i,
+                      PS fp' i (len - i))
+  where
+    len  = length s
+    incr = (`plusPtr` 1)
+    decr = (`plusPtr` (-1))
+
+    sep !i !p1 !p2
+       | i == len  = return p1
+       | f w       = do poke p1 w
+                        sep (i + 1) (incr p1) p2
+       | otherwise = do poke p2 w
+                        sep (i + 1) p1 (decr p2)
+      where
+        w = s `unsafeIndex` i
+
+    rev !p1 !p2
+      | p1 >= p2  = return ()
+      | otherwise = do a <- peek p1
+                       b <- peek p2
+                       poke p1 b
+                       poke p2 a
+                       rev (incr p1) (decr p2)
+
+-- --------------------------------------------------------------------
+-- Sarching for substrings
+
+-- |/O(n)/ The 'isPrefixOf' function takes two ByteStrings and returns 'True'
+-- if the first is a prefix of the second.
 isPrefixOf :: ByteString -> ByteString -> Bool
 isPrefixOf (PS x1 s1 l1) (PS x2 s2 l2)
     | l1 == 0   = True
index 0346d01..9747d10 100644 (file)
@@ -169,7 +169,7 @@ import Foreign.ForeignPtr       (newForeignPtr_)
 
 -- An alternative to Control.Exception (assert) for nhc98
 #ifdef __NHC__
-#define assert assertS "__FILE__ : __LINE__"
+#define assert  assertS "__FILE__ : __LINE__"
 assertS :: String -> Bool -> a -> a
 assertS _ True  = id
 assertS s False = error ("assertion failed at "++s)
index ad8938b..5a0b0bd 100644 (file)
@@ -1017,8 +1017,11 @@ filterNotByte w (LPS xs) = LPS (filterMap (P.filterNotByte w) xs)
 -- > partition p bs == (filter p xs, filter (not . p) xs)
 --
 partition :: (Word8 -> Bool) -> ByteString -> (ByteString, ByteString)
-partition f p = (filter f p, filter (not . f) p)
---TODO: use a better implementation
+partition _ Empty = (Empty, Empty)
+partition p (Chunk x xs) = (chunk t ts, chunk f fs)
+  where
+    (t,   f) = S.partition p x
+    (ts, fs) = partition   p xs
 
 -- ---------------------------------------------------------------------
 -- Searching for substrings