From 1f2aa9b110b56114ceacdc08193a7dcea02c8223 Mon Sep 17 00:00:00 2001
From: Andrew Martin <andrew.thaddeus@gmail.com>
Date: Sun, 23 Feb 2020 14:10:26 -0500
Subject: [PATCH] Add reversedOnto and commitsOntoChunks

---
 CHANGELOG.md                         |  1 +
 src/Data/ByteArray/Builder.hs        | 16 ++++++++++++++++
 src/Data/ByteArray/Builder/Unsafe.hs | 27 ++++++++++++++++++++++++++-
 3 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f57b3f0..956bd71 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,7 @@
 ## 0.3.4.0 -- 2020-??-??
 
 * Add `wordPaddedDec4`.
+* Add `reversedOnto` and `commitsOntoChunks`.
 
 ## 0.3.3.0 -- 2020-02-10
 
diff --git a/src/Data/ByteArray/Builder.hs b/src/Data/ByteArray/Builder.hs
index af89701..5cc9b9c 100644
--- a/src/Data/ByteArray/Builder.hs
+++ b/src/Data/ByteArray/Builder.hs
@@ -14,6 +14,7 @@ module Data.ByteArray.Builder
     -- * Evaluation
   , run
   , runOnto
+  , reversedOnto
   , putMany
   , putManyConsLength
     -- * Materialized Byte Sequences
@@ -113,6 +114,7 @@ import Data.ByteArray.Builder.Unsafe (Builder(Builder))
 import Data.ByteArray.Builder.Unsafe (BuilderState(BuilderState),pasteIO)
 import Data.ByteArray.Builder.Unsafe (Commits(Initial,Mutable,Immutable))
 import Data.ByteArray.Builder.Unsafe (reverseCommitsOntoChunks)
+import Data.ByteArray.Builder.Unsafe (commitsOntoChunks)
 import Data.ByteArray.Builder.Unsafe (stringUtf8,cstring)
 import Data.ByteArray.Builder.Unsafe (addCommitsLength,copyReverseCommits)
 import Data.ByteString.Short.Internal (ShortByteString(SBS))
@@ -160,6 +162,20 @@ runOnto hint@(I# hint# ) (Builder f) cs0 = runST $ do
       (# s1, Mutable bufX offX csX #)
   reverseCommitsOntoChunks cs0 cs
 
+-- | Variant of 'runOnto' that conses the additional chunks
+-- in reverse order. 
+reversedOnto :: 
+     Int -- ^ Size of initial chunk (use 4080 if uncertain)
+  -> Builder -- ^ Builder
+  -> Chunks
+  -> Chunks
+reversedOnto hint@(I# hint# ) (Builder f) cs0 = runST $ do
+  MutableByteArray buf0 <- PM.newByteArray hint
+  cs <- ST $ \s0 -> case f buf0 0# hint# Initial s0 of
+    (# s1, bufX, offX, _, csX #) ->
+      (# s1, Mutable bufX offX csX #)
+  commitsOntoChunks cs0 cs
+
 -- | Run a builder against lots of elements. This fills the same
 -- underlying buffer over and over again. Do not let the argument to
 -- the callback escape from the callback (i.e. do not write it to an
diff --git a/src/Data/ByteArray/Builder/Unsafe.hs b/src/Data/ByteArray/Builder/Unsafe.hs
index 54a480e..b4fd27f 100644
--- a/src/Data/ByteArray/Builder/Unsafe.hs
+++ b/src/Data/ByteArray/Builder/Unsafe.hs
@@ -18,6 +18,7 @@ module Data.ByteArray.Builder.Unsafe
   , fromEffect
     -- * Finalization
   , reverseCommitsOntoChunks
+  , commitsOntoChunks
   , copyReverseCommits
   , addCommitsLength
     -- * Safe Functions
@@ -115,7 +116,7 @@ addCommitsLength !acc (Mutable _ x cs) = addCommitsLength (acc + I# x) cs
 -- @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,
+-- and freezes any mutable byte arrays it encounters. Consequently,
 -- these must not be reused.
 reverseCommitsOntoChunks :: Chunks -> Commits s -> ST s Chunks
 reverseCommitsOntoChunks !xs Initial = pure xs
@@ -129,6 +130,30 @@ reverseCommitsOntoChunks !xs (Mutable buf len cs) = case len of
     arr <- PM.unsafeFreezeByteArray (MutableByteArray buf)
     reverseCommitsOntoChunks (ChunksCons (Bytes arr 0 (I# len)) xs) cs
 
+-- | Variant of 'reverseCommitsOntoChunks' that does not reverse
+-- the order of the commits. Since commits are built backwards by
+-- consing, this means that the chunks appended to the front will
+-- be backwards. Within each chunk, however, the bytes will be in
+-- the correct order.
+--
+-- Unlike 'reverseCommitsOntoChunks', this function is not tail
+-- recursive.
+commitsOntoChunks :: Chunks -> Commits s -> ST s Chunks
+commitsOntoChunks !xs0 cs0 = go cs0
+  where
+  go Initial = pure xs0
+  go (Immutable arr off len cs) = do
+    xs <- go cs
+    pure $! ChunksCons (Bytes (ByteArray arr) (I# off) (I# len)) xs
+  go (Mutable buf len cs) = case len of
+    -- Skip over empty byte arrays.
+    0# -> go cs
+    _ -> do
+      shrinkMutableByteArray (MutableByteArray buf) (I# len)
+      arr <- PM.unsafeFreezeByteArray (MutableByteArray buf)
+      xs <- go cs
+      pure $! ChunksCons (Bytes arr 0 (I# len)) xs
+
 -- | Copy the contents of the chunks into a mutable array, reversing
 -- the order of the chunks.
 -- Precondition: The destination must have enough space to house the