From 01bf4655e36d024243b664c91a06b11fee0f6d00 Mon Sep 17 00:00:00 2001
From: Andrew Martin <andrew.thaddeus@gmail.com>
Date: Sun, 22 Sep 2019 09:20:03 -0400
Subject: [PATCH] add benchmarks for encoding tree data structures

---
 bench/Main.hs                 | 10 +++++
 common/Word16Tree.hs          | 83 +++++++++++++++++++++++++++++++++++
 small-bytearray-builder.cabal |  3 ++
 test/Main.hs                  |  4 ++
 4 files changed, 100 insertions(+)
 create mode 100644 common/Word16Tree.hs

diff --git a/bench/Main.hs b/bench/Main.hs
index c4fefe4..d576496 100644
--- a/bench/Main.hs
+++ b/bench/Main.hs
@@ -13,6 +13,7 @@ import qualified Data.ByteArray.Builder.Bounded as U
 import qualified Cell
 import qualified SimpleCsv
 import qualified HexWord64
+import qualified Word16Tree
 
 main :: IO ()
 main = defaultMain
@@ -26,6 +27,15 @@ main = defaultMain
     [ bench "csv-no-escape" $ whnf
         (\x -> B.run 4080 (SimpleCsv.encodeRows x))
         Cell.cells
+    , bench "word-16-tree-small" $ whnf
+        (\x -> B.run 4080 (Word16Tree.encode x))
+        Word16Tree.exampleSmall
+    , bench "word-16-tree-2000" $ whnf
+        (\x -> B.run ((4096 * 16) - 16) (Word16Tree.encode x))
+        Word16Tree.example2000
+    , bench "word-16-tree-9000" $ whnf
+        (\x -> B.run ((4096 * 64) - 16) (Word16Tree.encode x))
+        Word16Tree.example9000
     ]
   ]
 
diff --git a/common/Word16Tree.hs b/common/Word16Tree.hs
new file mode 100644
index 0000000..ed042e0
--- /dev/null
+++ b/common/Word16Tree.hs
@@ -0,0 +1,83 @@
+{-# language BangPatterns #-}
+
+module Word16Tree
+  ( Word16Tree
+  , encode
+  , exampleSmall
+  , example2000
+  , example9000
+  , expectedSmall
+  ) where
+
+import Data.ByteArray.Builder as B
+import Data.Word (Word16)
+import Data.Primitive (ByteArray)
+import qualified Data.Bytes as Bytes
+
+data Word16Tree
+  = Branch !Word16Tree !Word16Tree
+  | Leaf {-# UNPACK #-} !Word16
+
+encode :: Word16Tree -> Builder
+encode (Leaf w) = B.word16PaddedUpperHex w
+encode (Branch a b) =
+  B.ascii '('
+  <>
+  encode a
+  <>
+  B.ascii ','
+  <>
+  encode b
+  <>
+  B.ascii ')'
+
+expectedSmall :: ByteArray
+expectedSmall = Bytes.toByteArray $ Bytes.fromAsciiString
+  "((AB59,(1F33,2E71)),((((FA9A,247B),890C),(0F13,((55BF,7CF1),389B))),1205))"
+
+exampleSmall :: Word16Tree
+exampleSmall = Branch
+  (Branch
+    (Leaf 0xAB59)
+    (Branch
+      (Leaf 0x1F33)
+      (Leaf 0x2E71)
+    )
+  )
+  (Branch
+    (Branch 
+      (Branch
+        (Branch
+          (Leaf 0xFA9A)
+          (Leaf 0x247B)
+        )
+        (Leaf 0x890C)
+      )
+      (Branch
+        (Leaf 0x0F13)
+        (Branch
+          (Branch
+            (Leaf 0x55BF)
+            (Leaf 0x7CF1)
+          )
+          (Leaf 0x389B)
+        )
+      )
+    )
+    (Leaf 0x1205)
+  )
+
+example2000 :: Word16Tree
+{-# noinline example2000 #-}
+example2000 = balanced 0 2000
+
+example9000 :: Word16Tree
+{-# noinline example9000 #-}
+example9000 = balanced 0 9000
+
+balanced :: Word16 -> Word16 -> Word16Tree
+balanced !off !n
+  | n == 0 = Leaf off
+  | n == 1 = Leaf (off + 1)
+  | otherwise = let x = div n 2 in
+      Branch (balanced off x) (balanced (off + x) (n - x))
diff --git a/small-bytearray-builder.cabal b/small-bytearray-builder.cabal
index 6025a5c..05f70ea 100644
--- a/small-bytearray-builder.cabal
+++ b/small-bytearray-builder.cabal
@@ -65,6 +65,7 @@ test-suite test
   ghc-options: -O2 -Wall
   other-modules:
     HexWord64
+    Word16Tree
   build-depends:
     , QuickCheck >=2.13.1 && <2.14
     , base >=4.12.0.0 && <5
@@ -88,6 +89,7 @@ benchmark bench
     , primitive
     , small-bytearray-builder
     , text-short
+    , byteslice
   ghc-options: -Wall -O2
   default-language: Haskell2010
   hs-source-dirs: bench, common
@@ -96,3 +98,4 @@ benchmark bench
     Cell
     HexWord64
     SimpleCsv
+    Word16Tree
diff --git a/test/Main.hs b/test/Main.hs
index 6f27852..ec602b9 100644
--- a/test/Main.hs
+++ b/test/Main.hs
@@ -28,6 +28,7 @@ import qualified Test.Tasty.HUnit as THU
 import qualified Test.Tasty.QuickCheck as TQC
 
 import qualified HexWord64
+import qualified Word16Tree
 
 main :: IO ()
 main = defaultMain tests
@@ -102,6 +103,9 @@ tests = testGroup "Tests"
         pack ("\"Hi\\r\\nLo\"") @=? run 1 (shortTextJsonString "Hi\r\nLo")
     , THU.testCase "shortTextJsonString-D" $
         pack ("\"Hi\\u001BLo\"") @=? run 1 (shortTextJsonString "Hi\ESCLo")
+    , THU.testCase "word-16-tree" $
+        Word16Tree.expectedSmall @=? run 1
+          (Word16Tree.encode Word16Tree.exampleSmall)
     ]
   , testGroup "alternate"
     [ TQC.testProperty "HexWord64" $ \x y ->