{-# LANGUAGE BangPatterns #-}

-- |

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

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

newtype MSparseVector s a = MSparseVector {forall s a. MSparseVector s a -> MVector s (Bool, a)
unMSparseVector :: MVector s (Bool, 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) (Bool, a)
vec <- Int -> m (MVector (PrimState m) (Bool, 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) (Bool, a) -> MSparseVector (PrimState m) a
forall s a. MVector s (Bool, a) -> MSparseVector s a
MSparseVector MVector (PrimState m) (Bool, a)
vec

insert ::
  (PrimMonad m) =>
  Int ->
  a ->
  MSparseVector (PrimState m) a ->
  m ()
insert :: forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> MSparseVector (PrimState m) a -> m ()
insert Int
index a
a (MSparseVector MVector (PrimState m) (Bool, a)
vec) = do
  let len :: Int
len = MVector (PrimState m) (Bool, a) -> Int
forall s a. MVector s a -> Int
MV.length MVector (PrimState m) (Bool, 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) (Bool, a) -> Int -> (Bool, a) -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector (PrimState m) (Bool, a)
vec Int
index (Bool
True, a
a)
    else do
      MVector (PrimState m) (Bool, a)
newVec <- Int -> (Bool, a) -> m (MVector (PrimState m) (Bool, 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) (Bool
False, a
forall a. HasCallStack => a
undefined)
      MVector (PrimState m) (Bool, a) -> Int -> (Bool, a) -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector (PrimState m) (Bool, a)
newVec Int
index (Bool
True, a
a)

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) (Bool, a)
vec) Int
index = do
  (Bool
present, a
val) <- MVector (PrimState m) (Bool, a) -> Int -> m (Bool, a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.read MVector (PrimState m) (Bool, a)
vec Int
index
  Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> m (Maybe a)) -> Maybe a -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ if Bool
present then a -> Maybe a
forall a. a -> Maybe a
Just a
val else Maybe a
forall a. Maybe a
Nothing
{-# INLINE read #-}

unsafeRead :: (PrimMonad m) => MSparseVector (PrimState m) a -> Int -> m a
unsafeRead :: forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> Int -> m a
unsafeRead (MSparseVector MVector (PrimState m) (Bool, a)
vec) Int
index = do
  (Bool
_, a
val) <- MVector (PrimState m) (Bool, a) -> Int -> m (Bool, a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.unsafeRead MVector (PrimState m) (Bool, a)
vec Int
index
  a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
val
{-# INLINE unsafeRead #-}

write :: (PrimMonad m) => MSparseVector (PrimState m) a -> Int -> Maybe a -> a -> m ()
write :: forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> Int -> Maybe a -> a -> m ()
write (MSparseVector MVector (PrimState m) (Bool, a)
vec) Int
index Maybe a
maybeVal a
defaultVal =
  case Maybe a
maybeVal of
    Maybe a
Nothing -> MVector (PrimState m) (Bool, a) -> Int -> (Bool, a) -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector (PrimState m) (Bool, a)
vec Int
index (Bool
False, a
defaultVal)
    Just a
val -> MVector (PrimState m) (Bool, a) -> Int -> (Bool, a) -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector (PrimState m) (Bool, a)
vec Int
index (Bool
True, a
val)
{-# INLINE write #-}

unsafeWrite :: (PrimMonad m) => MSparseVector (PrimState m) a -> Int -> a -> m ()
unsafeWrite :: forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> Int -> a -> m ()
unsafeWrite (MSparseVector MVector (PrimState m) (Bool, a)
vec) Int
index a
v = MVector (PrimState m) (Bool, a) -> Int -> (Bool, a) -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.unsafeWrite MVector (PrimState m) (Bool, a)
vec Int
index (Bool
True, a
v)
{-# INLINE unsafeWrite #-}

modify :: (PrimMonad m) => MSparseVector (PrimState m) a -> Int -> (Maybe a -> Maybe a) -> a -> m ()
modify :: forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a
-> Int -> (Maybe a -> Maybe a) -> a -> m ()
modify (MSparseVector MVector (PrimState m) (Bool, a)
vec) Int
index Maybe a -> Maybe a
f a
defaultVal = do
  !(Bool
present, a
val) <- MVector (PrimState m) (Bool, a) -> Int -> m (Bool, a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.read MVector (PrimState m) (Bool, a)
vec Int
index
  let currentVal :: Maybe a
currentVal = if Bool
present then a -> Maybe a
forall a. a -> Maybe a
Just a
val else Maybe a
forall a. Maybe a
Nothing
  case Maybe a -> Maybe a
f Maybe a
currentVal of
    Maybe a
Nothing -> MVector (PrimState m) (Bool, a) -> Int -> (Bool, a) -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector (PrimState m) (Bool, a)
vec Int
index (Bool
False, a
defaultVal)
    Just a
newVal -> MVector (PrimState m) (Bool, a) -> Int -> (Bool, a) -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector (PrimState m) (Bool, a)
vec Int
index (Bool
True, a
newVal)
{-# 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) (Bool, a)
vec) Int
index a -> a
f = do
  !(Bool
present, a
val) <- MVector (PrimState m) (Bool, a) -> Int -> m (Bool, a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.unsafeRead MVector (PrimState m) (Bool, a)
vec Int
index
  if Bool
present
    then MVector (PrimState m) (Bool, a) -> Int -> (Bool, a) -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.unsafeWrite MVector (PrimState m) (Bool, a)
vec Int
index (Bool
True, a -> a
f a
val)
    else () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
{-# 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) (Bool, a)
v) = do
  [(Bool, a)]
pairs <- Vector (Bool, a) -> [(Bool, a)]
forall a. Vector a -> [a]
V.toList (Vector (Bool, a) -> [(Bool, a)])
-> m (Vector (Bool, a)) -> m [(Bool, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) (Bool, a) -> m (Vector (Bool, a))
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.freeze MVector (PrimState m) (Bool, a)
v
  [Maybe a] -> m [Maybe a]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Maybe a] -> m [Maybe a]) -> [Maybe a] -> m [Maybe a]
forall a b. (a -> b) -> a -> b
$ ((Bool, a) -> Maybe a) -> [(Bool, a)] -> [Maybe a]
forall a b. (a -> b) -> [a] -> [b]
map (\(Bool
present, a
val) -> if Bool
present then a -> Maybe a
forall a. a -> Maybe a
Just a
val else Maybe a
forall a. Maybe a
Nothing) [(Bool, a)]
pairs
{-# INLINE toList #-}