From 4f8ea31b731ccd76a3130614e479caf87c9b62a9 Mon Sep 17 00:00:00 2001
From: Andrew Martin <andrew.thaddeus@gmail.com>
Date: Wed, 20 Nov 2019 14:30:51 -0500
Subject: [PATCH] Add Data.ByteArray.Builder.Unsafe.fromEffect

---
 CHANGELOG.md                         |  1 +
 small-bytearray-builder.cabal        |  2 +-
 src/Data/ByteArray/Builder/Unsafe.hs | 24 ++++++++++++++++++++++--
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1d5cebc..771ae75 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@
 * Export `reverseCommitsOntoChunks` from the `Unsafe` module.
 * Add `Semigroup` and `Monoid` instances for `Chunks`.
 * Add `consLengthLE32`.
+* Add `fromEffect` to the unsafe interface.
 
 ## 0.3.0.0 -- 2019-10-17
 
diff --git a/small-bytearray-builder.cabal b/small-bytearray-builder.cabal
index ca4e1e6..7ee0f9b 100644
--- a/small-bytearray-builder.cabal
+++ b/small-bytearray-builder.cabal
@@ -1,6 +1,6 @@
 cabal-version: 2.2
 name: small-bytearray-builder
-version: 0.3.0.0
+version: 0.3.1.0
 synopsis: Serialize to a small byte arrays
 description:
   This is similar to the builder facilities provided by
diff --git a/src/Data/ByteArray/Builder/Unsafe.hs b/src/Data/ByteArray/Builder/Unsafe.hs
index e7a4ac8..4cba197 100644
--- a/src/Data/ByteArray/Builder/Unsafe.hs
+++ b/src/Data/ByteArray/Builder/Unsafe.hs
@@ -10,6 +10,8 @@ module Data.ByteArray.Builder.Unsafe
   ( -- * Types
     Builder(..)
   , Commits(..)
+    -- * Construction
+  , fromEffect
     -- * Finalization
   , reverseCommitsOntoChunks
     -- * Safe Functions
@@ -26,7 +28,7 @@ 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 ((-#),(+#),(>#),(>=#))
 import GHC.Exts (Addr#,ByteArray#,MutableByteArray#,Int(I#),Ptr(Ptr))
 import GHC.Exts (IsString,Int#,State#)
 import GHC.ST (ST(ST))
@@ -119,6 +121,8 @@ goString (c : cs) buf0 off0 len0 cs0 s0 = case len0 ># 3# of
 -- We have to have a rule for both unpackCString# and unpackCStringUtf8#
 -- since GHC uses a different function based on whether or not non-ASCII
 -- codepoints are used in the string.
+-- TODO: The UTF-8 variant of this rule is unsound because GHC actually
+-- used Modified UTF-8.
 {-# RULES
 "Builder stringUtf8/cstring" forall s a b c d e.
   goString (unpackCString# s) a b c d e = goCString s a b c d e
@@ -140,10 +144,26 @@ goCString addr buf0 off0 len0 cs0 s0 = case Exts.indexWord8OffAddr# addr 0# of
     _ -> case Exts.writeWord8Array# buf0 off0 w s0 of
       s1 -> goCString (Exts.plusAddr# addr 1# ) buf0 (off0 +# 1# ) (len0 -# 1# ) cs0 s1
 
+fromEffect ::
+     Int -- ^ Maximum number of bytes the paste function needs
+  -> (forall s. MutableByteArray s -> Int -> ST s Int)
+     -- ^ Paste function. Takes a byte array and an offset and returns
+     -- the new offset and having pasted into the buffer.
+  -> Builder
+{-# inline fromEffect #-}
+fromEffect (I# req) f = Builder $ \buf0 off0 len0 cs0 s0 ->
+  let !(# s1, buf1, off1, len1, cs1 #) = case len0 >=# req of
+        1# -> (# s0, buf0, off0, len0, cs0 #)
+        _ -> let !(I# lenX) = max 4080 (I# req) in
+          case Exts.newByteArray# lenX s0 of
+            (# sX, bufX #) ->
+              (# sX, bufX, 0#, lenX, Mutable buf0 off0 cs0 #)
+   in case unST (f (MutableByteArray buf1) (I# off1)) s1 of
+        (# s2, I# off2 #) -> (# s2, buf1, off2, len1 -# (off2 -# off1), cs1 #)
+
 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)
-