-- |

-- Module      : Data.SparseVector.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.SparseVector.Mutable
  ( MSparseVector (..),
    empty,
    insert,
    read,
    unsafeRead,
    write,
    unsafeWrite,
    modify,
    unsafeModify,
    toList,
  )
where

import qualified Data.Vector as V
import Data.Vector.Mutable (MVector, PrimMonad (..))
import qualified Data.Vector.Mutable as MV
import Prelude hiding (read)

newtype MSparseVector s a = MSparseVector {forall s a. MSparseVector s a -> MVector s (Maybe a)
unMSparseVector :: MVector s (Maybe a)}

empty :: (PrimMonad m) => m (MSparseVector (PrimState m) a)
empty :: forall (m :: * -> *) a.
PrimMonad m =>
m (MSparseVector (PrimState m) a)
empty = do
  MVector (PrimState m) (Maybe a)
vec <- Int -> m (MVector (PrimState m) (Maybe a))
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
MV.new Int
0
  MSparseVector (PrimState m) a -> m (MSparseVector (PrimState m) a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (MSparseVector (PrimState m) a
 -> m (MSparseVector (PrimState m) a))
-> MSparseVector (PrimState m) a
-> m (MSparseVector (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ MVector (PrimState m) (Maybe a) -> MSparseVector (PrimState m) a
forall s a. MVector s (Maybe a) -> MSparseVector s a
MSparseVector MVector (PrimState m) (Maybe a)
vec

insert ::
  (PrimMonad m) =>
  Int ->
  a ->
  MSparseVector (PrimState m) a ->
  m (MSparseVector (PrimState m) a)
insert :: forall (m :: * -> *) a.
PrimMonad m =>
Int
-> a
-> MSparseVector (PrimState m) a
-> m (MSparseVector (PrimState m) a)
insert Int
index a
a (MSparseVector MVector (PrimState m) (Maybe a)
vec) = do
  let len :: Int
len = MVector (PrimState m) (Maybe a) -> Int
forall s a. MVector s a -> Int
MV.length MVector (PrimState m) (Maybe a)
vec
  if Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
index Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    then MVector (PrimState m) (Maybe a) -> Int -> Maybe a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector (PrimState m) (Maybe a)
vec Int
index (a -> Maybe a
forall a. a -> Maybe a
Just a
a) m ()
-> m (MSparseVector (PrimState m) a)
-> m (MSparseVector (PrimState m) a)
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> MSparseVector (PrimState m) a -> m (MSparseVector (PrimState m) a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (MVector (PrimState m) (Maybe a) -> MSparseVector (PrimState m) a
forall s a. MVector s (Maybe a) -> MSparseVector s a
MSparseVector MVector (PrimState m) (Maybe a)
vec)
    else do
      MVector (PrimState m) (Maybe a)
newVec <- Int -> Maybe a -> m (MVector (PrimState m) (Maybe a))
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MVector (PrimState m) a)
MV.replicate (Int
index Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Maybe a
forall a. Maybe a
Nothing
      MVector (PrimState m) (Maybe a) -> Int -> Maybe a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector (PrimState m) (Maybe a)
newVec Int
index (a -> Maybe a
forall a. a -> Maybe a
Just a
a)
      MSparseVector (PrimState m) a -> m (MSparseVector (PrimState m) a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (MSparseVector (PrimState m) a
 -> m (MSparseVector (PrimState m) a))
-> MSparseVector (PrimState m) a
-> m (MSparseVector (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ MVector (PrimState m) (Maybe a) -> MSparseVector (PrimState m) a
forall s a. MVector s (Maybe a) -> MSparseVector s a
MSparseVector MVector (PrimState m) (Maybe a)
newVec

read :: (PrimMonad m) => MSparseVector (PrimState m) a -> Int -> m (Maybe a)
read :: forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> Int -> m (Maybe a)
read (MSparseVector MVector (PrimState m) (Maybe a)
vec) = MVector (PrimState m) (Maybe a) -> Int -> m (Maybe a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.read MVector (PrimState m) (Maybe a)
vec
{-# INLINE read #-}

unsafeRead :: (PrimMonad m) => MSparseVector (PrimState m) a -> Int -> m (Maybe a)
unsafeRead :: forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> Int -> m (Maybe a)
unsafeRead (MSparseVector MVector (PrimState m) (Maybe a)
vec) = MVector (PrimState m) (Maybe a) -> Int -> m (Maybe a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.unsafeRead MVector (PrimState m) (Maybe a)
vec
{-# INLINE unsafeRead #-}

write :: (PrimMonad m) => MSparseVector (PrimState m) a -> Int -> Maybe a -> m ()
write :: forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> Int -> Maybe a -> m ()
write (MSparseVector MVector (PrimState m) (Maybe a)
vec) = MVector (PrimState m) (Maybe a) -> Int -> Maybe a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector (PrimState m) (Maybe a)
vec
{-# INLINE write #-}

unsafeWrite :: (PrimMonad m) => MSparseVector (PrimState m) a -> Int -> Maybe a -> m ()
unsafeWrite :: forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> Int -> Maybe a -> m ()
unsafeWrite (MSparseVector MVector (PrimState m) (Maybe a)
vec) = MVector (PrimState m) (Maybe a) -> Int -> Maybe a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.unsafeWrite MVector (PrimState m) (Maybe a)
vec
{-# INLINE unsafeWrite #-}

modify :: (PrimMonad m) => MSparseVector (PrimState m) a -> Int -> (Maybe a -> Maybe a) -> m ()
modify :: forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a
-> Int -> (Maybe a -> Maybe a) -> m ()
modify (MSparseVector MVector (PrimState m) (Maybe a)
vec) Int
index Maybe a -> Maybe a
f = do
  Maybe a
val <- MVector (PrimState m) (Maybe a) -> Int -> m (Maybe a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.read MVector (PrimState m) (Maybe a)
vec Int
index
  MVector (PrimState m) (Maybe a) -> Int -> Maybe a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector (PrimState m) (Maybe a)
vec Int
index (Maybe a -> Maybe a
f Maybe a
val)
{-# INLINE modify #-}

unsafeModify :: (PrimMonad m) => MSparseVector (PrimState m) a -> Int -> ( a ->  a) -> m ()
unsafeModify :: forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> Int -> (a -> a) -> m ()
unsafeModify (MSparseVector MVector (PrimState m) (Maybe a)
vec) Int
index a -> a
f = do
  Maybe a
val <- MVector (PrimState m) (Maybe a) -> Int -> m (Maybe a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.unsafeRead MVector (PrimState m) (Maybe a)
vec Int
index
  case Maybe a
val of
    Maybe a
Nothing -> () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    Just a
v -> MVector (PrimState m) (Maybe a) -> Int -> Maybe a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.unsafeWrite MVector (PrimState m) (Maybe a)
vec Int
index (a -> Maybe a
forall a. a -> Maybe a
Just (a -> a
f a
v))
{-# INLINE unsafeModify #-}

toList :: (PrimMonad m) => MSparseVector (PrimState m) a -> m [Maybe a]
toList :: forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> m [Maybe a]
toList (MSparseVector MVector (PrimState m) (Maybe a)
v) = Vector (Maybe a) -> [Maybe a]
forall a. Vector a -> [a]
V.toList (Vector (Maybe a) -> [Maybe a])
-> m (Vector (Maybe a)) -> m [Maybe a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) (Maybe a) -> m (Vector (Maybe a))
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.freeze MVector (PrimState m) (Maybe a)
v
{-# INLINE toList #-}