Vec.hs

-- Adapted from `grow-vector` as the package seems unmaintained
-- Module defines mutable vector that can grow in size automatically when an user
-- adds new elements at the end of vector.
--
-- We reallocate vector with 1.5x length to get amortized append.
{- HLINT ignore "Use camelCase" -}
module Mischief.ECS.Vec
  ( Vec (..),
    IOVec,

    -- * Quering info about vector
    length,
    null,
    capacity,

    -- * Creation
    new,
    newSized,

    -- * Quering subvectors
    slice,

    -- * Converting to immutable
    thaw,
    freeze,
    toList,

    -- * Capacity maninuplation
    ensure,
    ensureAppend,

    -- * Accessing individual elements
    read,
    write,
    unsafeRead,
    unsafeWrite,

    -- * Appending to vector
    pushBack,
    unsafePushBack,

    -- * modify an element
    modify,
    modify_,
    modifyM,
    modifyM_,

    -- * general utilities
    swap,
    tap,
    shrink,

    -- * O(1) amortized swap backed operations
    takeSwap,
    removeSwap,

    -- * cloning, very naive approach
    clone,
  )
where

import Control.Monad
import Control.Monad.Primitive
import Data.Foldable (for_)
import Data.Primitive.MutVar
import Data.Vector (Vector)
import Data.Vector qualified as Vector
import Data.Vector.Mutable (MVector)
import Data.Vector.Mutable qualified as MVector
import GHC.Generics
import GHC.Stack (HasCallStack)
import Prelude hiding (length, null, read)

-- | Normal rust-like vector with buffer (the MVector) + len (the len field) + cap (the `buffer` length)
data Vec s a = Vec
  { buffer :: !(MutVar s (MVector s a)),
    len :: !(MutVar s Int)
  }
  deriving (Generic)

type IOVec a = Vec RealWorld a

