{-# LANGUAGE BangPatterns #-}

module Data.Builder.ST
  ( Builder (..)
  , new
  , new1
  , push
  , freeze
  ) where

import Compat (unsafeShrinkAndFreeze)
import Control.Monad.ST (ST)
import Data.Chunks (Chunks (ChunksCons, ChunksNil))
import Data.Primitive (SmallMutableArray, newSmallArray, sizeofSmallArray, unsafeFreezeSmallArray, writeSmallArray)
import Foreign.Storable (sizeOf)

import qualified Data.Chunks as C

{- | Builder for an array of boxed elements. This type is appropriate
when constructing an array of unknown size in an effectful
(@ST@ or @IO@) setting. In a non-effectful setting, consider
the @Builder@ from @Data.Builder@ instead.

A 'Builder' must be used linearly. The type system does not
enforce this, so users must be careful when handling a 'Builder'.
-}
data Builder s a
  = Builder
      !(SmallMutableArray s a)
      !Int
      !Int
      !(Chunks a)

-- | Create a new 'Builder' with no elements in it.
new :: ST s (Builder s a)
new :: forall s a. ST s (Builder s a)
new = do
  SmallMutableArray s a
marr <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
initialLength a
forall a. a
errorThunk
  Builder s a -> ST s (Builder s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SmallMutableArray s a -> Int -> Int -> Chunks a -> Builder s a
forall s a.
SmallMutableArray s a -> Int -> Int -> Chunks a -> Builder s a
Builder SmallMutableArray s a
marr Int
0 Int
initialLength Chunks a
forall a. Chunks a
ChunksNil)

{- | Create a new 'Builder' with a single element. Useful when builder
creation is immidiately followed by 'push'. Note that:

> new >>= push x ≡ new1 x

But 'new1' performs slightly better.
-}
new1 :: a -> ST s (Builder s a)
new1 :: forall a s. a -> ST s (Builder s a)
new1 a
a0 = do
  SmallMutableArray s a
marr <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
initialLength a
a0
  Builder s a -> ST s (Builder s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SmallMutableArray s a -> Int -> Int -> Chunks a -> Builder s a
forall s a.
SmallMutableArray s a -> Int -> Int -> Chunks a -> Builder s a
Builder SmallMutableArray s a
marr Int
1 Int
initialLength Chunks a
forall a. Chunks a
ChunksNil)

{- | Push an element onto the end of the builder. This
is not strict in the element, so force it before pushing
it on to the builder if doing so is needed to prevent
space leaks.
-}
push ::
  -- | Element to push onto the end
  a ->
  -- | Builder, do not reuse this after pushing onto it
  Builder s a ->
  -- | New builder
  ST s (Builder s a)
push :: forall a s. a -> Builder s a -> ST s (Builder s a)
push a
a (Builder SmallMutableArray s a
marr Int
off Int
len Chunks a
cs) = case Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 of
  Bool
True -> do
    SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
marr Int
off a
a
    Builder s a -> ST s (Builder s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Builder s a -> ST s (Builder s a))
-> Builder s a -> ST s (Builder s a)
forall a b. (a -> b) -> a -> b
$! SmallMutableArray s a -> Int -> Int -> Chunks a -> Builder s a
forall s a.
SmallMutableArray s a -> Int -> Int -> Chunks a -> Builder s a
Builder SmallMutableArray s a
marr (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Chunks a
cs
  Bool
False -> do
    SmallArray a
arr <- SmallMutableArray (PrimState (ST s)) a -> ST s (SmallArray a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
marr
    let lenNew :: Int
lenNew = Int -> Int
nextLength (SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr)
    SmallMutableArray s a
marrNew <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
lenNew a
a
    let !csNew :: Chunks a
csNew = SmallArray a -> Chunks a -> Chunks a
forall a. SmallArray a -> Chunks a -> Chunks a
ChunksCons SmallArray a
arr Chunks a
cs
    Builder s a -> ST s (Builder s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Builder s a -> ST s (Builder s a))
-> Builder s a -> ST s (Builder s a)
forall a b. (a -> b) -> a -> b
$! SmallMutableArray s a -> Int -> Int -> Chunks a -> Builder s a
forall s a.
SmallMutableArray s a -> Int -> Int -> Chunks a -> Builder s a
Builder SmallMutableArray s a
marrNew Int
1 (Int
lenNew Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Chunks a
csNew

-- The sequence of sizes we create is:
--   64-bit: 6, 14, 30, 62, 126, 254, 254, 254...
--   32-bit: 6, 14, 30, 62, 126, 254, 510, 510, 510...
-- The goal is to have objects whose sizes are increasing
-- powers of 2 until we reach the size of a block (4KB).
-- A 254-element SmallArray on a 64-bit platform uses
-- exactly 4KB (header + ptrs + payload).
nextLength :: Int -> Int
nextLength :: Int -> Int
nextLength Int
i =
  if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
maxElementCount Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
smallArrayHeaderWords
    then Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
smallArrayHeaderWords
    else Int
maxElementCount Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
smallArrayHeaderWords

maxElementCount :: Int
maxElementCount :: Int
maxElementCount = Int -> Int -> Int
forall a. Integral a => a -> a -> a
div Int
4096 (Int -> Int
forall a. Storable a => a -> Int
sizeOf (Int
forall a. HasCallStack => a
undefined :: Int))

initialLength :: Int
initialLength :: Int
initialLength = Int
16 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
smallArrayHeaderWords

smallArrayHeaderWords :: Int
smallArrayHeaderWords :: Int
smallArrayHeaderWords = Int
2

{- | Convert a 'Builder' to 'Chunks'. The 'Builder' must not
be reused after this operation.
-}
freeze ::
  -- | Builder, do not reuse after freezing
  Builder s a ->
  ST s (Chunks a)
freeze :: forall s a. Builder s a -> ST s (Chunks a)
freeze (Builder SmallMutableArray s a
marr Int
off Int
_ Chunks a
cs) = do
  SmallArray a
arr <- SmallMutableArray s a -> Int -> ST s (SmallArray a)
forall s a. SmallMutableArray s a -> Int -> ST s (SmallArray a)
unsafeShrinkAndFreeze SmallMutableArray s a
marr Int
off
  Chunks a -> ST s (Chunks a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Chunks a -> ST s (Chunks a)) -> Chunks a -> ST s (Chunks a)
forall a b. (a -> b) -> a -> b
$! Chunks a -> Chunks a -> Chunks a
forall a. Chunks a -> Chunks a -> Chunks a
C.reverseOnto (SmallArray a -> Chunks a -> Chunks a
forall a. SmallArray a -> Chunks a -> Chunks a
ChunksCons SmallArray a
arr Chunks a
forall a. Chunks a
ChunksNil) Chunks a
cs

errorThunk :: a
{-# NOINLINE errorThunk #-}
errorThunk :: forall a. a
errorThunk = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"array-builder:Data.Builder.ST: error"