Add newBuilderState and closeBuilderState. Exposed commitDistance.

CHANGELOG.md

@@ -10,6 +10,8 @@ to ease the migration process.
 * Add `wordLEB128` and `word64LEB128`.
 * Add `integerDec` and `naturalDec`.
 * Add `word48PaddedLowerHex`.
+* Add `newBuilderState`, `closeBuilderState`, `commitDistance`, and
+  `commitDistance1` to the unsafe module.
 
 ## 0.3.4.0 -- 2020-02-27
 

src/Data/Bytes/Builder.hs

@@ -126,7 +126,7 @@ import Control.Exception (SomeException,toException)
 import Control.Monad.ST (ST,runST)
 import Control.Monad.IO.Class (MonadIO,liftIO)
 import Data.Bits (unsafeShiftR,unsafeShiftL,xor,finiteBitSize)
-import Data.Bytes.Builder.Unsafe (Builder(Builder))
+import Data.Bytes.Builder.Unsafe (Builder(Builder),commitDistance1)
 import Data.Bytes.Builder.Unsafe (BuilderState(BuilderState),pasteIO)
 import Data.Bytes.Builder.Unsafe (Commits(Initial,Mutable,Immutable))
 import Data.Bytes.Builder.Unsafe (reverseCommitsOntoChunks)
@@ -146,7 +146,7 @@ import Data.Word (Word64,Word32,Word16,Word8)
 import Foreign.C.String (CStringLen)
 import GHC.ByteOrder (ByteOrder(BigEndian,LittleEndian),targetByteOrder)
 import GHC.Exts (Int(I#),Char(C#),Int#,State#,ByteArray#,(>=#))
-import GHC.Exts (RealWorld,MutableByteArray#,(+#),(-#),(<#))
+import GHC.Exts (RealWorld,(+#),(-#),(<#))
 import GHC.Exts ((*#))
 import GHC.Integer.Logarithms.Compat (integerLog2#)
 import GHC.IO (IO(IO),stToIO)
@@ -1004,29 +1004,6 @@ consLength32BE = consLength Nat.constant (\x -> Bounded.word32BE (fromIntegral x
 consLength64BE :: Builder -> Builder
 consLength64BE = consLength Nat.constant (\x -> Bounded.word64BE (fromIntegral x))
 
--- Internal. This is like commitDistance, but you get to supply a
--- head of the commit list that has not yet been committed.
-commitDistance1 ::
-     MutableByteArray# s -- target
-  -> Int# -- offset into target
-  -> MutableByteArray# s -- head of array
-  -> Int# -- offset into head of array
-  -> Commits s
-  -> Int#
-commitDistance1 target offTarget buf0 offBuf cs =
-  case Exts.sameMutableByteArray# target buf0 of
-    1# -> offBuf -# offTarget
-    _ -> commitDistance target offBuf cs -# offTarget
-
-commitDistance :: MutableByteArray# s -> Int# -> Commits s -> Int#
-commitDistance _ !_ Initial = error "chunkDistance: chunk not found"
-commitDistance target !n (Immutable _ _ len cs) =
-  commitDistance target (n +# len) cs
-commitDistance target !n (Mutable buf len cs) =
-  case Exts.sameMutableByteArray# target buf of
-    1# -> n +# len
-    _ -> commitDistance target (n +# len) cs
-
 -- | Push the buffer currently being filled onto the chunk list,
 -- allocating a new active buffer of the requested size. This is
 -- helpful when a small builder is sandwhiched between two large

src/Data/Bytes/Builder/Unsafe.hs

@@ -16,11 +16,17 @@ module Data.Bytes.Builder.Unsafe
   , pasteIO
     -- * Construction
   , fromEffect
+    -- * Builder State
+  , newBuilderState
+  , closeBuilderState
     -- * Finalization
   , reverseCommitsOntoChunks
   , commitsOntoChunks
   , copyReverseCommits
   , addCommitsLength
+    -- * Commit Distance
+  , commitDistance
+  , commitDistance1
     -- * Safe Functions
     -- | These functions are actually completely safe, but they are defined
     -- here because they are used by typeclass instances. Import them from
@@ -58,12 +64,29 @@ newtype Builder
       (# State# s, MutableByteArray# s, Int#, Int#, Commits s #) -- all the same things
     )
 
+-- | A list of committed chunks along with the chunk currently being
+-- written to. This is kind of like a non-empty variant of 'Commmits'
+-- but with the additional invariant that the head chunk is a mutable
+-- byte array.
 data BuilderState s = BuilderState
   (MutableByteArray# s) -- buffer we are currently writing to
   Int# -- offset into the current buffer
   Int# -- number of bytes remaining in the current buffer
   !(Commits s) -- buffers and immutable byte slices that are already committed
 
+-- | Create an empty 'BuilderState' with a buffer of the given size.
+newBuilderState :: Int -> ST s (BuilderState s)
+{-# inline newBuilderState #-}
+newBuilderState n@(I# n# ) = do
+  MutableByteArray buf <- PM.newByteArray n
+  pure (BuilderState buf 0# n# Initial)
+
+-- | Push the active chunk onto the top of the commits.
+-- The @BuilderState@ argument must not be reused after being passed
+-- to this function. That is, its use must be affine.
+closeBuilderState :: BuilderState s -> Commits s
+closeBuilderState (BuilderState dst off _ cmts) = Mutable dst off cmts
+
 -- | Run a builder, performing an in-place update on the state.
 -- The @BuilderState@ argument must not be reused after being passed
 -- to this function. That is, its use must be affine.
@@ -258,3 +281,30 @@ unST (ST f) = f
 shrinkMutableByteArray :: MutableByteArray s -> Int -> ST s ()
 shrinkMutableByteArray (MutableByteArray arr) (I# sz) =
   primitive_ (Exts.shrinkMutableByteArray# arr sz)
+
+-- | Variant of commitDistance where you get to supply a
+-- head of the commit list that has not yet been committed.
+commitDistance1 ::
+     MutableByteArray# s -- target
+  -> Int# -- offset into target
+  -> MutableByteArray# s -- head of array
+  -> Int# -- offset into head of array
+  -> Commits s
+  -> Int#
+commitDistance1 target offTarget buf0 offBuf cs =
+  case Exts.sameMutableByteArray# target buf0 of
+    1# -> offBuf -# offTarget
+    _ -> commitDistance target offBuf cs -# offTarget
+
+-- | Compute the number of bytes between the last byte and the offset
+-- specified in a chunk. Precondition: the chunk must exist in the
+-- list of committed chunks. This relies on mutable byte arrays having
+-- identity (e.g. it uses @sameMutableByteArray#@).
+commitDistance :: MutableByteArray# s -> Int# -> Commits s -> Int#
+commitDistance !_ !_ Initial = errorWithoutStackTrace "chunkDistance: chunk not found"
+commitDistance target !n (Immutable _ _ len cs) =
+  commitDistance target (n +# len) cs
+commitDistance target !n (Mutable buf len cs) =
+  case Exts.sameMutableByteArray# target buf of
+    1# -> n +# len
+    _ -> commitDistance target (n +# len) cs