Avoid per-byte loop in cstring{,Utf8} builders

Data/ByteString/Builder.hs

@@ -257,12 +257,14 @@ import qualified Data.ByteString.Builder.Prim  as P
 import qualified Data.ByteString.Lazy.Internal as L
 import           Data.ByteString.Builder.ASCII
 import           Data.ByteString.Builder.RealFloat
+import           Data.ByteString.Internal (byteCountLiteral)
 
 import           Data.String (IsString(..))
 import           System.IO (Handle, IOMode(..), withBinaryFile)
 import           Foreign
 import           GHC.Base (unpackCString#, unpackCStringUtf8#,
                            unpackFoldrCString#, build)
+import           GHC.Ptr (Ptr(..))
 
 -- | Execute a 'Builder' and return the generated chunks as a 'L.LazyByteString'.
 -- The work is performed lazy, i.e., only when a chunk of the 'L.LazyByteString'
@@ -434,18 +436,20 @@ char8 :: Char -> Builder
 char8 = P.primFixed P.char8
 
 -- | Char8 encode a 'String'.
-{-# INLINE [1] string8 #-} -- phased to allow P.cstring rewrite
+{-# NOINLINE string8 #-}
 string8 :: String -> Builder
 string8 = P.primMapListFixed P.char8
 
 -- GHC desugars string literals with unpackCString# which the simplifier tends
 -- to promptly turn into build (unpackFoldrCString# s), so we match on both.
 {-# RULES
 "string8/unpackCString#" forall s.
-  string8 (unpackCString# s) = P.cstring s
+  string8 (unpackCString# s) =
+    asciiLiteralCopy (Ptr s) (byteCountLiteral s)
 
 "string8/unpackFoldrCString#" forall s.
-  string8 (build (unpackFoldrCString# s)) = P.cstring s
+  string8 (build (unpackFoldrCString# s)) =
+    asciiLiteralCopy (Ptr s) (byteCountLiteral s)
  #-}
 
 ------------------------------------------------------------------------------
@@ -461,19 +465,22 @@ charUtf8 = P.primBounded P.charUtf8
 --
 -- Note that 'stringUtf8' performs no codepoint validation and consequently may
 -- emit invalid UTF-8 if asked (e.g. single surrogates).
-{-# INLINE [1] stringUtf8 #-} -- phased to allow P.cstring rewrite
+{-# NOINLINE stringUtf8 #-}
 stringUtf8 :: String -> Builder
 stringUtf8 = P.primMapListBounded P.charUtf8
 
 {-# RULES
 "stringUtf8/unpackCStringUtf8#" forall s.
-  stringUtf8 (unpackCStringUtf8# s) = P.cstringUtf8 s
+  stringUtf8 (unpackCStringUtf8# s) =
+    modUtf8LitCopy (Ptr s) (byteCountLiteral s)
 
 "stringUtf8/unpackCString#" forall s.
-  stringUtf8 (unpackCString# s) = P.cstring s
+  stringUtf8 (unpackCString# s) =
+    asciiLiteralCopy (Ptr s) (byteCountLiteral s)
 
 "stringUtf8/unpackFoldrCString#" forall s.
-  stringUtf8 (build (unpackFoldrCString# s)) = P.cstring s
+  stringUtf8 (build (unpackFoldrCString# s)) =
+    asciiLiteralCopy (Ptr s) (byteCountLiteral s)
  #-}
 
 instance IsString Builder where

Data/ByteString/Builder/Internal.hs

@@ -1,5 +1,5 @@
 {-# LANGUAGE ScopedTypeVariables, CPP, BangPatterns, RankNTypes, TupleSections #-}
-{-# LANGUAGE Unsafe #-}
+{-# LANGUAGE MagicHash, ViewPatterns, Unsafe #-}
 {-# OPTIONS_HADDOCK not-home #-}
 -- | Copyright : (c) 2010 - 2011 Simon Meier
 -- License     : BSD3-style (see LICENSE)
@@ -84,6 +84,8 @@ module Data.ByteString.Builder.Internal (
   -- , sizedChunksInsert
 
   , byteStringCopy
+  , asciiLiteralCopy
+  , modUtf8LitCopy
   , byteStringInsert
   , byteStringThreshold
 
@@ -127,6 +129,8 @@ module Data.ByteString.Builder.Internal (
 ) where
 
 import           Control.Arrow (second)
+import           Control.Monad (when)
+import           Control.DeepSeq (NFData(..))
 
 import           Data.Semigroup (Semigroup(..))
 import           Data.List.NonEmpty (NonEmpty(..))
@@ -146,6 +150,11 @@ import           Foreign
 import           Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
 import           System.IO.Unsafe (unsafeDupablePerformIO)
 
+#if !(PURE_HASKELL || defined(USE_MEMCHR))
+import           Foreign.C.String (CString)
+import           GHC.Ptr (Ptr(..))
+#endif
+
 ------------------------------------------------------------------------------
 -- Buffers
 ------------------------------------------------------------------------------
@@ -154,11 +163,22 @@ import           System.IO.Unsafe (unsafeDupablePerformIO)
 data BufferRange = BufferRange {-# UNPACK #-} !(Ptr Word8)  -- First byte of range
                                {-# UNPACK #-} !(Ptr Word8)  -- First byte /after/ range
 
+-- | @since 0.12.1.0
+instance NFData BufferRange where
+  rnf !_ = ()
+
 -- | A 'Buffer' together with the 'BufferRange' of free bytes. The filled
 -- space starts at offset 0 and ends at the first free byte.
 data Buffer = Buffer {-# UNPACK #-} !(ForeignPtr Word8)
                      {-# UNPACK #-} !BufferRange
 
+-- | Like the @NFData@ instance for @StrictByteString@,
+-- this does not force the @ForeignPtrContents@ field
+-- of the underlying @ForeignPtr@.
+--
+-- @since 0.12.1.0
+instance NFData Buffer where
+  rnf !_ = ()
 
 -- | Combined size of the filled and free space in the buffer.
 {-# INLINE bufferSize #-}
@@ -876,6 +896,93 @@ byteStringInsert :: S.StrictByteString -> Builder
 byteStringInsert =
     \bs -> builder $ \k (BufferRange op _) -> return $ insertChunk op bs k
 
+
+------------------------------------------------------------------------------
+-- Raw CString encoding
+------------------------------------------------------------------------------
+
+-- | Builder for raw 'Addr#' pointers to null-terminated primitive ASCII
+-- strings that are free of embedded (overlong-encoded as the two-byte sequence
+-- @0xC0 0x80@) null characters.
+--
+-- @since 0.12.1.0
+{-# INLINABLE asciiLiteralCopy #-}
+asciiLiteralCopy :: Ptr Word8 -> Int -> Builder
+asciiLiteralCopy = \ !ip !len -> builder $ \k br -> do
+    let !ipe = ip `plusPtr` len
+    wrappedBytesCopyStep (BufferRange ip ipe) k br
+
+getNextEmbeddedNull :: Ptr Word8 -> Int -> IO (Ptr Word8)
+#if PURE_HASKELL || defined(USE_MEMCHR)
+getNextEmbeddedNull p len = do
+  c0loc <- S.memchr p 0xC0 (S.checkedCast len)
+  if c0loc == nullPtr
+    then pure c0loc
+    else do
+    let nextLoc = c0loc `plusPtr` 1 :: Ptr Word8
+    nextByte <- peek nextLoc
+    if nextByte == 0x80
+      then pure c0loc
+      else getNextEmbeddedNull nextLoc (p `minusPtr` nextLoc + len)
+
+#else
+getNextEmbeddedNull p _len = c_strstr (castPtr p) modifiedUtf8NUL
+
+-- | GHC represents @NUL@ in string literals via an overlong 2-byte encoding,
+-- which is part of "modified UTF-8" (GHC does not also implement CESU-8).
+modifiedUtf8NUL :: CString
+modifiedUtf8NUL = Ptr "\xc0\x80"#
+
+foreign import ccall unsafe "string.h strstr" c_strstr
+    :: CString -> CString -> IO (Ptr Word8)
+#endif
+
+
+-- | Builder for raw 'Addr#' pointers to null-terminated primitive UTF-8
+-- encoded strings that may contain embedded overlong-encodings (as the
+-- two-byte sequence @0xC0 0x80@) of null characters.
+--
+-- @since 0.12.1.0
+{-# INLINABLE modUtf8LitCopy #-}
+modUtf8LitCopy :: Ptr Word8 -> Int -> Builder
+modUtf8LitCopy = \ !ip !len -> builder $ \k br -> do
+    nullAt <- getNextEmbeddedNull ip len
+    modUtf8_step ip len nullAt k br
+
+modUtf8_step :: Ptr Word8 -> Int -> Ptr Word8 -> BuildStep r -> BuildStep r
+modUtf8_step !ip !len ((== nullPtr) -> True) k br =
+    -- Contains no encoded nulls, use simple copy codepath
+    wrappedBytesCopyStep (BufferRange ip ipe) k br
+  where
+    !ipe = ip `plusPtr` len
+modUtf8_step !ip !len !nullAt k (BufferRange op0 ope)
+    -- Copy as much of the null-free portion of the string as fits into the
+    -- available buffer space. If the string is long enough, we may have asked
+    -- for less than its full length, filling the buffer with the rest will go
+    -- into the next builder step.
+    | avail > nullFree = do
+        when (nullFree > 0) (copyBytes op0 ip nullFree)
+        pokeElemOff op0 nullFree 0
+        let used = nullFree + 2
+            len' = len - used
+            !ip' = ip `plusPtr` used
+            !op' = op0 `plusPtr` (nullFree + 1)
+        nullAt' <- getNextEmbeddedNull ip' len'
+        modUtf8_step ip' len' nullAt' k (BufferRange op' ope)
+    | avail > 0 = do
+        -- avail <= nullFree
+        copyBytes op0 ip avail
+        let len' = len - avail
+            !ip' = ip `plusPtr` avail
+            !op' = op0 `plusPtr` avail
+        return $ bufferFull 1 op' (modUtf8_step ip' len' nullAt k)
+    | otherwise =
+        return $ bufferFull 1 op0 (modUtf8_step ip len nullAt k)
+  where
+    !avail = ope `minusPtr` op0
+    !nullFree = nullAt `minusPtr` ip
+
+
 -- Short bytestrings
 ------------------------------------------------------------------------------
 

Data/ByteString/Builder/Prim.hs

@@ -579,7 +579,10 @@ primBounded w x =
 -- because it moves several variables out of the inner loop.
 {-# INLINE primMapListBounded #-}
 primMapListBounded :: BoundedPrim a -> [a] -> Builder
-primMapListBounded w xs0 =
+primMapListBounded w = \xs0 ->
+  -- We want this to inline when there is one arg, so that we can
+  -- specialise on the BoundedPrim "w".  So we move the \xs0 after the
+  -- "=" sign so that the INLINE pragma doesn't interfere with this.
     builder $ step xs0
   where
     step xs1 k (BufferRange op0 ope0) =
@@ -663,59 +666,29 @@ primMapLazyByteStringBounded w =
     L.foldrChunks (\x b -> primMapByteStringBounded w x `mappend` b) mempty
 
 
-------------------------------------------------------------------------------
--- Raw CString encoding
-------------------------------------------------------------------------------
-
--- | A null-terminated ASCII encoded 'Foreign.C.String.CString'.
--- Null characters are not representable.
+-- | Builder for raw 'Addr#' pointers to null-terminated primitive ASCII
+-- strings that are free of embedded (overlong-encoded as the two-byte sequence
+-- @0xC0 0x80@) null characters.
+--
+-- Deprecated since @bytestring-0.12.1.0@.
 --
 -- @since 0.11.0.0
+{-# DEPRECATED cstring "Use asciiLiteralCopy instead" #-}
 cstring :: Addr# -> Builder
-cstring =
-    \addr0 -> builder $ step addr0
-  where
-    step :: Addr# -> BuildStep r -> BuildStep r
-    step !addr !k br@(BufferRange op0@(Ptr op0#) ope)
-      | W8# ch == 0 = k br
-      | op0 == ope =
-          return $ bufferFull 1 op0 (step addr k)
-      | otherwise = do
-          IO $ \s -> case writeWord8OffAddr# op0# 0# ch s of
-                       s' -> (# s', () #)
-          let br' = BufferRange (op0 `plusPtr` 1) ope
-          step (addr `plusAddr#` 1#) k br'
-      where
-        !ch = indexWord8OffAddr# addr 0#
+cstring s = asciiLiteralCopy (Ptr s) (S.byteCountLiteral s)
+{-# INLINE cstring #-}
 
--- | A null-terminated UTF-8 encoded 'Foreign.C.String.CString'.
--- Null characters can be encoded as @0xc0 0x80@.
+-- | Builder for raw 'Addr#' pointers to null-terminated primitive UTF-8
+-- encoded strings that may contain embedded overlong-encodings (as the
+-- two-byte sequence @0xC0 0x80@) of null characters.
+--
+-- Deprecated since @bytestring-0.12.1.0@.
 --
 -- @since 0.11.0.0
+{-# DEPRECATED cstringUtf8 "Use modUtf8LitCopy instead" #-}
 cstringUtf8 :: Addr# -> Builder
-cstringUtf8 =
-    \addr0 -> builder $ step addr0
-  where
-    step :: Addr# -> BuildStep r -> BuildStep r
-    step !addr !k br@(BufferRange op0@(Ptr op0#) ope)
-      | W8# ch == 0 = k br
-      | op0 == ope =
-          return $ bufferFull 1 op0 (step addr k)
-        -- NULL is encoded as 0xc0 0x80
-      | W8# ch == 0xc0
-      , W8# (indexWord8OffAddr# addr 1#) == 0x80 = do
-          let !(W8# nullByte#) = 0
-          IO $ \s -> case writeWord8OffAddr# op0# 0# nullByte# s of
-                       s' -> (# s', () #)
-          let br' = BufferRange (op0 `plusPtr` 1) ope
-          step (addr `plusAddr#` 2#) k br'
-      | otherwise = do
-          IO $ \s -> case writeWord8OffAddr# op0# 0# ch s of
-                       s' -> (# s', () #)
-          let br' = BufferRange (op0 `plusPtr` 1) ope
-          step (addr `plusAddr#` 1#) k br'
-      where
-        !ch = indexWord8OffAddr# addr 0#
+cstringUtf8 s = modUtf8LitCopy (Ptr s) (S.byteCountLiteral s)
+{-# INLINE cstringUtf8 #-}
 
 ------------------------------------------------------------------------------
 -- Char8 encoding

Data/ByteString/Internal.hs

@@ -36,7 +36,7 @@ module Data.ByteString.Internal (
         unpackBytes, unpackAppendBytesLazy, unpackAppendBytesStrict,
         unpackChars, unpackAppendCharsLazy, unpackAppendCharsStrict,
         unsafePackAddress, unsafePackLenAddress,
-        unsafePackLiteral, unsafePackLenLiteral,
+        unsafePackLiteral, unsafePackLenLiteral, byteCountLiteral,
 
         -- * Low level imperative construction
         empty,