[commit: packages/bytestring] master, revert-46-patch-1, wip/nix-local-build: Improve the performance of `partition` for lazy and strict bytestrings (fae6927)

[git at git.haskell.org]

Repository : ssh://git@git.haskell.org/bytestring

On branches: master,revert-46-patch-1,wip/nix-local-build
Link       : http://git.haskell.org/packages/bytestring.git/commitdiff/fae6927d42b8c9d31e670adfb2b80bbe18bacff9

>---------------------------------------------------------------

commit fae6927d42b8c9d31e670adfb2b80bbe18bacff9
Author: Sean <burton.seanr at gmail.com>
Date:   Fri Sep 18 09:14:01 2015 +0100

    Improve the performance of `partition` for lazy and strict bytestrings


>---------------------------------------------------------------

fae6927d42b8c9d31e670adfb2b80bbe18bacff9
 Data/ByteString.hs          | 46 ++++++++++++++++++++++++++++++++++++---------
 Data/ByteString/Internal.hs |  2 +-
 Data/ByteString/Lazy.hs     |  7 +++++--
 3 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/Data/ByteString.hs b/Data/ByteString.hs
index 42263a6..c141441 100644
--- a/Data/ByteString.hs
+++ b/Data/ByteString.hs
@@ -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
diff --git a/Data/ByteString/Internal.hs b/Data/ByteString/Internal.hs
index 0346d01..9747d10 100644
--- a/Data/ByteString/Internal.hs
+++ b/Data/ByteString/Internal.hs
@@ -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)
diff --git a/Data/ByteString/Lazy.hs b/Data/ByteString/Lazy.hs
index ad8938b..5a0b0bd 100644
--- a/Data/ByteString/Lazy.hs
+++ b/Data/ByteString/Lazy.hs
@@ -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