-- | Return current capacity of the vector (amount of elements that it can fit without realloc)
capacity :: (PrimMonad m) => Vec (PrimState m) a -> m Int
capacity !v =
  MVector.length <$> readMutVar v.buffer
{-# INLINE capacity #-}

-- | Return current amount of elements in the vector
length :: (PrimMonad m) => Vec (PrimState m) a -> m Int
length !v = readMutVar v.len
{-# INLINE length #-}

-- | Return 'True' if there is no elements inside the vector
null :: (PrimMonad m) => Vec (PrimState m) a -> m Bool
null !v = (== 0) <$> length v

-- | Allocation of new growable vector with given capacity.
new :: (PrimMonad m) => Int -> m (Vec (PrimState m) a)
new = newSized 0
{-# INLINE new #-}

-- | Allocation of new growable vector with given filled size and capacity.
-- Elements is not initialized. Capacity must be greater than filled size.
newSized :: (PrimMonad m) => Int -> Int -> m (Vec (PrimState m) a)
newSized !n !cap = do
  !buffer <- MVector.new cap
  !buffer <- newMutVar buffer
  !len <- newMutVar n
  pure Vec {buffer, len}
{-# INLINEABLE newSized #-}

-- | Yield a part of mutable vector without copying it. The vector must contain at least i+n elements.
slice ::
  (PrimMonad m) =>
  -- | i starting index
  Int ->
  -- | n number of elements
  Int ->
  Vec (PrimState m) a ->
  m (Vec (PrimState m) a)
slice !i !n !v = do
  !len <- newMutVar n
  !mv <- readMutVar v.buffer
  !buffer <- newMutVar $! MVector.slice i n mv
  pure $! Vec {len, buffer}
{-# INLINEABLE slice #-}

-- | Convert immutable vector to grow mutable version. Doesn't allocate additonal memory for appending,
-- use 'ensure' to add capacity to the vector.
thaw ::
  (PrimMonad m) =>
  Vector a ->
  m (Vec (PrimState m) a)
thaw !u = do
  !buffer <- newMutVar =<< Vector.thaw u
  !len <- newMutVar $! Vector.length u
  pure Vec {buffer, len}
{-# INLINEABLE thaw #-}

-- | Freezing growable vector. It will contain only actual elements of the vector not including capacity
-- space, but you should call 'U.force' on resulting vector to not hold the allocated capacity of original
-- vector in memory.
freeze ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  m (Vector a)
freeze !v = do
  !n <- length v
  !mv <- readMutVar v.buffer
  Vector.freeze $ MVector.take n mv
{-# INLINEABLE freeze #-}

toList ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  m [a]
toList !vec = Vector.toList <$> freeze vec

ensure_not_oob ::
  (HasCallStack, PrimMonad m) =>
  String ->
  -- | the name of the function
  Int ->
  -- | The element we want to check
  Vec (PrimState m) a ->
  -- | The length of the vec
  m ()
ensure_not_oob !fname !i !vec = do
  !len <- length vec
  when (i < 0 || i >= len) $ do
    error $ mconcat [fname, ": index ", show i, " is out bounds ", show len]

-- | Ensure that grow vector has at least given capacity possibly with reallocation.
ensure ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  Int ->
  m ()
ensure !v !cap = do
  !current_cap <- capacity v
  unless (current_cap >= cap) $ do
    !buffer <- readMutVar v.buffer
    !grown <- MVector.grow buffer (cap - current_cap)
    writeMutVar v.buffer grown
{-# INLINEABLE ensure #-}

-- | Ensure that grow vector has enough space for additonal n elements.
-- We grow vector by 1.5 factor or by required elements count * 1.5.
ensureAppend ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  -- | Additional n elements
  Int ->
  m ()
ensureAppend !vec !i = do
  !len <- length vec
  !buf <- readMutVar vec.buffer
  let !cap = MVector.length buf
  unless (cap >= len + i) $ do
    -- ugly as shit code would like to fix
    let !newCap = ceiling $ max (growFactor * fromIntegral cap) (fromIntegral cap + growFactor * fromIntegral (len + i - cap))
    !new_buf <- MVector.grow buf (newCap - cap)
    writeMutVar vec.buffer new_buf
  where
    growFactor :: Double
    growFactor = 1.5
{-# INLINEABLE ensureAppend #-}

-- | Read element from vector at given index.
read ::
  (HasCallStack, PrimMonad m) =>
  Vec (PrimState m) a ->
  -- | Index of element. Must be in [0 .. length) range
  Int ->
  m a
read !vec !i = do
  ensure_not_oob "Vec.read" i vec
  !buf <- readMutVar vec.buffer
  MVector.unsafeRead buf i
{-# INLINEABLE read #-}

-- | Read element from vector at given index, without checking whether the index is inbounds
unsafeRead ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  -- | Index of element. Must be in [0 .. length) range
  Int ->
  m a
unsafeRead !vec !i = do
  !buf <- readMutVar vec.buffer
  MVector.unsafeRead buf i
{-# INLINEABLE unsafeRead #-}

-- | Write down element in the vector at given index.
write ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  -- | Index of element. Must be in [0 .. length) range
  Int ->
  a ->
  m ()
write !vec !i !value = do
  ensure_not_oob "Vec.write" i vec
  !buf <- readMutVar vec.buffer
  MVector.unsafeWrite buf i value
{-# INLINEABLE write #-}

-- | Write down element in the vector at given index.
unsafeWrite ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  -- | Index of element. Must be in [0 .. length) range
  Int ->
  a ->
  m ()
unsafeWrite !vec !i !value = do
  !buf <- readMutVar vec.buffer
  MVector.unsafeWrite buf i value
{-# INLINEABLE unsafeWrite #-}

modify ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  Int ->
  (a -> a) ->
  m a
modify !vec !i !f = do
  ensure_not_oob "Vec.modify" i vec
  !old_val <- read vec i
  write vec i (f old_val)
  pure old_val
{-# INLINEABLE modify #-}

modify_ ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  Int ->
  (a -> a) ->
  m ()
modify_ !vec !i !f = do
  ensure_not_oob "Vec.modify_" i vec
  !old_val <- read vec i
  write vec i (f old_val)
{-# INLINEABLE modify_ #-}

modifyM ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  Int ->
  (a -> m a) ->
  m a
modifyM !vec !i !f = do
  ensure_not_oob "Vec.modifyM" i vec
  !old_val <- read vec i
  write vec i =<< f old_val
  pure old_val
{-# INLINEABLE modifyM #-}

modifyM_ ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  Int ->
  (a -> m a) ->
  m ()
modifyM_ !vec !i !f = do
  ensure_not_oob "Vec.modifyM_" i vec
  !old_val <- read vec i
  write vec i =<< f old_val
{-# INLINEABLE modifyM_ #-}

tap ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  Int ->
  (a -> m ()) ->
  m ()
tap !vec !i !act = do
  ensure_not_oob "Vec.tap" i vec
  act =<< read vec i
{-# INLINE tap #-}

-- | O(1) amortized appending to vector
pushBack ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  a ->
  m ()
pushBack !vec !value = do
  ensureAppend vec 1
  unsafePushBack vec value
{-# INLINEABLE pushBack #-}

-- | O(1) amortized appending to vector. Doesn't reallocate vector, so
-- there must by capacity - length >= 1.
unsafePushBack ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  a ->
  m ()
unsafePushBack !vec !a = do
  !len <- length vec
  !buf <- readMutVar vec.buffer
  MVector.write buf len a
  writeMutVar vec.len (len + 1)
{-# INLINEABLE unsafePushBack #-}

swap ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  Int ->
  -- | index of element #1 to swap
  Int ->
  -- | index of element #2 to swap
  m ()
swap !vec !i !j = do
  when (i /= j) $ do
    ensure_not_oob "Vec.swap" i vec
    ensure_not_oob "Vec.swap" j vec
    !buffer <- readMutVar vec.buffer
    MVector.swap buffer i j

shrink ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  Int ->
  m ()
shrink !vec !amount = do
  !old_len <- length vec
  writeMutVar vec.len (old_len - amount)

takeSwap ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  Int ->
  m a
takeSwap !vec !i = do
  !value <- read vec i
  removeSwap vec i
  pure value

removeSwap ::
  (PrimMonad m) =>
  Vec (PrimState m) a ->
  -- | index of element to remove
  Int ->
  m ()
removeSwap !vec !i = do
  !len <- length vec
  swap vec i (len - 1)
  shrink vec 1

clone :: (PrimMonad m) => Vec (PrimState m) a -> m (Vec (PrimState m) a)
clone !vec = do
  !list <- toList vec
  !len <- length vec
  !newVec <- new len
  for_ list $ pushBack newVec
  pure newVec