Unboxed.hs

{-# LANGUAGE Strict #-}

module Data.Vector.Groan.Unboxed (
  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,

  -- * Conversion
  freeze,
  unsafeFreeze,
) where

import Control.Monad (forM_)
import Control.Monad.Primitive
import Control.Monad.ST.Strict
import Control.Monad.Trans.State.Strict
import Data.Monoid (First (..))
import Data.Primitive.MutVar
import Data.Vector qualified as V
import Data.Vector.Unboxed (Unbox)
import Data.Vector.Unboxed qualified as U
import Data.Vector.Unboxed.Mutable qualified as UM
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 :: !(UM.MVector s a)
  }

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

{-# INLINE runPrim #-}
runPrim :: (Unbox a, PrimMonad m) => Int -> a -> (Buffer (PrimState m) a -> m ()) -> m (U.Vector a)
runPrim initialCapacity z fn = do
  buffer <- new initialCapacity z
  fn buffer
  unsafeFreeze buffer

-- | Create a buffer with an initial (unfilled) capacity and a value for
-- uninitialised elements
{-# INLINE new #-}
new :: forall a m. (Unbox a, PrimMonad m) => Int -> a -> m (Buffer (PrimState m) a)
new initialCapacity uninitialisedElement = do
  buffer <- UM.replicate initialCapacity uninitialisedElement
  Buffer <$> newMutVar (BufferRep 0 buffer)

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

-- | Create a buffer from an array, copying its elements
{-# INLINE thaw #-}
thaw :: (Unbox a, PrimMonad m) => U.Vector a -> m (Buffer (PrimState m) a)
thaw input = do
  buffer <- U.thaw input
  Buffer <$> newMutVar (BufferRep (U.length input) buffer)

{-# INLINE read #-}

-- | Read an item from a buffer
read :: (Unbox a, PrimMonad m) => Buffer (PrimState m) a -> Int -> m a
read (Buffer ref) i = do
  BufferRep len buf <- readMutVar ref
  if i < len
    then UM.read 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 :: (Unbox a, PrimMonad m) => Buffer (PrimState m) a -> Int -> a -> m ()
write (Buffer ref) i a = do
  BufferRep len buf <- readMutVar ref
  if i < len
    then UM.write 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 :: (Unbox a, PrimMonad m) => Buffer (PrimState m) a -> a -> m Int
pushBack (Buffer ref) a = do
  BufferRep len dest <- readMutVar ref
  let capacity = UM.length dest
  if len < capacity
    then do
      UM.write dest len a
      writeMutVar ref (BufferRep (len + 1) dest)
    else do
      next <- UM.replicate (grow (capacity + 1)) a
      UM.copy (UM.slice 0 len next) dest
      UM.write next len a
      writeMutVar ref (BufferRep (len + 1) next)
  pure (len + 1)

{-# INLINE popBack #-}

-- | Pop the last element off of the buffer
popBack :: (Unbox a, 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 <- UM.read buf i
      pure (Just (el, i))
    else pure Nothing

-- | Freeze the buffer by copying
{-# INLINE freeze #-}
freeze :: (Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (U.Vector a)
freeze (Buffer ref) = do
  BufferRep len buf <- readMutVar ref
  U.freeze (UM.slice 0 len buf)

-- | Freeze the buffer into a 'Vector' without copying
{-# INLINE unsafeFreeze #-}
unsafeFreeze :: (Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m (U.Vector a)
unsafeFreeze (Buffer ref) = do
  BufferRep len buf <- readMutVar ref
  U.unsafeFreeze (UM.slice 0 len buf)

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

-- | 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 :: (Unbox a, PrimMonad m) => Buffer (PrimState m) a -> m Int
capacity (Buffer ref) = do
  BufferRep len buf <- readMutVar ref
  let capacity = UM.length buf
  pure (capacity - len)

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

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