ValueBuffer.hs

{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

-- | Simple 'Storable'-based append buffers in 'IO' and 'ST'.
--
-- The main properties are amortised constant 'push' and constant GC cost.
module Glean.Util.ValueBuffer
  ( Buffer, IOBuffer
  , new
  , push
  , appends
  , get
  , unsafeFreeze
  )
where

import Control.Monad
import Control.Monad.Primitive
import Data.Primitive.MutVar
import Data.Vector.Generic (Vector, Mutable)
import qualified Data.Vector.Generic as V
import qualified Data.Vector.Generic.Mutable as MV

data Impl v s a = Impl
  { _implData :: !(Mutable v s a)
  , _implLen :: {-# UNPACK #-} !Int
  }

-- | Append buffer type.
newtype Buffer v s a = Buffer { _fromBuffer :: MutVar s (Impl v s a) }

-- | 'Buffer' specialised to 'IO'.
type IOBuffer v = Buffer v RealWorld

-- | Create a new buffer with the given initial capacity.
new :: (PrimMonad m, Vector v a) => Int -> m (Buffer v (PrimState m) a)
{-# INLINABLE new #-}
new cap = do
  v <- MV.new $ max cap 1
  Buffer <$> newMutVar (Impl v 0)

-- | Push an element to the buffer, growing it if necessary.
push :: (PrimMonad m, Vector v a) => Buffer v (PrimState m) a -> a -> m ()
{-# INLINABLE push #-}
push (Buffer ref) x = do
  Impl v i <- readMutVar ref
  w <- if i < MV.length v
    then return v
    else MV.unsafeGrow v (MV.length v)
  MV.unsafeWrite w i x
  writeMutVar ref $! Impl w (i+1)

-- | Append multiple batches to the buffer.
--
-- NOTE: The list will be evaluated strictly.
appends
  :: (PrimMonad m, Vector v a) => Buffer v (PrimState m) a -> [v a] -> m ()
{-# INLINABLE appends #-}
appends _ [] = return ()
appends (Buffer ref) xs = do
  Impl v i <- readMutVar ref
  let !n = sum $ map V.length xs
  !w <- if i+n <= MV.length v
    then return v
    else MV.grow v $ max (i + n - MV.length v) (MV.length v)
  let app !i x = do
        let k = V.length x
        V.copy (MV.slice i k w) x
        return (i+k)
  i' <- foldM app i xs
  writeMutVar ref $ Impl w i'

-- | Return the underlying vector.
--
-- NOTE: Subsequent changes to the buffer might mutate the vector and vice
-- versa.
get
  :: (PrimMonad m, Vector v a)
  => Buffer v (PrimState m) a
  -> m (Mutable v (PrimState m) a)
{-# INLINABLE get #-}
get (Buffer ref) = do
  Impl v i <- readMutVar ref
  return $ MV.unsafeSlice 0 i v

-- | Freeze the buffer and return all writen elements as a vector.
unsafeFreeze
  :: (PrimMonad m, Vector v a) => Buffer v (PrimState m) a -> m (v a)
{-# INLINABLE unsafeFreeze #-}
unsafeFreeze = V.unsafeFreeze <=< get