List.hs

module Data.Vector.Groan.List where

import Control.Monad.Primitive
import Data.Foldable
import Data.IORef
import Data.Primitive.Array
import Data.Primitive.Contiguous qualified as C
import Data.Sequence (Seq (..))
import Data.Sequence qualified as Seq
import Data.Vector.Groan qualified as B
import Data.Vector.Groan.Unboxed qualified as U
import Prelude

newtype Buffer a = Buffer (IORef (Seq a))

type Capacity = Int
type Length = Int
type Index = Int

new :: Capacity -> IO (Buffer a)
new _ = Buffer <$> newIORef Seq.empty

replicate :: Length -> a -> IO (Buffer a)
replicate n a = Buffer <$> newIORef (Seq.replicate n a)

thaw :: Array a -> IO (Buffer a)
thaw arr = Buffer <$> newIORef (Seq.fromList (toList arr))

read :: Buffer a -> Index -> IO a
read (Buffer ref) i = do
  s <- readIORef ref
  case Seq.lookup i s of
    Nothing -> error ("index: out of range " <> show (i, Seq.length s))
    Just a -> pure a

write :: Buffer a -> Index -> a -> IO ()
write (Buffer ref) i a = do
  s <- readIORef ref
  if i < Seq.length s
    then writeIORef ref (Seq.update i a s)
    else error ("index: out of range " <> show (i, Seq.length s))

pushBack
  :: Buffer a
  -> a
  -> IO Length
  -- ^ New length
pushBack (Buffer ref) a = do
  modifyIORef' ref (:|> a)
  Seq.length <$> readIORef ref

pushFoldableBack
  :: (Foldable f)
  => Buffer a
  -> f a
  -> IO Length
  -- ^ New length
pushFoldableBack (Buffer ref) f = do
  modifyIORef' ref (\s -> foldl' (:|>) s f)
  Seq.length <$> readIORef ref

popBack :: Buffer a -> IO (Maybe (a, Length))
popBack (Buffer ref) =
  atomicModifyIORef ref \case
    xs :|> x -> (xs, Just (x, Seq.length xs))
    Empty -> (Empty, Nothing)

freeze :: Buffer a -> IO (Array a)
freeze (Buffer ref) = do
  s <- readIORef ref
  pure (C.fromListN (Seq.length s) (toList s))

size :: Buffer a -> IO Length
size (Buffer ref) = Seq.length <$> readIORef ref

capacity :: Buffer a -> IO Capacity
capacity = size