{-# 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
data Builder s a = Builder
!(SmallMutableArray s a)
!Int
!Int
!(Chunks a)
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)
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 ::
a
-> Builder s a
-> 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 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
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
= Int
2
freeze ::
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 <- 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"