Block/UArray/Builder: concat performance update

basement/Basement/Block/Base.hs

@@ -231,29 +231,26 @@ append a b
     !la = lengthBytes a
     !lb = lengthBytes b
 
-concat :: [Block ty] -> Block ty
-concat [] = empty
-concat l  =
-    case filterAndSum 0 [] l of
-        (_,[])            -> empty
-        (_,[x])           -> x
-        (totalLen,chunks) -> runST $ do
-            r <- unsafeNew Unpinned totalLen
-            doCopy r 0 chunks
-            unsafeFreeze r
+concat :: forall ty . [Block ty] -> Block ty
+concat original = runST $ do
+    r <- unsafeNew Unpinned total
+    goCopy r zero original
+    unsafeFreeze r
   where
-    -- TODO would go faster not to reverse but pack from the end instead
-    filterAndSum !totalLen acc []     = (totalLen, Data.List.reverse acc)
-    filterAndSum !totalLen acc (x:xs)
-        | len == 0  = filterAndSum totalLen acc xs
-        | otherwise = filterAndSum (len+totalLen) (x:acc) xs
-      where len = lengthBytes x
-
-    doCopy _ _ []     = return ()
-    doCopy r i (x:xs) = do
-        unsafeCopyBytesRO r i x 0 lx
-        doCopy r (i `offsetPlusE` lx) xs
-      where !lx = lengthBytes x
+    !total = size 0 original
+    -- size
+    size !sz []     = sz
+    size !sz (x:xs) = size (lengthBytes x + sz) xs
+
+    zero = Offset 0
+
+    goCopy r = loop
+      where
+        loop _  []      = pure ()
+        loop !i (x:xs) = do
+            unsafeCopyBytesRO r i x zero lx
+            loop (i `offsetPlusE` lx) xs
+          where !lx = lengthBytes x
 
 -- | Freeze a mutable block into a block.
 --

basement/Basement/Block/Builder.hs

