-- |
-- Description : A generic growable mutable vector with O(1) amortized append.
-- Copyright   : (c) Jeremy Nuttall, 2025
-- License     : BSD-3-Clause
-- Maintainer  : jeremy@jeremy-nuttall.com
-- Stability   : experimental
-- Portability : GHC
--
-- __WARNING:__ The functions in this module are generally unchecked and unsafe. Be careful to understand
-- and maintain invariants if using them. Misuse may result in undefined behavior.
--
-- Internal modules can change without warning between minor versions.
module Data.SparseSet.Generic.Mutable.Internal.GrowVec (
  GrowVec,
  withCapacity,
  new,
  length,
  capacity,
  snoc,
  readMaybe,
  unsafeRead,
  maximum,
  unsafeWrite,
  unsafeSwapRemove,
  mapM_,
  cleared,
  compact,
  freeze,
  unsafeFreeze,
)
where

import Control.DeepSeq (NFData)
import Control.Monad.Primitive
import Data.Typeable (Typeable)
import Data.Vector.Generic qualified as VG
import Data.Vector.Generic.Mutable qualified as VGM
import GHC.Generics (Generic)
import Prelude hiding (length, mapM_, maximum)

data GrowVec v s a = GrowVec {-# UNPACK #-} !Int (v s a)
  deriving (Int -> GrowVec v s a -> ShowS
[GrowVec v s a] -> ShowS
GrowVec v s a -> String
(Int -> GrowVec v s a -> ShowS)
-> (GrowVec v s a -> String)
-> ([GrowVec v s a] -> ShowS)
-> Show (GrowVec v s a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall (v :: * -> * -> *) s a.
Show (v s a) =>
Int -> GrowVec v s a -> ShowS
forall (v :: * -> * -> *) s a.
Show (v s a) =>
[GrowVec v s a] -> ShowS
forall (v :: * -> * -> *) s a.
Show (v s a) =>
GrowVec v s a -> String
$cshowsPrec :: forall (v :: * -> * -> *) s a.
Show (v s a) =>
Int -> GrowVec v s a -> ShowS
showsPrec :: Int -> GrowVec v s a -> ShowS
$cshow :: forall (v :: * -> * -> *) s a.
Show (v s a) =>
GrowVec v s a -> String
show :: GrowVec v s a -> String
$cshowList :: forall (v :: * -> * -> *) s a.
Show (v s a) =>
[GrowVec v s a] -> ShowS
showList :: [GrowVec v s a] -> ShowS
Show, (forall x. GrowVec v s a -> Rep (GrowVec v s a) x)
-> (forall x. Rep (GrowVec v s a) x -> GrowVec v s a)
-> Generic (GrowVec v s a)
forall x. Rep (GrowVec v s a) x -> GrowVec v s a
forall x. GrowVec v s a -> Rep (GrowVec v s a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall (v :: * -> * -> *) s a x.
Rep (GrowVec v s a) x -> GrowVec v s a
forall (v :: * -> * -> *) s a x.
GrowVec v s a -> Rep (GrowVec v s a) x
$cfrom :: forall (v :: * -> * -> *) s a x.
GrowVec v s a -> Rep (GrowVec v s a) x
from :: forall x. GrowVec v s a -> Rep (GrowVec v s a) x
$cto :: forall (v :: * -> * -> *) s a x.
Rep (GrowVec v s a) x -> GrowVec v s a
to :: forall x. Rep (GrowVec v s a) x -> GrowVec v s a
Generic, Typeable)

instance (NFData (v s a)) => NFData (GrowVec v s a)

-- | Create a new, empty vector with the given capacity
--
-- @since 0.1.0.0
withCapacity :: forall a v m. (PrimMonad m, VGM.MVector v a) => Int -> m (GrowVec v (PrimState m) a)
withCapacity :: forall a (v :: * -> * -> *) (m :: * -> *).
(PrimMonad m, MVector v a) =>
Int -> m (GrowVec v (PrimState m) a)
withCapacity Int
c = Int -> v (PrimState m) a -> GrowVec v (PrimState m) a
forall (v :: * -> * -> *) s a. Int -> v s a -> GrowVec v s a
GrowVec Int
0 (v (PrimState m) a -> GrowVec v (PrimState m) a)
-> m (v (PrimState m) a) -> m (GrowVec v (PrimState m) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m (v (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
VGM.new (Int -> Int
withMinCapacity Int
c)
{-# INLINE withCapacity #-}

-- | Create a new, empty vector with a default capcity
--
-- @since 0.1.0.0
new :: forall a v m. (PrimMonad m, VGM.MVector v a) => m (GrowVec v (PrimState m) a)
new :: forall a (v :: * -> * -> *) (m :: * -> *).
(PrimMonad m, MVector v a) =>
m (GrowVec v (PrimState m) a)
new = Int -> v (PrimState m) a -> GrowVec v (PrimState m) a
forall (v :: * -> * -> *) s a. Int -> v s a -> GrowVec v s a
GrowVec Int
0 (v (PrimState m) a -> GrowVec v (PrimState m) a)
-> m (v (PrimState m) a) -> m (GrowVec v (PrimState m) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m (v (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
VGM.new Int
16
{-# INLINE new #-}

-- | O(1) The logical length of the vector.
length :: forall a v s. GrowVec v s a -> Int
length :: forall a (v :: * -> * -> *) s. GrowVec v s a -> Int
length (GrowVec Int
l v s a
_) = Int
l
{-# INLINE length #-}

-- | O(1) The capacity of the vector.
capacity :: (VGM.MVector v a) => GrowVec v s a -> Int
capacity :: forall (v :: * -> * -> *) a s. MVector v a => GrowVec v s a -> Int
capacity (GrowVec Int
_ v s a
v) = v s a -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
VGM.length v s a
v
{-# INLINE capacity #-}

-- | Calculate the additional number of elements given the current length.
--
-- __INVARIANT__: length must be >= 2
--
-- @since 0.1.0.0
growthFactor :: Int -> Int
growthFactor :: Int -> Int
growthFactor Int
l = (Int
l Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
2) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
3
{-# INLINE growthFactor #-}

-- | O(1) amortized. Append to the vector, reallocating and copying if necessary.
--
-- Since this can't be done in-place, you must use the resulting vector in further computations.
--
-- @since 0.1.0.0
snoc
  :: (VGM.MVector v a, PrimMonad m) => GrowVec v (PrimState m) a -> a -> m (GrowVec v (PrimState m) a)
snoc :: forall (v :: * -> * -> *) a (m :: * -> *).
(MVector v a, PrimMonad m) =>
GrowVec v (PrimState m) a -> a -> m (GrowVec v (PrimState m) a)
snoc GrowVec v (PrimState m) a
gv a
a = do
  GrowVec v (PrimState m) a
gv' <- GrowVec v (PrimState m) a -> m (GrowVec v (PrimState m) a)
forall {f :: * -> *} {v :: * -> * -> *} {a}.
(PrimMonad f, MVector v a) =>
GrowVec v (PrimState f) a -> f (GrowVec v (PrimState f) a)
grow GrowVec v (PrimState m) a
gv
  GrowVec v (PrimState m) a -> Int -> a -> m ()
forall a (m :: * -> *) (v :: * -> * -> *).
(PrimMonad m, MVector v a) =>
GrowVec v (PrimState m) a -> Int -> a -> m ()
unsafeWrite GrowVec v (PrimState m) a
gv' (GrowVec v (PrimState m) a -> Int
forall a (v :: * -> * -> *) s. GrowVec v s a -> Int
length GrowVec v (PrimState m) a
gv) a
a
  GrowVec v (PrimState m) a -> m (GrowVec v (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure GrowVec v (PrimState m) a
gv'
 where
  grow :: GrowVec v (PrimState f) a -> f (GrowVec v (PrimState f) a)
grow (GrowVec Int
l v (PrimState f) a
v)
    | GrowVec v (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => GrowVec v s a -> Int
capacity GrowVec v (PrimState m) a
gv Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l = Int -> v (PrimState f) a -> GrowVec v (PrimState f) a
forall (v :: * -> * -> *) s a. Int -> v s a -> GrowVec v s a
GrowVec (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (v (PrimState f) a -> GrowVec v (PrimState f) a)
-> f (v (PrimState f) a) -> f (GrowVec v (PrimState f) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> v (PrimState f) a -> Int -> f (v (PrimState f) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m (v (PrimState m) a)
VGM.grow v (PrimState f) a
v (Int -> Int
growthFactor Int
l)
    | Bool
otherwise = GrowVec v (PrimState f) a -> f (GrowVec v (PrimState f) a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (GrowVec v (PrimState f) a -> f (GrowVec v (PrimState f) a))
-> GrowVec v (PrimState f) a -> f (GrowVec v (PrimState f) a)
forall a b. (a -> b) -> a -> b
$ Int -> v (PrimState f) a -> GrowVec v (PrimState f) a
forall (v :: * -> * -> *) s a. Int -> v s a -> GrowVec v s a
GrowVec (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) v (PrimState f) a
v
{-# INLINE snoc #-}

readMaybe
  :: forall a m v. (PrimMonad m, VGM.MVector v a) => GrowVec v (PrimState m) a -> Int -> m (Maybe a)
readMaybe :: forall a (m :: * -> *) (v :: * -> * -> *).
(PrimMonad m, MVector v a) =>
GrowVec v (PrimState m) a -> Int -> m (Maybe a)
readMaybe (GrowVec Int
l v (PrimState m) a
v) Int
i
  | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
l = Maybe a -> m (Maybe a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
  | Bool
otherwise = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> m a -> m (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead v (PrimState m) a
v Int
i
{-# INLINE readMaybe #-}

-- | O(1) Read from a position in the vector. This position must be less than the length of the vector. This is not checked.
--
-- @since 0.1.0.0
unsafeRead
  :: forall a m v. (PrimMonad m, VGM.MVector v a) => GrowVec v (PrimState m) a -> Int -> m a
unsafeRead :: forall a (m :: * -> *) (v :: * -> * -> *).
(PrimMonad m, MVector v a) =>
GrowVec v (PrimState m) a -> Int -> m a
unsafeRead (GrowVec Int
_ v (PrimState m) a
v) = v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead v (PrimState m) a
v
{-# INLINE unsafeRead #-}

-- | O(n) The maximum value in the vector. Useful for compaction.
--
-- @since 0.1.0.0
maximum
  :: (PrimMonad m, VGM.MVector v a, Ord a, Bounded a) => GrowVec v (PrimState m) a -> m (Maybe a)
maximum :: forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a, Ord a, Bounded a) =>
GrowVec v (PrimState m) a -> m (Maybe a)
maximum (GrowVec Int
l v (PrimState m) a
v)
  | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = Maybe a -> m (Maybe a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
  | Bool
otherwise = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> m a -> m (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> a -> a) -> a -> v (PrimState m) a -> m a
forall (m :: * -> *) (v :: * -> * -> *) a b.
(PrimMonad m, MVector v a) =>
(b -> a -> b) -> b -> v (PrimState m) a -> m b
VGM.foldl' a -> a -> a
forall a. Ord a => a -> a -> a
max a
forall a. Bounded a => a
minBound (Int -> Int -> v (PrimState m) a -> v (PrimState m) a
forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
VGM.unsafeSlice Int
0 Int
l v (PrimState m) a
v)
{-# INLINE maximum #-}

-- | O(1) Write to a position in the vector. This position must be less than the length of the vector. This is not checked.
--
-- @since 0.1.0.0
unsafeWrite
  :: forall a m v. (PrimMonad m, VGM.MVector v a) => GrowVec v (PrimState m) a -> Int -> a -> m ()
unsafeWrite :: forall a (m :: * -> *) (v :: * -> * -> *).
(PrimMonad m, MVector v a) =>
GrowVec v (PrimState m) a -> Int -> a -> m ()
unsafeWrite (GrowVec Int
_ v (PrimState m) a
v) = v (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite v (PrimState m) a
v
{-# INLINE unsafeWrite #-}

-- | O(1) Swap-and-pop an element in the vector
--
-- @since 0.1.0.0
unsafeSwapRemove
  :: forall a m v
   . (PrimMonad m, VGM.MVector v a)
  => GrowVec v (PrimState m) a
  -> Int
  -> m (a, GrowVec v (PrimState m) a)
unsafeSwapRemove :: forall a (m :: * -> *) (v :: * -> * -> *).
(PrimMonad m, MVector v a) =>
GrowVec v (PrimState m) a
-> Int -> m (a, GrowVec v (PrimState m) a)
unsafeSwapRemove (GrowVec Int
l v (PrimState m) a
v) Int
i = do
  a
old <- v (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.read v (PrimState m) a
v Int
i
  v (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> Int -> m ()
VGM.swap v (PrimState m) a
v Int
i (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  (a, GrowVec v (PrimState m) a) -> m (a, GrowVec v (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
old, Int -> v (PrimState m) a -> GrowVec v (PrimState m) a
forall (v :: * -> * -> *) s a. Int -> v s a -> GrowVec v s a
GrowVec (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) v (PrimState m) a
v)
{-# INLINE unsafeSwapRemove #-}

mapM_ :: (PrimMonad m, VGM.MVector v a) => (a -> m b) -> GrowVec v (PrimState m) a -> m ()
mapM_ :: forall (m :: * -> *) (v :: * -> * -> *) a b.
(PrimMonad m, MVector v a) =>
(a -> m b) -> GrowVec v (PrimState m) a -> m ()
mapM_ a -> m b
f (GrowVec Int
l v (PrimState m) a
v) = (a -> m b) -> v (PrimState m) a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a b.
(PrimMonad m, MVector v a) =>
(a -> m b) -> v (PrimState m) a -> m ()
VGM.mapM_ a -> m b
f (Int -> Int -> v (PrimState m) a -> v (PrimState m) a
forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
VGM.unsafeSlice Int
0 Int
l v (PrimState m) a
v)
{-# INLINE mapM_ #-}

-- | O(1) Create a new, empty vector by setting logical length to 0. This does not change the
-- underlying vector in any way.
--
-- @since 0.1.0.0
cleared :: forall a s v. GrowVec v s a -> GrowVec v s a
cleared :: forall a s (v :: * -> * -> *). GrowVec v s a -> GrowVec v s a
cleared (GrowVec Int
_ v s a
v) = Int -> v s a -> GrowVec v s a
forall (v :: * -> * -> *) s a. Int -> v s a -> GrowVec v s a
GrowVec Int
0 v s a
v
{-# INLINE cleared #-}

-- | O(n) Shrink the vector so that its capacity matches its current length
--
-- @since 0.1.0.0
compact
  :: (VGM.MVector v a, PrimMonad m) => GrowVec v (PrimState m) a -> m (GrowVec v (PrimState m) a)
compact :: forall (v :: * -> * -> *) a (m :: * -> *).
(MVector v a, PrimMonad m) =>
GrowVec v (PrimState m) a -> m (GrowVec v (PrimState m) a)
compact gv :: GrowVec v (PrimState m) a
gv@(GrowVec Int
l v (PrimState m) a
v)
  | GrowVec v (PrimState m) a -> Int
forall (v :: * -> * -> *) a s. MVector v a => GrowVec v s a -> Int
capacity GrowVec v (PrimState m) a
gv Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
l = GrowVec v (PrimState m) a -> m (GrowVec v (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure GrowVec v (PrimState m) a
gv
  | Bool
otherwise = do
      let l' :: Int
l' = Int -> Int
withMinCapacity Int
l
      v (PrimState m) a
v' <- v (PrimState m) a -> m (v (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> m (v (PrimState m) a)
VGM.clone (Int -> Int -> v (PrimState m) a -> v (PrimState m) a
forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
VGM.unsafeSlice Int
0 Int
l' v (PrimState m) a
v)
      GrowVec v (PrimState m) a -> m (GrowVec v (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> v (PrimState m) a -> GrowVec v (PrimState m) a
forall (v :: * -> * -> *) s a. Int -> v s a -> GrowVec v s a
GrowVec Int
l v (PrimState m) a
v')

freeze :: (PrimMonad m, VG.Vector v a) => GrowVec (VG.Mutable v) (PrimState m) a -> m (v a)
freeze :: forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
GrowVec (Mutable v) (PrimState m) a -> m (v a)
freeze (GrowVec Int
l Mutable v (PrimState m) a
v) = Mutable v (PrimState m) a -> m (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
VG.freeze (Int
-> Int -> Mutable v (PrimState m) a -> Mutable v (PrimState m) a
forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
VGM.unsafeSlice Int
0 Int
l Mutable v (PrimState m) a
v)
{-# INLINE freeze #-}

unsafeFreeze :: (PrimMonad m, VG.Vector v a) => GrowVec (VG.Mutable v) (PrimState m) a -> m (v a)
unsafeFreeze :: forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
GrowVec (Mutable v) (PrimState m) a -> m (v a)
unsafeFreeze (GrowVec Int
l Mutable v (PrimState m) a
v) = Mutable v (PrimState m) a -> m (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
VG.unsafeFreeze (Int
-> Int -> Mutable v (PrimState m) a -> Mutable v (PrimState m) a
forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
VGM.unsafeSlice Int
0 Int
l Mutable v (PrimState m) a
v)
{-# INLINE unsafeFreeze #-}

--------------------------------------------------------------------------------
-- Utilities
--------------------------------------------------------------------------------
withMinCapacity :: Int -> Int
withMinCapacity :: Int -> Int
withMinCapacity Int
c = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
c Int
4
{-# INLINE withMinCapacity #-}