{-# LANGUAGE CPP #-}
{-# LANGUAGE PatternSynonyms #-}

-- |
-- Description : Mutable sparse arrays, suitable for sparse set implementation.
-- 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.MutableSparseArray (
  MutableSparseArray,
  withCapacity,
  new,
  contains,
  lookup,
  unsafeInsert,
  delete,
  unsafeDelete,
  clear,
  unsafeCompactTo,
  freeze,
  unsafeFreeze,
)
where

import Control.DeepSeq (NFData)
import Control.Monad (when)
import Control.Monad.Primitive
import Data.Maybe (isJust)
import Data.Typeable (Typeable)
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Primitive qualified as VP
import Data.Vector.Primitive.Mutable qualified as VPM
import GHC.Generics (Generic)
import Prelude hiding (lookup, maximum)

pattern ABSURD :: Int
pattern $mABSURD :: forall {r}. Int -> ((# #) -> r) -> ((# #) -> r) -> r
$bABSURD :: Int
ABSURD = -1

-- | Mutable sparse integer array parameterized by its state token.
--
-- @since 0.1.0.0
newtype MutableSparseArray s = MutableSparseArray
  {forall s. MutableSparseArray s -> MVector s Int
getSparseArray :: VPM.MVector s Int}
  deriving newtype (MutableSparseArray s -> ()
(MutableSparseArray s -> ()) -> NFData (MutableSparseArray s)
forall s. MutableSparseArray s -> ()
forall a. (a -> ()) -> NFData a
$crnf :: forall s. MutableSparseArray s -> ()
rnf :: MutableSparseArray s -> ()
NFData)
  deriving stock ((forall x. MutableSparseArray s -> Rep (MutableSparseArray s) x)
-> (forall x. Rep (MutableSparseArray s) x -> MutableSparseArray s)
-> Generic (MutableSparseArray s)
forall x. Rep (MutableSparseArray s) x -> MutableSparseArray s
forall x. MutableSparseArray s -> Rep (MutableSparseArray s) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall s x. Rep (MutableSparseArray s) x -> MutableSparseArray s
forall s x. MutableSparseArray s -> Rep (MutableSparseArray s) x
$cfrom :: forall s x. MutableSparseArray s -> Rep (MutableSparseArray s) x
from :: forall x. MutableSparseArray s -> Rep (MutableSparseArray s) x
$cto :: forall s x. Rep (MutableSparseArray s) x -> MutableSparseArray s
to :: forall x. Rep (MutableSparseArray s) x -> MutableSparseArray s
Generic, Typeable)

-- | Create a new, empty array from a given capacity.
--
-- @since 0.1.0.0
withCapacity :: (PrimMonad m) => Int -> m (MutableSparseArray (PrimState m))
withCapacity :: forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableSparseArray (PrimState m))
withCapacity Int
rc = ST (PrimState m) (MutableSparseArray (PrimState m))
-> m (MutableSparseArray (PrimState m))
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim do
  let c :: Int
c = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
rc Int
4
  MVector (PrimState m) Int
arr <- Int
-> ST (PrimState m) (MVector (PrimState (ST (PrimState m))) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MVector (PrimState m) a)
VPM.new Int
c
  Bool -> ST (PrimState m) () -> ST (PrimState m) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (ST (PrimState m) () -> ST (PrimState m) ())
-> ST (PrimState m) () -> ST (PrimState m) ()
forall a b. (a -> b) -> a -> b
$ Int
-> Int
-> MVector (PrimState (ST (PrimState m))) Int
-> ST (PrimState m) ()
forall (m :: * -> *) (v :: * -> * -> *).
(PrimMonad m, MVector v Int) =>
Int -> Int -> v (PrimState m) Int -> m ()
fillArray Int
0 Int
c MVector (PrimState m) Int
MVector (PrimState (ST (PrimState m))) Int
arr
  MutableSparseArray (PrimState m)
-> ST (PrimState m) (MutableSparseArray (PrimState m))
forall a. a -> ST (PrimState m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MutableSparseArray (PrimState m)
 -> ST (PrimState m) (MutableSparseArray (PrimState m)))
-> MutableSparseArray (PrimState m)
-> ST (PrimState m) (MutableSparseArray (PrimState m))
forall a b. (a -> b) -> a -> b
$ MVector (PrimState m) Int -> MutableSparseArray (PrimState m)
forall s. MVector s Int -> MutableSparseArray s
MutableSparseArray MVector (PrimState m) Int
arr
{-# INLINE withCapacity #-}

-- | Create a new, empty array.
--
-- @since 0.1.0.0
new :: (PrimMonad m) => m (MutableSparseArray (PrimState m))
new :: forall (m :: * -> *).
PrimMonad m =>
m (MutableSparseArray (PrimState m))
new = Int -> m (MutableSparseArray (PrimState m))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableSparseArray (PrimState m))
withCapacity Int
32
{-# INLINE new #-}

contains :: (PrimMonad m) => MutableSparseArray (PrimState m) -> Int -> m Bool
contains :: forall (m :: * -> *).
PrimMonad m =>
MutableSparseArray (PrimState m) -> Int -> m Bool
contains MutableSparseArray (PrimState m)
arr Int
i = Maybe Int -> Bool
forall a. Maybe a -> Bool
isJust (Maybe Int -> Bool) -> m (Maybe Int) -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutableSparseArray (PrimState m) -> Int -> m (Maybe Int)
forall (m :: * -> *).
PrimMonad m =>
MutableSparseArray (PrimState m) -> Int -> m (Maybe Int)
lookup MutableSparseArray (PrimState m)
arr Int
i
{-# INLINE contains #-}

lookup :: (PrimMonad m) => MutableSparseArray (PrimState m) -> Int -> m (Maybe Int)
#if MIN_VERSION_vector(0,13,0)
lookup :: forall (m :: * -> *).
PrimMonad m =>
MutableSparseArray (PrimState m) -> Int -> m (Maybe Int)
lookup (MutableSparseArray MVector (PrimState m) Int
arr) Int
i = (Maybe Int -> (Int -> Maybe Int) -> Maybe Int
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> Maybe Int
msaReprToMaybe) (Maybe Int -> Maybe Int) -> m (Maybe Int) -> m (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) Int -> Int -> m (Maybe Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MVector (PrimState m) a -> Int -> m (Maybe a)
VPM.readMaybe MVector (PrimState m) Int
arr Int
i
#else
lookup (MutableSparseArray arr) i
  | i < 0 || i >= VPM.length arr = pure Nothing
  | otherwise = msaReprToMaybe <$> VPM.unsafeRead arr i
#endif
{-# INLINE lookup #-}

unsafeInsert
  :: (PrimMonad m)
  => MutableSparseArray (PrimState m)
  -> Int
  -> Int
  -> m (MutableSparseArray (PrimState m))
unsafeInsert :: forall (m :: * -> *).
PrimMonad m =>
MutableSparseArray (PrimState m)
-> Int -> Int -> m (MutableSparseArray (PrimState m))
unsafeInsert (MutableSparseArray MVector (PrimState m) Int
arr) Int
i Int
v
  | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = [Char] -> m (MutableSparseArray (PrimState m))
forall a. HasCallStack => [Char] -> a
error ([Char] -> m (MutableSparseArray (PrimState m)))
-> [Char] -> m (MutableSparseArray (PrimState m))
forall a b. (a -> b) -> a -> b
$ [Char]
"Negative index " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i
  | Bool
otherwise = do
      let len :: Int
len = MVector (PrimState m) Int -> Int
forall a s. Prim a => MVector s a -> Int
VPM.length MVector (PrimState m) Int
arr
          growBy :: Int
growBy = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int
len 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)
      MVector (PrimState m) Int
mArr <-
        if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len
          then do
            MVector (PrimState m) Int
r <- MVector (PrimState m) Int -> Int -> m (MVector (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
VPM.unsafeGrow MVector (PrimState m) Int
arr Int
growBy
            Int -> Int -> MVector (PrimState m) Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *).
(PrimMonad m, MVector v Int) =>
Int -> Int -> v (PrimState m) Int -> m ()
fillArray Int
len Int
growBy MVector (PrimState m) Int
r
            MVector (PrimState m) Int -> m (MVector (PrimState m) Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector (PrimState m) Int
r
          else MVector (PrimState m) Int -> m (MVector (PrimState m) Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector (PrimState m) Int
arr

      MVector (PrimState m) Int -> Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VPM.unsafeWrite MVector (PrimState m) Int
mArr Int
i Int
v
      MutableSparseArray (PrimState m)
-> m (MutableSparseArray (PrimState m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MutableSparseArray (PrimState m)
 -> m (MutableSparseArray (PrimState m)))
-> MutableSparseArray (PrimState m)
-> m (MutableSparseArray (PrimState m))
forall a b. (a -> b) -> a -> b
$ MVector (PrimState m) Int -> MutableSparseArray (PrimState m)
forall s. MVector s Int -> MutableSparseArray s
MutableSparseArray MVector (PrimState m) Int
mArr
{-# INLINE unsafeInsert #-}

delete :: (PrimMonad m) => MutableSparseArray (PrimState m) -> Int -> m (Maybe Int)
delete :: forall (m :: * -> *).
PrimMonad m =>
MutableSparseArray (PrimState m) -> Int -> m (Maybe Int)
delete (MutableSparseArray MVector (PrimState m) Int
arr) Int
i
  | Int
i 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
>= MVector (PrimState m) Int -> Int
forall a s. Prim a => MVector s a -> Int
VPM.length MVector (PrimState m) Int
arr = Maybe Int -> m (Maybe Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Int
forall a. Maybe a
Nothing
  | Bool
otherwise = Int -> Maybe Int
msaReprToMaybe (Int -> Maybe Int) -> m Int -> m (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) Int -> Int -> Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MVector (PrimState m) a -> Int -> a -> m a
VPM.unsafeExchange MVector (PrimState m) Int
arr Int
i Int
ABSURD
{-# INLINE delete #-}

-- | Currently checks that the index is not negative, but this may change in the future
--
-- @since 0.1.0.0
unsafeDelete :: (PrimMonad m) => MutableSparseArray (PrimState m) -> Int -> m (Maybe Int)
unsafeDelete :: forall (m :: * -> *).
PrimMonad m =>
MutableSparseArray (PrimState m) -> Int -> m (Maybe Int)
unsafeDelete (MutableSparseArray MVector (PrimState m) Int
arr) Int
i
  | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = [Char] -> m (Maybe Int)
forall a. HasCallStack => [Char] -> a
error ([Char] -> m (Maybe Int)) -> [Char] -> m (Maybe Int)
forall a b. (a -> b) -> a -> b
$ [Char]
"Negative index " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i
  | Bool
otherwise = Int -> Maybe Int
msaReprToMaybe (Int -> Maybe Int) -> m Int -> m (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) Int -> Int -> Int -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MVector (PrimState m) a -> Int -> a -> m a
VPM.unsafeExchange MVector (PrimState m) Int
arr Int
i Int
ABSURD
{-# INLINE unsafeDelete #-}

clear :: (PrimMonad m) => MutableSparseArray (PrimState m) -> m ()
clear :: forall (m :: * -> *).
PrimMonad m =>
MutableSparseArray (PrimState m) -> m ()
clear (MutableSparseArray MVector (PrimState m) Int
arr) = MVector (PrimState m) Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MVector (PrimState m) a -> a -> m ()
VPM.set MVector (PrimState m) Int
arr Int
ABSURD
{-# INLINE clear #-}

unsafeCompactTo
  :: (PrimMonad m) => MutableSparseArray (PrimState m) -> Int -> m (MutableSparseArray (PrimState m))
unsafeCompactTo :: forall (m :: * -> *).
PrimMonad m =>
MutableSparseArray (PrimState m)
-> Int -> m (MutableSparseArray (PrimState m))
unsafeCompactTo (MutableSparseArray MVector (PrimState m) Int
arr) Int
len
  | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = [Char] -> m (MutableSparseArray (PrimState m))
forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot compact to negative capacity"
  | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= MVector (PrimState m) Int -> Int
forall a s. Prim a => MVector s a -> Int
VPM.length MVector (PrimState m) Int
arr = MutableSparseArray (PrimState m)
-> m (MutableSparseArray (PrimState m))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MutableSparseArray (PrimState m)
 -> m (MutableSparseArray (PrimState m)))
-> MutableSparseArray (PrimState m)
-> m (MutableSparseArray (PrimState m))
forall a b. (a -> b) -> a -> b
$ MVector (PrimState m) Int -> MutableSparseArray (PrimState m)
forall s. MVector s Int -> MutableSparseArray s
MutableSparseArray MVector (PrimState m) Int
arr
  | Bool
otherwise = MVector (PrimState m) Int -> MutableSparseArray (PrimState m)
forall s. MVector s Int -> MutableSparseArray s
MutableSparseArray (MVector (PrimState m) Int -> MutableSparseArray (PrimState m))
-> m (MVector (PrimState m) Int)
-> m (MutableSparseArray (PrimState m))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) Int -> m (MVector (PrimState m) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MVector (PrimState m) a -> m (MVector (PrimState m) a)
VPM.clone (Int
-> Int -> MVector (PrimState m) Int -> MVector (PrimState m) Int
forall a s. Prim a => Int -> Int -> MVector s a -> MVector s a
VPM.slice Int
0 Int
len MVector (PrimState m) Int
arr)
{-# INLINE unsafeCompactTo #-}

freeze :: (PrimMonad m) => MutableSparseArray (PrimState m) -> m (VP.Vector Int)
freeze :: forall (m :: * -> *).
PrimMonad m =>
MutableSparseArray (PrimState m) -> m (Vector Int)
freeze (MutableSparseArray MVector (PrimState m) Int
arr) = MVector (PrimState m) Int -> m (Vector Int)
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VP.freeze MVector (PrimState m) Int
arr
{-# INLINE freeze #-}

unsafeFreeze :: (PrimMonad m) => MutableSparseArray (PrimState m) -> m (VP.Vector Int)
unsafeFreeze :: forall (m :: * -> *).
PrimMonad m =>
MutableSparseArray (PrimState m) -> m (Vector Int)
unsafeFreeze (MutableSparseArray MVector (PrimState m) Int
arr) = MVector (PrimState m) Int -> m (Vector Int)
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VP.unsafeFreeze MVector (PrimState m) Int
arr
{-# INLINE unsafeFreeze #-}

--------------------------------------------------------------------------------
-- Utilities
--------------------------------------------------------------------------------
msaReprToMaybe :: Int -> Maybe Int
msaReprToMaybe :: Int -> Maybe Int
msaReprToMaybe Int
v
  | Int
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
ABSURD = Maybe Int
forall a. Maybe a
Nothing
  | Bool
otherwise = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
v
{-# INLINE msaReprToMaybe #-}

fillArray :: (PrimMonad m, VGM.MVector v Int) => Int -> Int -> v (PrimState m) Int -> m ()
fillArray :: forall (m :: * -> *) (v :: * -> * -> *).
(PrimMonad m, MVector v Int) =>
Int -> Int -> v (PrimState m) Int -> m ()
fillArray Int
len Int
growBy v (PrimState m) Int
arr = ST (PrimState m) () -> m ()
forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (ST (PrimState m) () -> m ()) -> ST (PrimState m) () -> m ()
forall a b. (a -> b) -> a -> b
$ v (PrimState m) Int -> Int -> ST (PrimState m) ()
forall s. v s Int -> Int -> ST s ()
forall (v :: * -> * -> *) a s. MVector v a => v s a -> a -> ST s ()
VGM.basicSet (Int -> Int -> v (PrimState m) Int -> v (PrimState m) Int
forall s. Int -> Int -> v s Int -> v s Int
forall (v :: * -> * -> *) a s.
MVector v a =>
Int -> Int -> v s a -> v s a
VGM.basicUnsafeSlice Int
len Int
growBy v (PrimState m) Int
arr) Int
ABSURD
{-# INLINE fillArray #-}