{-# 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 :: forall s a. ST s (Builder s a)
new = do
  SmallMutableArray s a
marr <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
initialLength forall a. a
errorThunk
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall s a.
SmallMutableArray s a -> Int -> Int -> Chunks a -> Builder s a
Builder SmallMutableArray s a
marr Int
0 Int
initialLength 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 <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
initialLength a
a0
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall s a.
SmallMutableArray s a -> Int -> Int -> Chunks a -> Builder s a
Builder SmallMutableArray s a
marr Int
1 Int
initialLength 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 :: 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 forall a. Ord a => a -> a -> Bool
> Int
0 of
  Bool
True -> do
    forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
marr Int
off a
a
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! forall s a.
SmallMutableArray s a -> Int -> Int -> Chunks a -> Builder s a
Builder SmallMutableArray s a
marr (Int
off forall a. Num a => a -> a -> a
+ Int
1) (Int
len forall a. Num a => a -> a -> a
- Int
1) Chunks a
cs
  Bool
False -> do
    SmallArray a
arr <- forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray SmallMutableArray s a
marr
    let lenNew :: Int
lenNew = Int -> Int
nextLength (forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr)
    SmallMutableArray s a
marrNew <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
lenNew a
a
    let !csNew :: Chunks a
csNew = forall a. SmallArray a -> Chunks a -> Chunks a
ChunksCons SmallArray a
arr Chunks a
cs
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! forall s a.
SmallMutableArray s a -> Int -> Int -> Chunks a -> Builder s a
Builder SmallMutableArray s a
marrNew Int
1 (Int
lenNew 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 forall a. Ord a => a -> a -> Bool
< Int
maxElementCount forall a. Num a => a -> a -> a
- Int
smallArrayHeaderWords
  then Int
i forall a. Num a => a -> a -> a
* Int
2 forall a. Num a => a -> a -> a
+ Int
smallArrayHeaderWords
  else Int
maxElementCount forall a. Num a => a -> a -> a
- Int
smallArrayHeaderWords

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

initialLength :: Int
initialLength :: Int
initialLength = Int
16 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 :: 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 <- forall s a. SmallMutableArray s a -> Int -> ST s (SmallArray a)
unsafeShrinkAndFreeze SmallMutableArray s a
marr Int
off
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! forall a. Chunks a -> Chunks a -> Chunks a
C.reverseOnto (forall a. SmallArray a -> Chunks a -> Chunks a
ChunksCons SmallArray a
arr forall a. Chunks a
ChunksNil) Chunks a
cs

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