Groan.hs

{-# LANGUAGE Strict #-}

module Data.Vector.Groan (
  Buffer,

  -- * Running computations that allocate to a buffer
  run,
  runPrim,

  -- * Creating buffers
  new,
  replicate,
  thaw,

  -- * Reading
  read,

  -- ** Size and capacity
  size,
  capacity,

  -- * Mutable operations

  -- ** Pop
  popBack,

  -- ** Writes
  write,

  -- *** Appends
  pushBack,
  pushAllBack,
  pushFoldableBack,

  -- * Conversion
  freeze,
  unsafeFreeze,
  unsafeFreezeVector,
) where

import Control.Monad (forM_)
import Control.Monad.Primitive
import Control.Monad.ST.Strict
import Control.Monad.Trans.State.Strict
import Data.Primitive.Array
import Data.Primitive.Contiguous qualified as C
import Data.Primitive.MutVar
import Data.Vector (Vector)
import Data.Vector qualified as V
import Prelude hiding (read, replicate)

-- | Single-threaded buffer
newtype Buffer s a = Buffer (MutVar s (BufferRep s a))

data BufferRep s a = BufferRep
  { length :: {-# UNPACK #-} !Int
  , buffer :: {-# UNPACK #-} !(MutableArray s a)
  }

{-# INLINE run #-}
run :: Int -> (forall s. Buffer s a -> ST s ()) -> Vector a
run initialCapacity fn = runST do
  buffer <- new initialCapacity
  fn buffer
  unsafeFreezeVector buffer

{-# INLINE runPrim #-}
runPrim :: (PrimMonad m) => Int -> (Buffer (PrimState m) a -> m ()) -> m (Vector a)
runPrim initialCapacity fn = do
  buffer <- new initialCapacity
  fn buffer
  unsafeFreezeVector buffer

-- | Create a buffer with an initial (unfilled) capacity
{-# INLINE new #-}
new :: forall a m. (PrimMonad m) => Int -> m (Buffer (PrimState m) a)
new initialCapacity = do
  buffer <- newArray initialCapacity uninitialisedElement
  Buffer <$> newMutVar (BufferRep 0 buffer)

-- | Create a buffer with n initial elements
{-# INLINE replicate #-}
replicate :: (PrimMonad m) => Int -> a -> m (Buffer (PrimState m) a)
replicate size initialElem = do
  buffer <- newArray size initialElem
  Buffer <$> newMutVar (BufferRep size buffer)

-- | Create a buffer from an array, copying its elements
{-# INLINE thaw #-}
thaw :: (PrimMonad m) => Array a -> m (Buffer (PrimState m) a)
thaw input = do
  buffer <- thawArray input 0 (sizeofArray input)
  Buffer <$> newMutVar (BufferRep (sizeofArray input) buffer)

{-# INLINE read #-}

-- | Read an item from a buffer
read :: (PrimMonad m) => Buffer (PrimState m) a -> Int -> m a
read (Buffer ref) i = do
  BufferRep len buf <- readMutVar ref
  if i < len
    then readArray buf i
    else error ("index: out of range " <> show (i, len))

{-# INLINE write #-}

-- | Write an item to a buffer
--
-- The index given must be within the buffer's existing *length* - not capacity.
write :: (PrimMonad m) => Buffer (PrimState m) a -> Int -> a -> m ()
write (Buffer ref) i a = do
  BufferRep len buf <- readMutVar ref
  if i < len
    then writeArray buf i a
    else error ("index: out of range " <> show (i, len))

{-# INLINE pushBack #-}

-- | Push an element to end of the buffer, returning the new length
pushBack :: (PrimMonad m) => Buffer (PrimState m) a -> a -> m Int
pushBack (Buffer ref) a = do
  BufferRep len dest <- readMutVar ref
  let capacity = sizeofMutableArray dest
  if len < capacity
    then do
      writeArray dest len a
      writeMutVar ref (BufferRep (len + 1) dest)
    else do
      next <- newArray (grow (capacity + 1)) uninitialisedElement
      copyMutableArray next 0 dest 0 (sizeofMutableArray dest)
      writeArray next len a
      writeMutVar ref (BufferRep (len + 1) next)
  pure (len + 1)

{-# INLINE pushFoldableBack #-}

-- | Push multiple elements to the end of the buffer, returning the new length
pushFoldableBack :: (Foldable f, PrimMonad m) => Buffer (PrimState m) a -> f a -> m Int
pushFoldableBack b@(Buffer ref) f
  | null f = size b
pushFoldableBack (Buffer ref) f = do
  BufferRep {length = len, buffer = dest} <- readMutVar ref
  let capacity = sizeofMutableArray dest
  let nextLen = len + Prelude.length f
  if nextLen < capacity
    then do
      -- just write
      writtenLen <- (`execStateT` len) do
        forM_ f \x -> StateT \i -> do
          writeArray dest i x
          pure ((), i + 1)
      -- assert (writtenLen == nextLen)
      writeMutVar ref (BufferRep writtenLen dest)
    else do
      -- grow then write
      dest <- C.resize dest (grow nextLen)
      writtenLen <- (`execStateT` len) do
        forM_ f \x -> StateT \i -> do
          writeArray dest i x
          pure ((), i + 1)
      -- assert (writtenLen == nextLen)
      writeMutVar ref (BufferRep writtenLen dest)
  pure nextLen

{-# INLINE pushAllBack #-}

-- | Push multiple elements to the end of the buffer
pushAllBack :: (PrimMonad m) => Buffer (PrimState m) a -> C.Sliced Array a -> m Int
pushAllBack b@(Buffer ref) arr
  | C.null arr = size b
pushAllBack (Buffer ref) arr = do
  BufferRep len dest <- readMutVar ref
  let capacity = sizeofMutableArray dest
  let nextLen = len + C.size arr
  if nextLen < capacity
    then do
      -- just write
      C.copy dest len arr
      writeMutVar ref (BufferRep nextLen dest)
    else do
      -- grow then write
      dest <- C.resize dest (grow nextLen)
      C.copy dest len arr
      writeMutVar ref (BufferRep nextLen dest)
  pure nextLen

{-# INLINE popBack #-}

-- | Pop the last element off of the buffer
popBack :: (PrimMonad m) => Buffer (PrimState m) a -> m (Maybe (a, Int))
popBack (Buffer ref) = do
  BufferRep len buf <- readMutVar ref
  let i = len - 1
  if i >= 0
    then do
      writeMutVar ref (BufferRep i buf)
      el <- readArray buf i
      pure (Just (el, i))
    else pure Nothing

-- | Freeze the buffer by copying
{-# INLINE freeze #-}
freeze :: (PrimMonad m) => Buffer (PrimState m) a -> m (Array a)
freeze (Buffer ref) = do
  BufferRep len buf <- readMutVar ref
  C.freeze (C.sliceMut buf 0 len)

-- | Freeze the buffer without copying
{-# INLINE unsafeFreeze #-}
unsafeFreeze :: (PrimMonad m) => Buffer (PrimState m) a -> m (C.Sliced C.Array a)
unsafeFreeze (Buffer ref) = do
  BufferRep len buf <- readMutVar ref
  frz <- C.unsafeFreeze buf
  pure (C.slice frz 0 len)

-- | Freeze the buffer into a 'Vector' without copying
{-# INLINE unsafeFreezeVector #-}
unsafeFreezeVector :: (PrimMonad m) => Buffer (PrimState m) a -> m (Vector a)
unsafeFreezeVector (Buffer ref) = do
  BufferRep len buf <- readMutVar ref
  frz <- C.unsafeFreeze buf
  pure (V.unsafeFromArraySlice frz 0 len)

--------------------------------------------------------------------------------

-- | The current length of the buffer
{-# INLINE size #-}
size :: (PrimMonad m) => Buffer (PrimState m) a -> m Int
size (Buffer ref) = do
  BufferRep len _buf <- readMutVar ref
  pure len

-- | The current capacity (as in, remaining free space) of the buffer
{-# INLINE capacity #-}
capacity :: (PrimMonad m) => Buffer (PrimState m) a -> m Int
capacity (Buffer ref) = do
  BufferRep len buf <- readMutVar ref
  let capacity = sizeofMutableArray buf
  return (capacity - len)

--------------------------------------------------------------------------------

{-# INLINE grow #-}
grow :: Int -> Int
grow i = i * 2

{-# NOINLINE uninitialisedElement #-}
uninitialisedElement :: a
uninitialisedElement = error "Buffer: Not initialised"