@@ -60,6 +60,8 @@ instance Monoid Builder where
     {-# INLINE mempty #-}
     mappend = append
     {-# INLINABLE mappend #-}
+    mconcat = concat
+    {-# INLINABLE mconcat #-}
 
 -- | create an empty builder
 --
@@ -79,6 +81,15 @@ append (Builder size1 (Action action1)) (Builder size2 (Action action2)) =
     size = size1 + size2
 {-# INLINABLE append #-}
 
+-- | concatenate the list of builder
+concat :: [Builder] -> Builder
+concat = loop 0 (Action $ \_ !off -> pure off)
+  where
+    loop !sz acc          []                              = Builder sz acc
+    loop !sz (Action acc) (Builder !s (Action action):xs) =
+       loop (sz + s) (Action $ \arr off -> acc arr off >>= action arr) xs
+{-# INLINABLE concat #-}
+
 -- | run the given builder and return the generated block
 run :: PrimMonad prim => Builder -> prim (Block Word8)
 run (Builder sz action) = do

basement/Basement/UArray/Base.hs

@@ -605,29 +605,26 @@ append a b
     !la = length a
     !lb = length b
 
-concat :: PrimType ty => [UArray ty] -> UArray ty
-concat [] = empty
-concat l  =
-    case filterAndSum (CountOf 0) [] l of
-        (_,[])            -> empty
-        (_,[x])           -> x
-        (totalLen,chunks) -> runST $ do
-            r <- new totalLen
-            doCopy r (Offset 0) chunks
-            unsafeFreeze r
+concat :: forall ty . PrimType ty => [UArray ty] -> UArray ty
+concat original = runST $ do
+    r <- new total
+    goCopy r 0 original
+    unsafeFreeze r
   where
-    -- TODO would go faster not to reverse but pack from the end instead
-    filterAndSum !totalLen acc []     = (totalLen, List.reverse acc)
-    filterAndSum !totalLen acc (x:xs)
-        | len == CountOf 0 = filterAndSum totalLen acc xs
-        | otherwise      = filterAndSum (len+totalLen) (x:acc) xs
-      where len = length x
-
-    doCopy _ _ []     = return ()
-    doCopy r i (x:xs) = do
-        unsafeCopyAtRO r i x (Offset 0) lx
-        doCopy r (i `offsetPlusE` lx) xs
-      where lx = length x
+    !total = size 0 original
+    -- size
+    size !sz []     = sz
+    size !sz (x:xs) = size (length x + sz) xs
+
+    zero = Offset 0
+
+    goCopy r = loop
+      where
+        loop _  []      = pure ()
+        loop !i (x:xs) = do
+            unsafeCopyAtRO r i x zero lx
+            loop (i `offsetPlusE` lx) xs
+          where !lx = length x
 
 -- | Create a Block from a UArray.
 --

benchs/Fake/ByteString.hs

@@ -17,6 +17,7 @@ module Fake.ByteString
     , readInt
     , readInteger
     , unpack
+    , concat
     ) where
 
 import Prelude (undefined, Maybe(..))
@@ -41,6 +42,8 @@ foldr _ _ _ = undefined
 and _ _ = undefined
 all _ _ = undefined
 any _ _ = undefined
+concat :: [ByteString] -> ByteString
+concat _ = undefined
 unpack :: ByteString -> [Word8]
 unpack = undefined
 

benchs/Fake/Vector.hs

@@ -15,6 +15,7 @@ module Fake.Vector
     , and
     , all
     , any
+    , concat
     ) where
 
 import Prelude (undefined)
@@ -40,3 +41,4 @@ foldr _ _ _ = undefined
 and _ _ = undefined
 all _ _ = undefined
 any _ _ = undefined
+concat = undefined

benchs/Main.hs

@@ -13,6 +13,7 @@ import Foundation.String.Read
 import Foundation.String
 import BenchUtil.Common
 import BenchUtil.RefData
+import qualified Basement.Block.Builder as Builder
 
 import Sys
 
@@ -179,6 +180,8 @@ benchsByteArray = bgroup "ByteArray"
     , benchFoldr
     , benchReverse
     , benchFilter
+    , benchMonoidConcat
+    , benchBuilderBlock
     , benchAll
     , benchSort
     , benchSort32
@@ -204,11 +207,36 @@ benchsByteArray = bgroup "ByteArray"
         t = ByteString.pack dat
         v = Vector.fromList dat
 
+    diffListByteArray :: ([UArray Word8] -> a)
+                      -> ([Block Word8] -> b)
+                      -> ([ByteString.ByteString] -> c)
+                      -> ([Vector.Vector Word8] -> d)
+                      -> [[Word8]]
+                      -> [Benchmark]
+    diffListByteArray uarrayBench blockBench bsBench vectorBench dat =
+        [ bench "[UArray_W8]" $ whnf uarrayBench s
+        , bench "[Block_W8]" $ whnf blockBench s'
+#ifdef BENCH_ALL
+        , bench "[ByteString]" $ whnf bsBench t
+        , bench "[Vector_W8]" $ whnf vectorBench v
+#endif
+        ]
+      where
+        s = fromList <$> dat
+        s' = fromList <$> dat
+        t = ByteString.pack <$> dat
+        v = Vector.fromList <$> dat
+
     allDat =
         [ ("bs20", rdBytes20)
         , ("bs200", rdBytes200)
         , ("bs2000", rdBytes2000)
         ]
+    allListDat =
+        [ ("listBs20", Prelude.replicate 20 rdBytes20)
+        , ("listBs200", Prelude.replicate 200 rdBytes200)
+        , ("listBs2000", Prelude.replicate 2000 rdBytes2000)
+        ]
     allDatSuffix s = fmap (first (\x -> x <> "-" <> s)) allDat
 
     benchLength = bgroup "Length" $
@@ -270,6 +298,15 @@ benchsByteArray = bgroup "ByteArray"
                                                     (ByteString.filter (> 100))
                                                     (Vector.filter (> 100)) dat) allDat
 
+    benchMonoidConcat = bgroup "Monoid/mconcat" $
+        fmap (\(n, dat) -> bgroup n $ diffListByteArray mconcat mconcat ByteString.concat Vector.concat dat) allListDat
+    benchBuilderBlock = bgroup "Monoid/builder" $
+        [ bench "[block Word8]" $ whnf builderConcat (Prelude.replicate 2000 (fromList rdBytes2000))
+        ]
+      where
+        builderConcat :: [Block Word8] -> Block Word8
+        builderConcat l = runST $ Builder.run $ mconcat (fmap Builder.emit l)
+
     benchSort = bgroup "Sort" $ fmap (\(n, dat) ->
         bgroup n $
             [ bench "UArray_W8" $ whnf uarrayBench (fromList dat)