Further simplification

Compute the cstring literal's length in the caller, with the builder
functions now taking boxed arguments  Performance is the same, and the
code is cleaner.

Author: hs-viktor

Data/ByteString/Builder.hs

@@ -263,12 +263,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 lazy 'L.ByteString'.
 -- The work is performed lazy, i.e., only when a chunk of the lazy 'L.ByteString'
@@ -440,18 +442,20 @@ char8 :: Char -> Builder
 char8 = P.primFixed P.char8
 
 -- | Char8 encode a 'String'.
-{-# INLINE [1] string8 #-} -- phased to allow cstringLiteral rewrite
+{-# INLINE [1] string8 #-} -- phased to allow literal cstring rewrites
 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) = cstringLiteral s
+  string8 (unpackCString# s) =
+    ascLiteralCopy (Ptr s) (byteCountLiteral s)
 
 "string8/unpackFoldrCString#" forall s.
-  string8 (build (unpackFoldrCString# s)) = cstringLiteral s
+  string8 (build (unpackFoldrCString# s)) =
+    ascLiteralCopy (Ptr s) (byteCountLiteral s)
  #-}
 
 ------------------------------------------------------------------------------
@@ -467,19 +471,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 cstringLiteral rewrite
+{-# INLINE [1] stringUtf8 #-} -- phased to allow literal cstring rewrites
 stringUtf8 :: String -> Builder
 stringUtf8 = P.primMapListBounded P.charUtf8
 
 {-# RULES
 "stringUtf8/unpackCStringUtf8#" forall s.
-  stringUtf8 (unpackCStringUtf8# s) = cstringUtf8Literal s
+  stringUtf8 (unpackCStringUtf8# s) =
+    modUtf8LitCopy (Ptr s) (byteCountLiteral s)
 
 "stringUtf8/unpackCString#" forall s.
-  stringUtf8 (unpackCString# s) = cstringLiteral s
+  stringUtf8 (unpackCString# s) =
+    ascLiteralCopy (Ptr s) (byteCountLiteral s)
 
 "stringUtf8/unpackFoldrCString#" forall s.
-  stringUtf8 (build (unpackFoldrCString# s)) = cstringLiteral s
+  stringUtf8 (build (unpackFoldrCString# s)) =
+    ascLiteralCopy (Ptr s) (byteCountLiteral s)
  #-}
 
 instance IsString Builder where

Data/ByteString/Builder/Internal.hs

@@ -84,10 +84,10 @@ module Data.ByteString.Builder.Internal (
   -- , sizedChunksInsert
 
   , byteStringCopy
+  , ascLiteralCopy
+  , modUtf8LitCopy
   , byteStringInsert
   , byteStringThreshold
-  , cstringLiteral
-  , cstringUtf8Literal
 
   , lazyByteStringCopy
   , lazyByteStringInsert
@@ -872,62 +872,52 @@ byteStringInsert =
 -- @0xC0 0x80@) null characters.
 --
 -- @since 0.11.5.0
-{-# INLINABLE cstringLiteral #-}
-cstringLiteral :: Addr# -> Builder
-cstringLiteral = \addr -> builder $ \k br -> do
-    let ip   = Ptr addr
-#if __GLASGOW_HASKELL__ >= 811
-        ipe  = Ptr (addr `plusAddr#` (cstringLength# addr))
-#else
-    !ipe <- plusPtr ip . fromIntegral <$> S.c_strlen ip
-#endif
+{-# INLINABLE ascLiteralCopy #-}
+ascLiteralCopy :: Ptr Word8 -> Int -> Builder
+ascLiteralCopy = \ !ip !len -> builder $ \k br -> do
+    let !ipe = ip `plusPtr` len
     wrappedBytesCopyStep (BufferRange ip ipe) k br
 
 -- | 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.11.5.0
-{-# INLINABLE cstringUtf8Literal #-}
-cstringUtf8Literal :: Addr# -> Builder
-cstringUtf8Literal = \addr0 -> builder $ \k br -> do
-#if __GLASGOW_HASKELL__ >= 811
-    let len = I# (cstringLength# addr0)
-#else
-    len <- fromIntegral <$> S.c_strlen (Ptr addr0)
-#endif
-    nullAt <- c_strstr (Ptr addr0) (Ptr "\xc0\x80"#)
-    cstringUtf8_step (Ptr addr0) len nullAt k br
-
-cstringUtf8_step :: Ptr Word8-> Int -> Ptr Word8 -> BuildStep r -> BuildStep r
-cstringUtf8_step !ip !len ((== nullPtr) -> True) k br =
+{-# INLINABLE modUtf8LitCopy #-}
+modUtf8LitCopy :: Ptr Word8 -> Int -> Builder
+modUtf8LitCopy = \ !ip !len -> builder $ \k br -> do
+    nullAt <- c_strstr (castPtr ip) (Ptr "\xc0\x80"#)
+    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
-cstringUtf8_step !ip !len !nullAt k (BufferRange op0 ope)
+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) (S.memcpy op0 ip nullFree)
+        when (nullFree > 0) (copyBytes op0 ip nullFree)
         pokeElemOff op0 nullFree 0
-        let !op' = op0 `plusPtr` (nullFree + 1)
-            nread = nullFree + 2
-            !ip' = ip `plusPtr` nread
-            len' = len - nread
+        let used = nullFree + 2
+            len' = len - used
+            !ip' = ip `plusPtr` used
+            !op' = op0 `plusPtr` (nullFree + 1)
         nullAt' <- c_strstr ip' (Ptr "\xc0\x80"#)
-        cstringUtf8_step ip' len' nullAt' k (BufferRange op' ope)
+        modUtf8_step ip' len' nullAt' k (BufferRange op' ope)
     | avail > 0 = do
         -- avail <= nullFree
-        S.memcpy op0 ip avail
-        let !op' = op0 `plusPtr` avail
+        copyBytes op0 ip avail
+        let len' = len - avail
             !ip' = ip `plusPtr` avail
-            len' = len - avail
-        return $ bufferFull 1 op' (cstringUtf8_step ip' len' nullAt k)
+            !op' = op0 `plusPtr` avail
+        return $ bufferFull 1 op' (modUtf8_step ip' len' nullAt k)
     | otherwise =
-        return $ bufferFull 1 op0 (cstringUtf8_step ip len nullAt k)
+        return $ bufferFull 1 op0 (modUtf8_step ip len nullAt k)
   where
     !avail = ope `minusPtr` op0
     !nullFree = nullAt `minusPtr` ip

Data/ByteString/Builder/Prim.hs

@@ -668,19 +668,19 @@ primMapLazyByteStringBounded w =
 -- @0xC0 0x80@) null characters.
 --
 -- @since 0.11.0.0
-{-# DEPRECATED cstring "Use cstringLiteral instead" #-}
+{-# DEPRECATED cstring "Use ascLiteralCopy instead" #-}
 cstring :: Addr# -> Builder
-cstring = cstringLiteral
+cstring s = ascLiteralCopy (Ptr s) (S.byteCountLiteral s)
 {-# INLINE cstring #-}
 
 -- | 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.11.0.0
-{-# DEPRECATED cstringUtf8 "Use cstringUtf8Literal instead" #-}
+{-# DEPRECATED cstringUtf8 "Use modUtf8LitCopy instead" #-}
 cstringUtf8 :: Addr# -> Builder
-cstringUtf8 = cstringUtf8Literal
+cstringUtf8 s = modUtf8LitCopy (Ptr s) (S.byteCountLiteral s)
 {-# INLINE cstringUtf8 #-}
 
 ------------------------------------------------------------------------------

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,

Data/ByteString/Internal/Type.hs

@@ -39,7 +39,7 @@ module Data.ByteString.Internal.Type (
         unpackBytes, unpackAppendBytesLazy, unpackAppendBytesStrict,
         unpackChars, unpackAppendCharsLazy, unpackAppendCharsStrict,
         unsafePackAddress, unsafePackLenAddress,
-        unsafePackLiteral, unsafePackLenLiteral,
+        unsafePackLiteral, unsafePackLenLiteral, byteCountLiteral,
 
         -- * Low level imperative construction
         empty,
@@ -436,13 +436,22 @@ unsafePackLenAddress len addr# = do
 -- @since 0.11.1.0
 unsafePackLiteral :: Addr# -> ByteString
 unsafePackLiteral addr# =
+  unsafePackLenLiteral (byteCountLiteral addr#) addr#
+{-# INLINE unsafePackLiteral #-}
+
+-- | Byte count of null-terminated primitive literal string excluding the
+-- terminating null byte.
+byteCountLiteral :: Addr# -> Int
+byteCountLiteral addr# =
 #if __GLASGOW_HASKELL__ >= 811
-  unsafePackLenLiteral (I# (cstringLength# addr#)) addr#
+  I# (cstringLength# addr#)
 #else
-  let len = accursedUnutterablePerformIO (c_strlen (Ptr addr#))
-   in unsafePackLenLiteral (fromIntegral len) addr#
+  fromIntegral (pure_strlen (Ptr addr#))
+
+foreign import ccall unsafe "string.h strlen" pure_strlen
+    :: CString -> CSize
 #endif
-{-# INLINE unsafePackLiteral #-}
+{-# INLINE byteCountLiteral #-}
 
 
 -- | See 'unsafePackLiteral'. This function is similar,

bench/BenchAll.hs

@@ -25,6 +25,7 @@ import qualified Data.ByteString                       as S
 import qualified Data.ByteString.Char8                 as S8
 import qualified Data.ByteString.Lazy                  as L
 import qualified Data.ByteString.Lazy.Char8            as L8
+import           Data.ByteString.Internal              (byteCountLiteral)
 
 import           Data.ByteString.Builder
 import           Data.ByteString.Builder.Extra         (byteStringCopy,
@@ -250,12 +251,13 @@ largeTraversalInput = S.concat (replicate 10 byteStringData)
 smallTraversalInput :: S.ByteString
 smallTraversalInput = S8.pack "The quick brown fox"
 
-asciiLit, utf8Lit :: Ptr Word8
-asciiLit = Ptr "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"#
-utf8Lit  = Ptr "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx\xc0\x80xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"#
+ascBuf, utfBuf :: Ptr Word8
+ascBuf = Ptr "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"#
+utfBuf = Ptr "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx\xc0\x80xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"#
 
-addrOf :: Ptr Word8 -> Addr#
-addrOf (Ptr addr) = addr
+asclit, utflit :: Ptr Word8 -> Builder
+asclit str@(Ptr addr) = BI.ascLiteralCopy str (byteCountLiteral addr)
+utflit str@(Ptr addr) = BI.modUtf8LitCopy str (byteCountLiteral addr)
 
 main :: IO ()
 main = do
@@ -266,11 +268,11 @@ main = do
         , benchB' "ensureFree 8"  ()  (const (ensureFree 8))
         , benchB' "intHost 1"     1   intHost
         , benchB' "UTF-8 String (naive)" "hello world\0" fromString
-        , benchB' "UTF-8 String"  () $ \() -> BI.cstringUtf8Literal "hello world\xc0\x80"#
+        , benchB' "UTF-8 String"  () $ \() -> utflit (Ptr "hello world\xc0\x80"#)
         , benchB' "String (naive)" "hello world!" fromString
-        , benchB' "String"        () $ \() -> BI.cstringLiteral "hello world!"#
-        , benchB' "AsciiLit"      () $ \() -> BI.cstringLiteral (addrOf asciiLit)
-        , benchB' "Utf8Lit"       () $ \() -> BI.cstringUtf8Literal (addrOf utf8Lit)
+        , benchB' "String"        () $ \() -> asclit (Ptr "hello world!"#)
+        , benchB' "AsciiLit"      () $ \() -> asclit ascBuf
+        , benchB' "Utf8Lit"       () $ \() -> utflit utfBuf
         ]
 
       , bgroup "Encoding wrappers"

tests/builder/Data/ByteString/Builder/Tests.hs

@@ -50,6 +50,7 @@ import           Control.Exception (evaluate)
 import           System.IO (openTempFile, hPutStr, hClose, hSetBinaryMode, hSetEncoding, utf8, hSetNewlineMode, noNewlineTranslation)
 import           Foreign (ForeignPtr, withForeignPtr, castPtr)
 import           Foreign.C.String (withCString)
+import           GHC.Ptr (Ptr(..))
 import           Numeric (showFFloat)
 import           System.Posix.Internals (c_unlink)
 
@@ -994,14 +995,18 @@ testsUtf8 =
 testCString :: [TestTree]
 testCString =
     [ testProperty "cstring" $
-        toLazyByteString (BI.cstringLiteral "hello world!"#) ==
+        toLazyByteString (asclit (Ptr "hello world!"#)) ==
         LC.pack "hello" `L.append` L.singleton 0x20
                         `L.append` LC.pack "world!"
     , testProperty "cstringUtf8" $
-        toLazyByteString (BI.cstringUtf8Literal "hello\xc0\x80\xc0\x80world\xc0\x80!"#) ==
+        toLazyByteString (utflit (Ptr "hello\xc0\x80\xc0\x80world\xc0\x80!"#)) ==
         LC.pack "hello" `L.append` L.singleton 0x00
                         `L.append` L.singleton 0x00
                         `L.append` LC.pack "world"
                         `L.append` L.singleton 0x00
                         `L.append` LC.singleton '!'
     ]
+
+asclit, utflit :: Ptr Word8 -> Builder
+asclit str@(Ptr addr) = BI.ascLiteralCopy str (S.byteCountLiteral addr)
+utflit str@(Ptr addr) = BI.modUtf8LitCopy str (S.byteCountLiteral addr)