{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}

-- |

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

import qualified Data.Vector.Unboxed as V
import Data.Vector.Unboxed.Mutable (MVector, PrimMonad (..), Unbox)
import qualified Data.Vector.Unboxed.Mutable as MV
import Data.Maybe
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, Unbox a) => m (MSparseVector (PrimState m) a)
empty :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
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, Unbox a) =>
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, Unbox a) =>
  Int ->
  a ->
  MSparseVector (PrimState m) a ->
  m ()
insert :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
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 a s. Unbox 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, Unbox a) =>
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, Unbox a) =>
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)
      -- Copy old data

      let copyLoop :: Int -> m ()
copyLoop Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            | Bool
otherwise = do
                (Bool, a)
val <- MVector (PrimState m) (Bool, a) -> Int -> m (Bool, a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MV.read MVector (PrimState m) (Bool, a)
MVector (PrimState m) (Bool, a)
vec Int
i
                MVector (PrimState m) (Bool, a) -> Int -> (Bool, a) -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector (PrimState m) (Bool, a)
MVector (PrimState m) (Bool, a)
newVec Int
i (Bool, a)
val
                Int -> m ()
copyLoop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      Int -> m ()
forall {m :: * -> *}.
(PrimState m ~ PrimState m, PrimMonad m) =>
Int -> m ()
copyLoop Int
0
      MVector (PrimState m) (Bool, a) -> Int -> (Bool, a) -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector (PrimState m) (Bool, a)
newVec Int
index (Bool
True, a
a)

read :: (PrimMonad m, Unbox a) => MSparseVector (PrimState m) a -> Int -> m (Maybe a)
read :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MSparseVector (PrimState m) a -> Int -> m (Maybe a)
read (MSparseVector MVector (PrimState m) (Bool, a)
vec) 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) (Bool, a) -> Int
forall a s. Unbox a => MVector s a -> Int
MV.length MVector (PrimState m) (Bool, a)
vec = 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
  | Bool
otherwise = do
      (Bool
present, a
val) <- MVector (PrimState m) (Bool, a) -> Int -> m (Bool, a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MV.read MVector (PrimState m) (Bool, a)
vec Int
i
      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, Unbox a) => MSparseVector (PrimState m) a -> Int -> m a
unsafeRead :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
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, Unbox a) =>
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, Unbox a) => MSparseVector (PrimState m) a -> Int -> Maybe a -> a -> m ()
write :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MSparseVector (PrimState m) a -> Int -> Maybe a -> a -> m ()
write (MSparseVector MVector (PrimState m) (Bool, a)
vec) Int
i Maybe a
maybeVal a
defaultVal
  | 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) (Bool, a) -> Int
forall a s. Unbox a => MVector s a -> Int
MV.length MVector (PrimState m) (Bool, a)
vec = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return () -- Out of bounds, ignore

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

unsafeWrite :: (PrimMonad m, Unbox a) => MSparseVector (PrimState m) a -> Int -> a -> m ()
unsafeWrite :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
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, Unbox a) =>
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, Unbox a) => MSparseVector (PrimState m) a -> Int -> (Maybe a -> Maybe a) -> a -> m ()
modify :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
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
  | Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= MVector (PrimState m) (Bool, a) -> Int
forall a s. Unbox a => MVector s a -> Int
MV.length MVector (PrimState m) (Bool, a)
vec = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return () -- Out of bounds, ignore

  | Bool
otherwise = do
      !(Bool
present, a
val) <- MVector (PrimState m) (Bool, a) -> Int -> m (Bool, a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
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, Unbox a) =>
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, Unbox a) =>
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, Unbox a) => MSparseVector (PrimState m) a -> Int -> (a -> a) -> m ()
unsafeModify :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
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, Unbox a) =>
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, Unbox a) =>
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, Unbox a) => MSparseVector (PrimState m) a -> m [Maybe a]
toList :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
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. Unbox 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 a (m :: * -> *).
(Unbox 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 #-}