<commit_before>module Main.BinaryTreeWithSize where

import Main.Prelude hiding (traverse_, fold, empty)
import Foreign
import qualified Data.ByteString as A
import qualified Data.Vector.Storable.Mutable as B


data Builder =
  Builder !Int (BinaryTree Bytes)

instance Monoid Builder where
  mempty = empty
  mappend = append

data BinaryTree a =
  Void |
  Leaf a |
  Branch (BinaryTree a) (BinaryTree a)
  deriving (Functor, Foldable, Traversable)

empty :: Builder
empty =
  Builder 0 (Void)

append :: Builder -> Builder -> Builder
append (Builder length1 tree1) (Builder length2 tree2) =
  Builder (length1 + length2) (Branch tree1 tree2)

bytes :: Bytes -> Builder
bytes bytes =
  Builder (A.length bytes) (Leaf bytes)

lengthOf :: Builder -> Int
lengthOf (Builder length tree) =
  length

bytesOf :: Builder -> Bytes
bytesOf (Builder length tree) =
  fill length tree

fill len t0 = unsafePerformIO $ do
  ptr <- mallocArray len
  let
    go !i (Branch l r) = do
      ll <- go i l
      go ll r

    go !i (Leaf a) = writeBS 0 i
      where
        writeBS !j !ij
          | j < A.length a = pokeElemOff ptr ij (A.index a j) >> writeBS (j+1) (ij+1)
          | otherwise      = return ij

    go !i Void    = return 0

  go 0 t0
  A.packCStringLen (castPtr ptr, len)<commit_msg>Update BinaryTreeWithSize<commit_after>module Main.BinaryTreeWithSize where

import Main.Prelude hiding (traverse_, fold, empty)
import qualified Data.ByteString as A
import qualified Data.ByteString.Internal as B
import qualified Main.Bytes as C
import qualified Foreign as D


data Builder =
  Builder !Int (BinaryTree Bytes)

instance Monoid Builder where
  mempty = empty
  mappend = append

instance IsString Builder where
  fromString = bytes . fromString

data BinaryTree a =
  Void |
  Leaf a |
  Branch (BinaryTree a) (BinaryTree a)
  deriving (Functor, Foldable, Traversable)

empty :: Builder
empty =
  Builder 0 (Void)

append :: Builder -> Builder -> Builder
append (Builder length1 tree1) (Builder length2 tree2) =
  Builder (length1 + length2) (Branch tree1 tree2)

bytes :: Bytes -> Builder
bytes bytes =
  Builder (A.length bytes) (Leaf bytes)

lengthOf :: Builder -> Int
lengthOf (Builder length tree) =
  length

-- |
-- FIXME: Seriously needs some optimization!
-- It must perform better than \"thruList\", not worse.
-- 
-- See
-- http://hackage.haskell.org/package/bytestring-0.10.6.0/docs/src/Data.ByteString.Internal.html#unsafeCreate
bytesOf :: Builder -> Bytes
bytesOf (Builder length tree) =
  B.unsafeCreate length $ void . pokeTree tree

pokeTree :: BinaryTree Bytes -> D.Ptr Word8 -> IO (D.Ptr Word8)
pokeTree tree ptr =
  case tree of
    Void -> pure ptr
    Leaf bytes -> C.poke bytes ptr
    Branch tree1 tree2 -> pokeTree tree1 ptr >>= pokeTree tree2
{'fraction_non_alphanumeric': 0.09159754778218536, 'fraction_numerical': 0.01298232960692391, 'mean_word_length': 3.361635220125786, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 22, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7642132', 'n_tokens_mistral': 899, 'n_tokens_neox': 881, 'n_words': 324}