-- |
-- 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,
    write,
    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

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

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 #-}