{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

-- |
-- Module      : Data.SparseVector
-- 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
  ( -- * Sparse vectors
    SparseVector (..),

    -- * Construction
    empty,

    -- ** Operations
    insert,
    lookup,
    delete,
    mapWithKey,
    mapAccum,

    -- ** Intersection
    intersection,
    intersectionWith,
    intersectionWithKey,
    intersectionVec,
    intersectionVecWith,
    intersectionVecWithKey,

    -- ** Conversion
    fromList,
    toList,
    fromVector,
    toVector,

    -- ** Mutations
    freeze,
    thaw,
  )
where

import Control.DeepSeq
import Control.Monad
import Control.Monad.State.Strict
import Data.Maybe
import Data.SparseVector.Mutable (MSparseVector (..))
import Data.Vector (Vector)
import qualified Data.Vector as V
import Data.Vector.Mutable (PrimMonad (..))
import Prelude hiding (lookup)

-- | Sparse n-dimensional vector.
--
-- A sparse vector is defined as a @Vector (Maybe a)@,
-- where @Maybe a@ is a cell for an element in the sparse vector.
--
-- Inserting elements at some dimension @n@ will grow the vector up to @n@,
-- using @Nothing@ to create empty cells.
newtype SparseVector a = SparseVector {forall a. SparseVector a -> Vector (Maybe a)
unSparseVector :: Vector (Maybe a)}
  deriving (Int -> SparseVector a -> ShowS
[SparseVector a] -> ShowS
SparseVector a -> String
(Int -> SparseVector a -> ShowS)
-> (SparseVector a -> String)
-> ([SparseVector a] -> ShowS)
-> Show (SparseVector a)
forall a. Show a => Int -> SparseVector a -> ShowS
forall a. Show a => [SparseVector a] -> ShowS
forall a. Show a => SparseVector a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> SparseVector a -> ShowS
showsPrec :: Int -> SparseVector a -> ShowS
$cshow :: forall a. Show a => SparseVector a -> String
show :: SparseVector a -> String
$cshowList :: forall a. Show a => [SparseVector a] -> ShowS
showList :: [SparseVector a] -> ShowS
Show, SparseVector a -> SparseVector a -> Bool
(SparseVector a -> SparseVector a -> Bool)
-> (SparseVector a -> SparseVector a -> Bool)
-> Eq (SparseVector a)
forall a. Eq a => SparseVector a -> SparseVector a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => SparseVector a -> SparseVector a -> Bool
== :: SparseVector a -> SparseVector a -> Bool
$c/= :: forall a. Eq a => SparseVector a -> SparseVector a -> Bool
/= :: SparseVector a -> SparseVector a -> Bool
Eq, (forall a b. (a -> b) -> SparseVector a -> SparseVector b)
-> (forall a b. a -> SparseVector b -> SparseVector a)
-> Functor SparseVector
forall a b. a -> SparseVector b -> SparseVector a
forall a b. (a -> b) -> SparseVector a -> SparseVector b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> SparseVector a -> SparseVector b
fmap :: forall a b. (a -> b) -> SparseVector a -> SparseVector b
$c<$ :: forall a b. a -> SparseVector b -> SparseVector a
<$ :: forall a b. a -> SparseVector b -> SparseVector a
Functor, (forall m. Monoid m => SparseVector m -> m)
-> (forall m a. Monoid m => (a -> m) -> SparseVector a -> m)
-> (forall m a. Monoid m => (a -> m) -> SparseVector a -> m)
-> (forall a b. (a -> b -> b) -> b -> SparseVector a -> b)
-> (forall a b. (a -> b -> b) -> b -> SparseVector a -> b)
-> (forall b a. (b -> a -> b) -> b -> SparseVector a -> b)
-> (forall b a. (b -> a -> b) -> b -> SparseVector a -> b)
-> (forall a. (a -> a -> a) -> SparseVector a -> a)
-> (forall a. (a -> a -> a) -> SparseVector a -> a)
-> (forall a. SparseVector a -> [a])
-> (forall a. SparseVector a -> Bool)
-> (forall a. SparseVector a -> Int)
-> (forall a. Eq a => a -> SparseVector a -> Bool)
-> (forall a. Ord a => SparseVector a -> a)
-> (forall a. Ord a => SparseVector a -> a)
-> (forall a. Num a => SparseVector a -> a)
-> (forall a. Num a => SparseVector a -> a)
-> Foldable SparseVector
forall a. Eq a => a -> SparseVector a -> Bool
forall a. Num a => SparseVector a -> a
forall a. Ord a => SparseVector a -> a
forall m. Monoid m => SparseVector m -> m
forall a. SparseVector a -> Bool
forall a. SparseVector a -> Int
forall a. SparseVector a -> [a]
forall a. (a -> a -> a) -> SparseVector a -> a
forall m a. Monoid m => (a -> m) -> SparseVector a -> m
forall b a. (b -> a -> b) -> b -> SparseVector a -> b
forall a b. (a -> b -> b) -> b -> SparseVector a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
$cfold :: forall m. Monoid m => SparseVector m -> m
fold :: forall m. Monoid m => SparseVector m -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> SparseVector a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> SparseVector a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> SparseVector a -> m
foldMap' :: forall m a. Monoid m => (a -> m) -> SparseVector a -> m
$cfoldr :: forall a b. (a -> b -> b) -> b -> SparseVector a -> b
foldr :: forall a b. (a -> b -> b) -> b -> SparseVector a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> SparseVector a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> SparseVector a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> SparseVector a -> b
foldl :: forall b a. (b -> a -> b) -> b -> SparseVector a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> SparseVector a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> SparseVector a -> b
$cfoldr1 :: forall a. (a -> a -> a) -> SparseVector a -> a
foldr1 :: forall a. (a -> a -> a) -> SparseVector a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> SparseVector a -> a
foldl1 :: forall a. (a -> a -> a) -> SparseVector a -> a
$ctoList :: forall a. SparseVector a -> [a]
toList :: forall a. SparseVector a -> [a]
$cnull :: forall a. SparseVector a -> Bool
null :: forall a. SparseVector a -> Bool
$clength :: forall a. SparseVector a -> Int
length :: forall a. SparseVector a -> Int
$celem :: forall a. Eq a => a -> SparseVector a -> Bool
elem :: forall a. Eq a => a -> SparseVector a -> Bool
$cmaximum :: forall a. Ord a => SparseVector a -> a
maximum :: forall a. Ord a => SparseVector a -> a
$cminimum :: forall a. Ord a => SparseVector a -> a
minimum :: forall a. Ord a => SparseVector a -> a
$csum :: forall a. Num a => SparseVector a -> a
sum :: forall a. Num a => SparseVector a -> a
$cproduct :: forall a. Num a => SparseVector a -> a
product :: forall a. Num a => SparseVector a -> a
Foldable, Functor SparseVector
Foldable SparseVector
(Functor SparseVector, Foldable SparseVector) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> SparseVector a -> f (SparseVector b))
-> (forall (f :: * -> *) a.
    Applicative f =>
    SparseVector (f a) -> f (SparseVector a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> SparseVector a -> m (SparseVector b))
-> (forall (m :: * -> *) a.
    Monad m =>
    SparseVector (m a) -> m (SparseVector a))
-> Traversable SparseVector
forall (t :: * -> *).
(Functor t, Foldable t) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a.
Monad m =>
SparseVector (m a) -> m (SparseVector a)
forall (f :: * -> *) a.
Applicative f =>
SparseVector (f a) -> f (SparseVector a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> SparseVector a -> m (SparseVector b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> SparseVector a -> f (SparseVector b)
$ctraverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> SparseVector a -> f (SparseVector b)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> SparseVector a -> f (SparseVector b)
$csequenceA :: forall (f :: * -> *) a.
Applicative f =>
SparseVector (f a) -> f (SparseVector a)
sequenceA :: forall (f :: * -> *) a.
Applicative f =>
SparseVector (f a) -> f (SparseVector a)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> SparseVector a -> m (SparseVector b)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> SparseVector a -> m (SparseVector b)
$csequence :: forall (m :: * -> *) a.
Monad m =>
SparseVector (m a) -> m (SparseVector a)
sequence :: forall (m :: * -> *) a.
Monad m =>
SparseVector (m a) -> m (SparseVector a)
Traversable, SparseVector a -> ()
(SparseVector a -> ()) -> NFData (SparseVector a)
forall a. NFData a => SparseVector a -> ()
forall a. (a -> ()) -> NFData a
$crnf :: forall a. NFData a => SparseVector a -> ()
rnf :: SparseVector a -> ()
NFData)

instance Semigroup (SparseVector a) where
  SparseVector Vector (Maybe a)
v1 <> :: SparseVector a -> SparseVector a -> SparseVector a
<> SparseVector Vector (Maybe a)
v2 =
    let (Vector (Maybe a)
lhs, Vector (Maybe a)
rhs) = if Vector (Maybe a) -> Int
forall a. Vector a -> Int
V.length Vector (Maybe a)
v1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Vector (Maybe a) -> Int
forall a. Vector a -> Int
V.length Vector (Maybe a)
v2 then (Vector (Maybe a)
v1, Vector (Maybe a)
v2) else (Vector (Maybe a)
v2, Vector (Maybe a)
v1)
     in Vector (Maybe a) -> SparseVector a
forall a. Vector (Maybe a) -> SparseVector a
SparseVector (Vector (Maybe a) -> SparseVector a)
-> Vector (Maybe a) -> SparseVector a
forall a b. (a -> b) -> a -> b
$ Vector (Maybe a) -> Vector (Int, Maybe a) -> Vector (Maybe a)
forall a. Vector a -> Vector (Int, a) -> Vector a
V.update Vector (Maybe a)
lhs (Vector (Maybe a) -> Vector (Int, Maybe a)
forall a. Vector a -> Vector (Int, a)
V.indexed Vector (Maybe a)
rhs)
  {-# INLINE (<>) #-}

instance Monoid (SparseVector a) where
  mempty :: SparseVector a
mempty = SparseVector a
forall a. SparseVector a
empty
  {-# INLINE mempty #-}

-- | Empty sparse vector.
empty :: SparseVector a
empty :: forall a. SparseVector a
empty = Vector (Maybe a) -> SparseVector a
forall a. Vector (Maybe a) -> SparseVector a
SparseVector Vector (Maybe a)
forall a. Vector a
V.empty
{-# INLINE empty #-}

-- | Insert an element at a given index into a `SparseVector`.
--
-- Inserting elements at some dimension @n@ will grow the vector up to @n@,
-- using @Nothing@ to create empty cells.
--
-- >>> insert 0 'a' empty
-- SparseVector {unSparseVector = [Just 'a']}
--
-- >>> insert 2 'b' empty
-- SparseVector {unSparseVector = [Nothing,Nothing,Just 'b']}
insert :: Int -> a -> SparseVector a -> SparseVector a
insert :: forall a. Int -> a -> SparseVector a -> SparseVector a
insert Int
index a
a (SparseVector Vector (Maybe a)
vec) =
  let len :: Int
len = Vector (Maybe a) -> Int
forall a. Vector a -> Int
V.length Vector (Maybe a)
vec
   in Vector (Maybe a) -> SparseVector a
forall a. Vector (Maybe a) -> SparseVector a
SparseVector (Vector (Maybe a) -> SparseVector a)
-> Vector (Maybe a) -> SparseVector a
forall a b. (a -> b) -> a -> b
$
        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 Vector (Maybe a) -> [(Int, Maybe a)] -> Vector (Maybe a)
forall a. Vector a -> [(Int, a)] -> Vector a
V.unsafeUpd Vector (Maybe a)
vec [(Int
index, a -> Maybe a
forall a. a -> Maybe a
Just a
a)]
          else Vector (Maybe a) -> Maybe a -> Vector (Maybe a)
forall a. Vector a -> a -> Vector a
V.snoc (Vector (Maybe a)
vec Vector (Maybe a) -> Vector (Maybe a) -> Vector (Maybe a)
forall a. Vector a -> Vector a -> Vector a
V.++ Int -> Maybe a -> Vector (Maybe a)
forall a. Int -> a -> Vector a
V.replicate (Int
index Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len) Maybe a
forall a. Maybe a
Nothing) (a -> Maybe a
forall a. a -> Maybe a
Just a
a)
{-# INLINE insert #-}

-- | Lookup an element at a given index in a `SparseVector`.
lookup :: Int -> SparseVector a -> Maybe a
lookup :: forall a. Int -> SparseVector a -> Maybe a
lookup Int
i (SparseVector Vector (Maybe a)
v) = Maybe (Maybe a) -> Maybe a
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (Maybe (Maybe a) -> Maybe a) -> Maybe (Maybe a) -> Maybe a
forall a b. (a -> b) -> a -> b
$ Vector (Maybe a)
v Vector (Maybe a) -> Int -> Maybe (Maybe a)
forall a. Vector a -> Int -> Maybe a
V.!? Int
i
{-# INLINE lookup #-}

-- | Delete an index from a `SparseVector`, replacing its cell with @Nothing@.
delete :: Int -> SparseVector a -> SparseVector a
delete :: forall a. Int -> SparseVector a -> SparseVector a
delete Int
index (SparseVector Vector (Maybe a)
vec) =
  Vector (Maybe a) -> SparseVector a
forall a. Vector (Maybe a) -> SparseVector a
SparseVector (Vector (Maybe a) -> SparseVector a)
-> Vector (Maybe a) -> SparseVector a
forall a b. (a -> b) -> a -> b
$ Vector (Maybe a) -> [(Int, Maybe a)] -> Vector (Maybe a)
forall a. Vector a -> [(Int, a)] -> Vector a
V.unsafeUpd Vector (Maybe a)
vec [(Int
index, Maybe a
forall a. Maybe a
Nothing)]
{-# INLINE delete #-}

mapWithKey :: (Int -> a -> b) -> SparseVector a -> SparseVector b
mapWithKey :: forall a b. (Int -> a -> b) -> SparseVector a -> SparseVector b
mapWithKey Int -> a -> b
f (SparseVector Vector (Maybe a)
v) =
  let go :: (Int, Maybe a) -> Maybe b
go (Int
i, Just a
a) = b -> Maybe b
forall a. a -> Maybe a
Just (b -> Maybe b) -> b -> Maybe b
forall a b. (a -> b) -> a -> b
$ Int -> a -> b
f Int
i a
a
      go (Int, Maybe a)
_ = Maybe b
forall a. Maybe a
Nothing
   in Vector (Maybe b) -> SparseVector b
forall a. Vector (Maybe a) -> SparseVector a
SparseVector ((Int, Maybe a) -> Maybe b
go ((Int, Maybe a) -> Maybe b)
-> Vector (Int, Maybe a) -> Vector (Maybe b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector (Maybe a) -> Vector (Int, Maybe a)
forall a. Vector a -> Vector (Int, a)
V.indexed Vector (Maybe a)
v)
{-# INLINE mapWithKey #-}

mapAccum :: (a -> b -> (a, c)) -> a -> SparseVector b -> (a, SparseVector c)
mapAccum :: forall a b c.
(a -> b -> (a, c)) -> a -> SparseVector b -> (a, SparseVector c)
mapAccum a -> b -> (a, c)
f a
a (SparseVector Vector (Maybe b)
v) =
  let f' :: Maybe b -> m (Maybe c)
f' (Just b
b) = do
        a
acc <- m a
forall s (m :: * -> *). MonadState s m => m s
get
        let (a
acc', c
c) = a -> b -> (a, c)
f a
acc b
b
        a -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put a
acc'
        Maybe c -> m (Maybe c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (c -> Maybe c
forall a. a -> Maybe a
Just c
c)
      f' Maybe b
Nothing = Maybe c -> m (Maybe c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe c
forall a. Maybe a
Nothing
      (Vector (Maybe c)
v', a
a') = State a (Vector (Maybe c)) -> a -> (Vector (Maybe c), a)
forall s a. State s a -> s -> (a, s)
runState ((Maybe b -> StateT a Identity (Maybe c))
-> Vector (Maybe b) -> State a (Vector (Maybe c))
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM Maybe b -> StateT a Identity (Maybe c)
forall {m :: * -> *}. MonadState a m => Maybe b -> m (Maybe c)
f' Vector (Maybe b)
v) a
a
   in (a
a', Vector (Maybe c) -> SparseVector c
forall a. Vector (Maybe a) -> SparseVector a
SparseVector Vector (Maybe c)
v')

intersection :: SparseVector a -> SparseVector b -> SparseVector a
intersection :: forall a b. SparseVector a -> SparseVector b -> SparseVector a
intersection = (a -> b -> a) -> SparseVector a -> SparseVector b -> SparseVector a
forall a b c.
(a -> b -> c) -> SparseVector a -> SparseVector b -> SparseVector c
intersectionWith a -> b -> a
forall a b. a -> b -> a
const

intersectionWith :: (a -> b -> c) -> SparseVector a -> SparseVector b -> SparseVector c
intersectionWith :: forall a b c.
(a -> b -> c) -> SparseVector a -> SparseVector b -> SparseVector c
intersectionWith = (Int -> a -> b -> c)
-> SparseVector a -> SparseVector b -> SparseVector c
forall a b c.
(Int -> a -> b -> c)
-> SparseVector a -> SparseVector b -> SparseVector c
intersectionWithKey ((Int -> a -> b -> c)
 -> SparseVector a -> SparseVector b -> SparseVector c)
-> ((a -> b -> c) -> Int -> a -> b -> c)
-> (a -> b -> c)
-> SparseVector a
-> SparseVector b
-> SparseVector c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b -> c) -> Int -> a -> b -> c
forall a b. a -> b -> a
const

intersectionWithKey :: (Int -> a -> b -> c) -> SparseVector a -> SparseVector b -> SparseVector c
intersectionWithKey :: forall a b c.
(Int -> a -> b -> c)
-> SparseVector a -> SparseVector b -> SparseVector c
intersectionWithKey Int -> a -> b -> c
f (SparseVector Vector (Maybe a)
a) (SparseVector Vector (Maybe b)
b) =
  let (Vector (Maybe a)
as, Vector (Maybe b)
bs) =
        if Vector (Maybe a) -> Int
forall a. Vector a -> Int
V.length Vector (Maybe a)
a Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Vector (Maybe b) -> Int
forall a. Vector a -> Int
V.length Vector (Maybe b)
b
          then (Vector (Maybe a)
a, Vector (Maybe b)
b Vector (Maybe b) -> Vector (Maybe b) -> Vector (Maybe b)
forall a. Vector a -> Vector a -> Vector a
V.++ Int -> Maybe b -> Vector (Maybe b)
forall a. Int -> a -> Vector a
V.replicate (Vector (Maybe a) -> Int
forall a. Vector a -> Int
V.length Vector (Maybe a)
a Int -> Int -> Int
forall a. Num a => a -> a -> a
- Vector (Maybe b) -> Int
forall a. Vector a -> Int
V.length Vector (Maybe b)
b) Maybe b
forall a. Maybe a
Nothing)
          else (Vector (Maybe a)
a Vector (Maybe a) -> Vector (Maybe a) -> Vector (Maybe a)
forall a. Vector a -> Vector a -> Vector a
V.++ Int -> Maybe a -> Vector (Maybe a)
forall a. Int -> a -> Vector a
V.replicate (Vector (Maybe b) -> Int
forall a. Vector a -> Int
V.length Vector (Maybe b)
b Int -> Int -> Int
forall a. Num a => a -> a -> a
- Vector (Maybe a) -> Int
forall a. Vector a -> Int
V.length Vector (Maybe a)
a) Maybe a
forall a. Maybe a
Nothing, Vector (Maybe b)
b)
      go :: (Int, (Maybe a, Maybe b)) -> Maybe c
go (Int
i, (Just a
a', Just b
b')) = c -> Maybe c
forall a. a -> Maybe a
Just (c -> Maybe c) -> c -> Maybe c
forall a b. (a -> b) -> a -> b
$ Int -> a -> b -> c
f Int
i a
a' b
b'
      go (Int, (Maybe a, Maybe b))
_ = Maybe c
forall a. Maybe a
Nothing
   in Vector (Maybe c) -> SparseVector c
forall a. Vector (Maybe a) -> SparseVector a
SparseVector (Vector (Maybe c) -> SparseVector c)
-> (Vector (Maybe a, Maybe b) -> Vector (Maybe c))
-> Vector (Maybe a, Maybe b)
-> SparseVector c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, (Maybe a, Maybe b)) -> Maybe c)
-> Vector (Int, (Maybe a, Maybe b)) -> Vector (Maybe c)
forall a b. (a -> b) -> Vector a -> Vector b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int, (Maybe a, Maybe b)) -> Maybe c
go (Vector (Int, (Maybe a, Maybe b)) -> Vector (Maybe c))
-> (Vector (Maybe a, Maybe b) -> Vector (Int, (Maybe a, Maybe b)))
-> Vector (Maybe a, Maybe b)
-> Vector (Maybe c)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Maybe a, Maybe b) -> Vector (Int, (Maybe a, Maybe b))
forall a. Vector a -> Vector (Int, a)
V.indexed (Vector (Maybe a, Maybe b) -> SparseVector c)
-> Vector (Maybe a, Maybe b) -> SparseVector c
forall a b. (a -> b) -> a -> b
$ Vector (Maybe a) -> Vector (Maybe b) -> Vector (Maybe a, Maybe b)
forall a b. Vector a -> Vector b -> Vector (a, b)
V.zip Vector (Maybe a)
as Vector (Maybe b)
bs

intersectionVec :: SparseVector a -> SparseVector b -> Vector a
intersectionVec :: forall a b. SparseVector a -> SparseVector b -> Vector a
intersectionVec = (a -> b -> a) -> SparseVector a -> SparseVector b -> Vector a
forall a b c.
(a -> b -> c) -> SparseVector a -> SparseVector b -> Vector c
intersectionVecWith a -> b -> a
forall a b. a -> b -> a
const
{-# INLINE intersectionVec #-}

intersectionVecWith :: (a -> b -> c) -> SparseVector a -> SparseVector b -> Vector c
intersectionVecWith :: forall a b c.
(a -> b -> c) -> SparseVector a -> SparseVector b -> Vector c
intersectionVecWith = (Int -> a -> b -> c)
-> SparseVector a -> SparseVector b -> Vector c
forall a b c.
(Int -> a -> b -> c)
-> SparseVector a -> SparseVector b -> Vector c
intersectionVecWithKey ((Int -> a -> b -> c)
 -> SparseVector a -> SparseVector b -> Vector c)
-> ((a -> b -> c) -> Int -> a -> b -> c)
-> (a -> b -> c)
-> SparseVector a
-> SparseVector b
-> Vector c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b -> c) -> Int -> a -> b -> c
forall a b. a -> b -> a
const
{-# INLINE intersectionVecWith #-}

intersectionVecWithKey :: (Int -> a -> b -> c) -> SparseVector a -> SparseVector b -> Vector c
intersectionVecWithKey :: forall a b c.
(Int -> a -> b -> c)
-> SparseVector a -> SparseVector b -> Vector c
intersectionVecWithKey Int -> a -> b -> c
f (SparseVector Vector (Maybe a)
a) (SparseVector Vector (Maybe b)
b) = (Int -> (Maybe a, Maybe b) -> Maybe c)
-> Vector (Maybe a, Maybe b) -> Vector c
forall a b. (Int -> a -> Maybe b) -> Vector a -> Vector b
V.imapMaybe Int -> (Maybe a, Maybe b) -> Maybe c
go (Vector (Maybe a, Maybe b) -> Vector c)
-> Vector (Maybe a, Maybe b) -> Vector c
forall a b. (a -> b) -> a -> b
$ Vector (Maybe a) -> Vector (Maybe b) -> Vector (Maybe a, Maybe b)
forall a b. Vector a -> Vector b -> Vector (a, b)
V.zip Vector (Maybe a)
a Vector (Maybe b)
b
  where
    go :: Int -> (Maybe a, Maybe b) -> Maybe c
go Int
i (Just a
a', Just b
b') = c -> Maybe c
forall a. a -> Maybe a
Just (c -> Maybe c) -> c -> Maybe c
forall a b. (a -> b) -> a -> b
$ Int -> a -> b -> c
f Int
i a
a' b
b'
    go Int
_ (Maybe a, Maybe b)
_ = Maybe c
forall a. Maybe a
Nothing
{-# INLINE intersectionVecWithKey #-}

fromList :: [(Int, a)] -> SparseVector a
fromList :: forall a. [(Int, a)] -> SparseVector a
fromList = ((Int, a) -> SparseVector a -> SparseVector a)
-> SparseVector a -> [(Int, a)] -> SparseVector a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((Int -> a -> SparseVector a -> SparseVector a)
-> (Int, a) -> SparseVector a -> SparseVector a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Int -> a -> SparseVector a -> SparseVector a
forall a. Int -> a -> SparseVector a -> SparseVector a
insert) SparseVector a
forall a. SparseVector a
empty

toList :: SparseVector a -> [Maybe a]
toList :: forall a. SparseVector a -> [Maybe a]
toList (SparseVector Vector (Maybe a)
v) = Vector (Maybe a) -> [Maybe a]
forall a. Vector a -> [a]
V.toList Vector (Maybe a)
v
{-# INLINE toList #-}

fromVector :: Vector a -> SparseVector a
fromVector :: forall a. Vector a -> SparseVector a
fromVector = Vector (Maybe a) -> SparseVector a
forall a. Vector (Maybe a) -> SparseVector a
SparseVector (Vector (Maybe a) -> SparseVector a)
-> (Vector a -> Vector (Maybe a)) -> Vector a -> SparseVector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Maybe a) -> Vector a -> Vector (Maybe a)
forall a b. (a -> b) -> Vector a -> Vector b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Maybe a
forall a. a -> Maybe a
Just

toVector :: SparseVector a -> Vector a
toVector :: forall a. SparseVector a -> Vector a
toVector (SparseVector Vector (Maybe a)
v) = Vector (Maybe a) -> Vector a
forall a. Vector (Maybe a) -> Vector a
V.catMaybes Vector (Maybe a)
v

-- | Freeze a `MSparseVector` into a `SparseVector`.
freeze :: (PrimMonad m) => MSparseVector (PrimState m) a -> m (SparseVector a)
freeze :: forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> m (SparseVector a)
freeze (MSparseVector MVector (PrimState m) (Maybe a)
vec) = do
  Vector (Maybe a)
vec' <- 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)
vec
  SparseVector a -> m (SparseVector a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (SparseVector a -> m (SparseVector a))
-> SparseVector a -> m (SparseVector a)
forall a b. (a -> b) -> a -> b
$ Vector (Maybe a) -> SparseVector a
forall a. Vector (Maybe a) -> SparseVector a
SparseVector Vector (Maybe a)
vec'
{-# INLINE freeze #-}

-- | Unfreeze a `SparseVector` into a `MSparseVector`.
thaw :: (PrimMonad m) => SparseVector a -> m (MSparseVector (PrimState m) a)
thaw :: forall (m :: * -> *) a.
PrimMonad m =>
SparseVector a -> m (MSparseVector (PrimState m) a)
thaw (SparseVector Vector (Maybe a)
vec) = do
  MVector (PrimState m) (Maybe a)
vec' <- Vector (Maybe a) -> m (MVector (PrimState m) (Maybe a))
forall (m :: * -> *) a.
PrimMonad m =>
Vector a -> m (MVector (PrimState m) a)
V.thaw Vector (Maybe a)
vec
  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'
{-# INLINE thaw #-}