From e9de3f30db77d73238320d794e584a788f0a0420 Mon Sep 17 00:00:00 2001
From: Andrew Martin <andrew.thaddeus@gmail.com>
Date: Thu, 27 Feb 2020 12:17:23 -0500
Subject: [PATCH] Add copy2 to help avoid multiple length checks

---
 CHANGELOG.md                  |  1 +
 src/Data/ByteArray/Builder.hs | 20 ++++++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index bef093f..9f68755 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,7 @@
 * Add `reversedOnto` and `commitsOntoChunks`.
 * Add `ascii(2|3|4|5|6)`.
 * Add `cstringLen` unbounded builder.
+* Add `copy2`.
 
 ## 0.3.3.0 -- 2020-02-10
 
diff --git a/src/Data/ByteArray/Builder.hs b/src/Data/ByteArray/Builder.hs
index 2b96d33..5aab322 100644
--- a/src/Data/ByteArray/Builder.hs
+++ b/src/Data/ByteArray/Builder.hs
@@ -20,6 +20,7 @@ module Data.ByteArray.Builder
     -- * Materialized Byte Sequences
   , bytes
   , copy
+  , copy2
   , insert
   , byteArray
   , shortByteString
@@ -384,6 +385,25 @@ cstringLen (Exts.Ptr src#, I# slen# ) = Builder
   where
   !(I# newSz) = max (I# slen#) 4080
 
+-- | Create a builder from two byte sequences. This always results in two
+-- calls to @memcpy@. This is beneficial when the byte sequences are
+-- known to be small (less than 256 bytes).
+copy2 :: Bytes -> Bytes -> Builder
+copy2 (Bytes (ByteArray srcA# ) (I# soffA# ) (I# slenA# ))
+      (Bytes (ByteArray srcB# ) (I# soffB# ) (I# slenB# )) = Builder
+  (\buf0 off0 len0 cs0 s0 -> case len0 <# slen# of
+    1# -> case Exts.newByteArray# newSz s0 of
+        (# s1, buf1 #) -> case Exts.copyByteArray# srcA# soffA# buf1 0# slenA# s1 of
+          s2 -> case Exts.copyByteArray# srcB# soffB# buf1 slenA# slenB# s2 of
+            s3 -> (# s3, buf1, slen#, newSz -# slen#, Mutable buf0 off0 cs0 #)
+    _ -> let !s1 = Exts.copyByteArray# srcA# soffA# buf0 off0 slenA# s0
+             !s2 = Exts.copyByteArray# srcB# soffB# buf0 (off0 +# slenA# ) slenB# s1 in
+      (# s2, buf0, off0 +# slen#, len0 -# slen#, cs0 #)
+  )
+  where
+  !slen# = slenA# +# slenB#
+  !(I# newSz) = max (I# slen#) 4080
+
 -- | Create a builder from a byte sequence. This never calls @memcpy@.
 -- Instead, it pushes a chunk that references the argument byte sequence.
 -- This wastes the remaining space in the active chunk, so it may adversely