-- |

-- Module      : Data.SparseSet.Strict.Mutable

-- Copyright   : (c) Matt Hunzinger, 2025

-- License     : BSD-style (see the LICENSE file in the distribution)

--

-- Maintainer  : matt@hunzinger.me

-- Stability   : provisional

-- Portability : non-portable (GHC extensions)

module Data.SparseSet.Strict.Mutable
  ( -- * Mutable sparse sets

    MSparseSet (..),

    -- ** Construction

    empty,

    -- ** Operations

    read,
    unsafeRead,
    write,
    unsafeWrite,
    unsafeModify,

    -- ** Conversion

    toList,

    -- ** Re-exports

    PrimMonad (..),
  )
where

import Data.SparseVector.Strict.Mutable (MSparseVector)
import qualified Data.SparseVector.Strict.Mutable as SV
import Data.Vector.Strict.Mutable (MVector, PrimMonad (..))
import qualified Data.Vector.Strict.Mutable as V
import Prelude hiding (lookup, read)

data MSparseSet s i a = MSparseSet
  { forall s i a. MSparseSet s i a -> MVector s a
dense :: {-# UNPACK #-} !(MVector s a),
    forall s i a. MSparseSet s i a -> MSparseVector s i
sparse :: !(MSparseVector s i)
  }

empty :: (PrimMonad m) => m (MSparseSet (PrimState m) i a)
empty :: forall (m :: * -> *) i a.
PrimMonad m =>
m (MSparseSet (PrimState m) i a)
empty = do
  MVector (PrimState m) a
d <- Int -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
V.new Int
0
  MVector (PrimState m) a
-> MSparseVector (PrimState m) i -> MSparseSet (PrimState m) i a
forall s i a. MVector s a -> MSparseVector s i -> MSparseSet s i a
MSparseSet MVector (PrimState m) a
d (MSparseVector (PrimState m) i -> MSparseSet (PrimState m) i a)
-> m (MSparseVector (PrimState m) i)
-> m (MSparseSet (PrimState m) i a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (MSparseVector (PrimState m) i)
forall (m :: * -> *) a.
PrimMonad m =>
m (MSparseVector (PrimState m) a)
SV.empty
{-# INLINE empty #-}

read :: (PrimMonad m, Integral i) => MSparseSet (PrimState m) i a -> Int -> m (Maybe a)
read :: forall (m :: * -> *) i a.
(PrimMonad m, Integral i) =>
MSparseSet (PrimState m) i a -> Int -> m (Maybe a)
read (MSparseSet MVector (PrimState m) a
d MSparseVector (PrimState m) i
s) Int
i = do
  Maybe i
m <- MSparseVector (PrimState m) i -> Int -> m (Maybe i)
forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> Int -> m (Maybe a)
SV.read MSparseVector (PrimState m) i
s Int
i
  case Maybe i
m of
    Maybe i
Nothing -> Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
    Just i
i' -> 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
<$> MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
V.read MVector (PrimState m) a
d (i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
i')
{-# INLINE read #-}

unsafeRead :: (PrimMonad m, Integral i) => MSparseSet (PrimState m) i a -> Int -> m a
unsafeRead :: forall (m :: * -> *) i a.
(PrimMonad m, Integral i) =>
MSparseSet (PrimState m) i a -> Int -> m a
unsafeRead (MSparseSet MVector (PrimState m) a
d MSparseVector (PrimState m) i
s) Int
i = do
  i
i' <- MSparseVector (PrimState m) i -> Int -> m i
forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> Int -> m a
SV.unsafeRead MSparseVector (PrimState m) i
s Int
i
  MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
V.unsafeRead MVector (PrimState m) a
d (i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
i')
{-# INLINE unsafeRead #-}

write :: (PrimMonad m, Integral i) => MSparseSet (PrimState m) i a -> Int -> a -> m ()
write :: forall (m :: * -> *) i a.
(PrimMonad m, Integral i) =>
MSparseSet (PrimState m) i a -> Int -> a -> m ()
write (MSparseSet MVector (PrimState m) a
d MSparseVector (PrimState m) i
s) Int
i a
a = do
  Maybe i
m <- MSparseVector (PrimState m) i -> Int -> m (Maybe i)
forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> Int -> m (Maybe a)
SV.read MSparseVector (PrimState m) i
s Int
i
  case Maybe i
m of
    Just i
i' -> MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
V.write MVector (PrimState m) a
d (i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
i') a
a
    Maybe i
Nothing -> () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
{-# INLINE write #-}

unsafeWrite :: (PrimMonad m, Integral i) => MSparseSet (PrimState m) i a -> Int -> a -> m ()
unsafeWrite :: forall (m :: * -> *) i a.
(PrimMonad m, Integral i) =>
MSparseSet (PrimState m) i a -> Int -> a -> m ()
unsafeWrite (MSparseSet MVector (PrimState m) a
d MSparseVector (PrimState m) i
s) Int
i a
a = do
  i
i' <- MSparseVector (PrimState m) i -> Int -> m i
forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> Int -> m a
SV.unsafeRead MSparseVector (PrimState m) i
s Int
i
  MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
V.unsafeWrite MVector (PrimState m) a
d (i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
i') a
a
{-# INLINE unsafeWrite #-}

unsafeModify :: (PrimMonad m, Integral i) => MSparseSet (PrimState m) i a -> Int -> (a -> a) -> m ()
unsafeModify :: forall (m :: * -> *) i a.
(PrimMonad m, Integral i) =>
MSparseSet (PrimState m) i a -> Int -> (a -> a) -> m ()
unsafeModify (MSparseSet MVector (PrimState m) a
d MSparseVector (PrimState m) i
s) Int
i a -> a
f = do
  i
i' <- MSparseVector (PrimState m) i -> Int -> m i
forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> Int -> m a
SV.unsafeRead MSparseVector (PrimState m) i
s Int
i
  MVector (PrimState m) a -> (a -> a) -> Int -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
V.unsafeModify MVector (PrimState m) a
d a -> a
f (i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
i')
{-# INLINE unsafeModify #-}

toList :: (PrimMonad m, Integral i) => MSparseSet (PrimState m) i a -> m [Maybe (i, a)]
toList :: forall (m :: * -> *) i a.
(PrimMonad m, Integral i) =>
MSparseSet (PrimState m) i a -> m [Maybe (i, a)]
toList MSparseSet (PrimState m) i a
s = do
  [Maybe i]
as <- MSparseVector (PrimState m) i -> m [Maybe i]
forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> m [Maybe a]
SV.toList (MSparseVector (PrimState m) i -> m [Maybe i])
-> MSparseVector (PrimState m) i -> m [Maybe i]
forall a b. (a -> b) -> a -> b
$ MSparseSet (PrimState m) i a -> MSparseVector (PrimState m) i
forall s i a. MSparseSet s i a -> MSparseVector s i
sparse MSparseSet (PrimState m) i a
s
  (Maybe i -> m (Maybe (i, a))) -> [Maybe i] -> m [Maybe (i, a)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Maybe i -> m (Maybe (i, a))
forall {f :: * -> *} {a}.
(PrimState f ~ PrimState m, PrimMonad f, Integral a) =>
Maybe a -> f (Maybe (a, a))
go [Maybe i]
as
  where
    go :: Maybe a -> f (Maybe (a, a))
go (Just a
i) = (\a
a -> (a, a) -> Maybe (a, a)
forall a. a -> Maybe a
Just (a
i, a
a)) (a -> Maybe (a, a)) -> f a -> f (Maybe (a, a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState f) a -> Int -> f a
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
V.unsafeRead (MSparseSet (PrimState f) i a -> MVector (PrimState f) a
forall s i a. MSparseSet s i a -> MVector s a
dense MSparseSet (PrimState m) i a
MSparseSet (PrimState f) i a
s) (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
i)
    go Maybe a
Nothing = Maybe (a, a) -> f (Maybe (a, a))
forall a. a -> f a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (a, a)
forall a. Maybe a
Nothing
{-# INLINE toList #-}