Export reverseCommitsOntoChunks from the Unsafe module

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 * Add big-endian and little-endian parsers for `Word128`. This includes
   both the single and multiple element variants.
+* Export `reverseCommitsOntoChunks` from the `Unsafe` module.
 
 ## 0.3.0.0 -- 2019-10-17
 

src/Data/ByteArray/Builder.hs

@@ -97,13 +97,13 @@ module Data.ByteArray.Builder
   , flush
   ) where
 
-import Control.Monad.Primitive (primitive_)
 import Control.Monad.ST (ST,runST)
 import Data.ByteArray.Builder.Unsafe (Builder(Builder))
 import Data.ByteArray.Builder.Unsafe (Commits(Initial,Mutable,Immutable))
+import Data.ByteArray.Builder.Unsafe (reverseCommitsOntoChunks)
 import Data.ByteArray.Builder.Unsafe (stringUtf8,cstring)
 import Data.ByteString.Short.Internal (ShortByteString(SBS))
-import Data.Bytes.Chunks (Chunks(..))
+import Data.Bytes.Chunks (Chunks(ChunksNil))
 import Data.Bytes.Types (Bytes(Bytes))
 import Data.Char (ord)
 import Data.Int (Int64,Int32,Int16,Int8)
@@ -134,21 +134,7 @@ run hint@(I# hint# ) (Builder f) = runST $ do
   cs <- ST $ \s0 -> case f buf0 0# hint# Initial s0 of
     (# s1, bufX, offX, _, csX #) ->
       (# s1, Mutable bufX offX csX #)
-  commitsOntoChunks ChunksNil cs
-
--- Internal. This freezes all the mutable byte arrays in-place,
--- so be careful. It also reverses the chunks since everything
--- is backwards.
-commitsOntoChunks :: Chunks -> Commits s -> ST s Chunks
-commitsOntoChunks !xs Initial = pure xs
-commitsOntoChunks !xs (Immutable arr off len cs) =
-  commitsOntoChunks (ChunksCons (Bytes (ByteArray arr) (I# off) (I# len)) xs) cs
-commitsOntoChunks !xs (Mutable buf len cs) = case len of
-  0# -> commitsOntoChunks xs cs
-  _ -> do
-    shrinkMutableByteArray (MutableByteArray buf) (I# len)
-    arr <- PM.unsafeFreezeByteArray (MutableByteArray buf)
-    commitsOntoChunks (ChunksCons (Bytes arr 0 (I# len)) xs) cs
+  reverseCommitsOntoChunks ChunksNil cs
 
 -- | Convert a bounded builder to an unbounded one. If the size
 -- is a constant, use @Arithmetic.Nat.constant@ as the first argument
@@ -623,10 +609,6 @@ char c = fromBounded Nat.constant (Bounded.char c)
 unST :: ST s a -> State# s -> (# State# s, a #)
 unST (ST f) = f
 
-shrinkMutableByteArray :: MutableByteArray s -> Int -> ST s ()
-shrinkMutableByteArray (MutableByteArray arr) (I# sz) =
-  primitive_ (Exts.shrinkMutableByteArray# arr sz)
-
 -- | Requires exactly 8 bytes. Dump the octets of a 64-bit
 -- signed integer in a little-endian fashion.
 int64LE :: Int64 -> Builder

src/Data/ByteArray/Builder/Unsafe.hs

@@ -10,6 +10,8 @@ module Data.ByteArray.Builder.Unsafe
   ( -- * Types
     Builder(..)
   , Commits(..)
+    -- * Finalization
+  , reverseCommitsOntoChunks
     -- * Safe Functions
     -- | These functions are actually completely safe, but they are defined
     -- here because they are used by typeclass instances. Import them from
@@ -18,17 +20,21 @@ module Data.ByteArray.Builder.Unsafe
   , cstring
   ) where
 
-import Data.Primitive (MutableByteArray(MutableByteArray),ByteArray)
+import Control.Monad.Primitive (primitive_)
+import Data.Bytes.Chunks (Chunks(ChunksCons))
+import Data.Bytes.Types (Bytes(Bytes))
+import Data.Primitive (MutableByteArray(..),ByteArray(..))
 import Foreign.C.String (CString)
+import GHC.Base (unpackCString#,unpackCStringUtf8#)
 import GHC.Exts ((-#),(+#),(>#))
 import GHC.Exts (Addr#,ByteArray#,MutableByteArray#,Int(I#),Ptr(Ptr))
-import GHC.Exts (IsString,Int#,State#,MutableByteArray#)
+import GHC.Exts (IsString,Int#,State#)
 import GHC.ST (ST(ST))
-import GHC.Base (unpackCString#,unpackCStringUtf8#)
 
 import qualified GHC.Exts as Exts
 import qualified Data.ByteArray.Builder.Bounded as Bounded
 import qualified Data.ByteArray.Builder.Bounded.Unsafe as UnsafeBounded
+import qualified Data.Primitive as PM
 
 -- | An unmaterialized sequence of bytes that may be pasted
 -- into a mutable byte array.
@@ -68,6 +74,24 @@ data Commits s
       !(Commits s)
   | Initial
 
+-- | Cons the chunks from a list of @Commits@ onto an initial
+-- @Chunks@ list (this argument is often @ChunksNil@). This reverses
+-- the order of the chunks, which is desirable since builders assemble
+-- @Commits@ with the chunks backwards. This performs an in-place shrink
+-- and freezes on any mutable byte arrays it encounters. Consequently,
+-- these must not be reused.
+reverseCommitsOntoChunks :: Chunks -> Commits s -> ST s Chunks
+reverseCommitsOntoChunks !xs Initial = pure xs
+reverseCommitsOntoChunks !xs (Immutable arr off len cs) =
+  reverseCommitsOntoChunks (ChunksCons (Bytes (ByteArray arr) (I# off) (I# len)) xs) cs
+reverseCommitsOntoChunks !xs (Mutable buf len cs) = case len of
+  -- Skip over empty byte arrays.
+  0# -> reverseCommitsOntoChunks xs cs
+  _ -> do
+    shrinkMutableByteArray (MutableByteArray buf) (I# len)
+    arr <- PM.unsafeFreezeByteArray (MutableByteArray buf)
+    reverseCommitsOntoChunks (ChunksCons (Bytes arr 0 (I# len)) xs) cs
+
 -- | Create a builder from a cons-list of 'Char'. These
 -- are be UTF-8 encoded.
 stringUtf8 :: String -> Builder
@@ -118,3 +142,8 @@ goCString addr buf0 off0 len0 cs0 s0 = case Exts.indexWord8OffAddr# addr 0# of
 
 unST :: ST s a -> State# s -> (# State# s, a #)
 unST (ST f) = f
+
+shrinkMutableByteArray :: MutableByteArray s -> Int -> ST s ()
+shrinkMutableByteArray (MutableByteArray arr) (I# sz) =
+  primitive_ (Exts.shrinkMutableByteArray# arr sz)
+