{-# language BangPatterns #-}

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

import Compat (unsafeShrinkAndFreeze)
import Control.Monad.ST (ST)
import Data.Chunks (Chunks(ChunksNil,ChunksCons))
import Data.Primitive (SmallMutableArray)
import Data.Primitive (newSmallArray,writeSmallArray,unsafeFreezeSmallArray)
import Data.Primitive (sizeofSmallArray)
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 :: 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 (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 :: 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 (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 ::
     a -- ^ Element to push onto the end
  -> Builder s a -- ^ Builder, do not reuse this after pushing onto it
  -> ST s (Builder s a) -- ^ New builder
push :: 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 (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 (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 s a -- ^ Builder, do not reuse after freezing
  -> ST s (Chunks a)
freeze :: 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 (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 :: a
errorThunk = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"array-builder:Data.Builder.ST: error"