{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}

-- |

-- Module      : Data.SparseSet.Strict

-- 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.SparseSet.Strict
  ( -- * Sparse sets

    SparseSet (..),

    -- ** Construction

    empty,

    -- ** Operations

    insert,
    lookup,
    delete,

    -- ** Intersection

    intersection,
    intersectionWith,
    intersectionVec,

    -- ** Conversion

    toList,
    freeze,
    unsafeFreeze,
    thaw,
    unsafeThaw,
  )
where

import Control.DeepSeq
import Data.SparseSet.Strict.Mutable (MSparseSet (MSparseSet), PrimMonad (..))
import Data.SparseVector.Strict (SparseVector)
import qualified Data.SparseVector.Strict as SV
import Data.Vector.Strict (Vector)
import qualified Data.Vector.Strict as V
import qualified Data.Vector.Strict.Mutable as MV
import GHC.Generics (Generic)
import Prelude hiding (lookup)

data SparseSet i a = SparseSet
  { forall i a. SparseSet i a -> Vector a
dense :: {-# UNPACK #-} !(Vector a),
    forall i a. SparseSet i a -> SparseVector i
sparse :: !(SparseVector i)
  }
  deriving (Int -> SparseSet i a -> ShowS
[SparseSet i a] -> ShowS
SparseSet i a -> String
(Int -> SparseSet i a -> ShowS)
-> (SparseSet i a -> String)
-> ([SparseSet i a] -> ShowS)
-> Show (SparseSet i a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall i a. (Show a, Show i) => Int -> SparseSet i a -> ShowS
forall i a. (Show a, Show i) => [SparseSet i a] -> ShowS
forall i a. (Show a, Show i) => SparseSet i a -> String
$cshowsPrec :: forall i a. (Show a, Show i) => Int -> SparseSet i a -> ShowS
showsPrec :: Int -> SparseSet i a -> ShowS
$cshow :: forall i a. (Show a, Show i) => SparseSet i a -> String
show :: SparseSet i a -> String
$cshowList :: forall i a. (Show a, Show i) => [SparseSet i a] -> ShowS
showList :: [SparseSet i a] -> ShowS
Show, SparseSet i a -> SparseSet i a -> Bool
(SparseSet i a -> SparseSet i a -> Bool)
-> (SparseSet i a -> SparseSet i a -> Bool) -> Eq (SparseSet i a)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall i a. (Eq a, Eq i) => SparseSet i a -> SparseSet i a -> Bool
$c== :: forall i a. (Eq a, Eq i) => SparseSet i a -> SparseSet i a -> Bool
== :: SparseSet i a -> SparseSet i a -> Bool
$c/= :: forall i a. (Eq a, Eq i) => SparseSet i a -> SparseSet i a -> Bool
/= :: SparseSet i a -> SparseSet i a -> Bool
Eq, (forall x. SparseSet i a -> Rep (SparseSet i a) x)
-> (forall x. Rep (SparseSet i a) x -> SparseSet i a)
-> Generic (SparseSet i a)
forall x. Rep (SparseSet i a) x -> SparseSet i a
forall x. SparseSet i a -> Rep (SparseSet i a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall i a x. Rep (SparseSet i a) x -> SparseSet i a
forall i a x. SparseSet i a -> Rep (SparseSet i a) x
$cfrom :: forall i a x. SparseSet i a -> Rep (SparseSet i a) x
from :: forall x. SparseSet i a -> Rep (SparseSet i a) x
$cto :: forall i a x. Rep (SparseSet i a) x -> SparseSet i a
to :: forall x. Rep (SparseSet i a) x -> SparseSet i a
Generic, SparseSet i a -> ()
(SparseSet i a -> ()) -> NFData (SparseSet i a)
forall a. (a -> ()) -> NFData a
forall i a. (NFData a, NFData i) => SparseSet i a -> ()
$crnf :: forall i a. (NFData a, NFData i) => SparseSet i a -> ()
rnf :: SparseSet i a -> ()
NFData)

empty :: SparseSet i a
empty :: forall i a. SparseSet i a
empty = Vector a -> SparseVector i -> SparseSet i a
forall i a. Vector a -> SparseVector i -> SparseSet i a
SparseSet Vector a
forall a. Vector a
V.empty SparseVector i
forall a. SparseVector a
SV.empty
{-# INLINE empty #-}

insert :: (Integral i) => i -> a -> SparseSet i a -> SparseSet i a
insert :: forall i a. Integral i => i -> a -> SparseSet i a -> SparseSet i a
insert i
i a
a SparseSet i a
s =
  case Int -> SparseVector i -> Maybe i
forall a. Int -> SparseVector a -> Maybe a
SV.lookup (i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
i) (SparseVector i -> Maybe i) -> SparseVector i -> Maybe i
forall a b. (a -> b) -> a -> b
$ SparseSet i a -> SparseVector i
forall i a. SparseSet i a -> SparseVector i
sparse SparseSet i a
s of
    Just i
denseIndex ->
      SparseSet i a
s {dense = V.modify (\MVector s a
v -> MVector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write MVector s a
MVector (PrimState (ST s)) a
v (i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
denseIndex) a
a) (dense s)}
    Maybe i
Nothing ->
      SparseSet
        { dense :: Vector a
dense = Vector a -> a -> Vector a
forall a. Vector a -> a -> Vector a
V.snoc (SparseSet i a -> Vector a
forall i a. SparseSet i a -> Vector a
dense SparseSet i a
s) a
a,
          sparse :: SparseVector i
sparse = Int -> i -> SparseVector i -> SparseVector i
forall a. Int -> a -> SparseVector a -> SparseVector a
SV.insert (i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
i) (Int -> i
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> i) -> Int -> i
forall a b. (a -> b) -> a -> b
$ Vector a -> Int
forall a. Vector a -> Int
V.length (Vector a -> Int) -> Vector a -> Int
forall a b. (a -> b) -> a -> b
$ SparseSet i a -> Vector a
forall i a. SparseSet i a -> Vector a
dense SparseSet i a
s) (SparseSet i a -> SparseVector i
forall i a. SparseSet i a -> SparseVector i
sparse SparseSet i a
s)
        }
{-# INLINE insert #-}

lookup :: (Integral i) => SparseSet i a -> i -> Maybe a
lookup :: forall i a. Integral i => SparseSet i a -> i -> Maybe a
lookup SparseSet i a
s i
i =
  case Int -> SparseVector i -> Maybe i
forall a. Int -> SparseVector a -> Maybe a
SV.lookup (i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
i) (SparseVector i -> Maybe i) -> SparseVector i -> Maybe i
forall a b. (a -> b) -> a -> b
$ SparseSet i a -> SparseVector i
forall i a. SparseSet i a -> SparseVector i
sparse SparseSet i a
s of
    Just i
denseIndex -> a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> (Int -> a) -> Int -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.unsafeIndex (SparseSet i a -> Vector a
forall i a. SparseSet i a -> Vector a
dense SparseSet i a
s) (Int -> Maybe a) -> Int -> Maybe a
forall a b. (a -> b) -> a -> b
$ i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
denseIndex
    Maybe i
Nothing -> Maybe a
forall a. Maybe a
Nothing
{-# INLINE lookup #-}

delete :: (Integral i) => i -> SparseSet i a -> SparseSet i a
delete :: forall i a. Integral i => i -> SparseSet i a -> SparseSet i a
delete i
i SparseSet i a
s =
  case Int -> SparseVector i -> Maybe i
forall a. Int -> SparseVector a -> Maybe a
SV.lookup (i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
i) (SparseVector i -> Maybe i) -> SparseVector i -> Maybe i
forall a b. (a -> b) -> a -> b
$ SparseSet i a -> SparseVector i
forall i a. SparseSet i a -> SparseVector i
sparse SparseSet i a
s of
    Just i
denseIndex ->
      SparseSet
        { dense :: Vector a
dense =
            Int -> Vector a -> Vector a
forall a. Int -> Vector a -> Vector a
V.take (i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
denseIndex) (SparseSet i a -> Vector a
forall i a. SparseSet i a -> Vector a
dense SparseSet i a
s)
              Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
V.++ Int -> Vector a -> Vector a
forall a. Int -> Vector a -> Vector a
V.drop (i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
denseIndex Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (SparseSet i a -> Vector a
forall i a. SparseSet i a -> Vector a
dense SparseSet i a
s),
          sparse :: SparseVector i
sparse = Int -> SparseVector i -> SparseVector i
forall a. Int -> SparseVector a -> SparseVector a
SV.delete (i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
i) (SparseSet i a -> SparseVector i
forall i a. SparseSet i a -> SparseVector i
sparse SparseSet i a
s)
        }
    Maybe i
Nothing -> SparseSet i a
s
{-# INLINE delete #-}

intersection ::
  (Integral i) =>
  SparseSet i a ->
  SparseSet i b ->
  SparseSet i a
intersection :: forall i a b.
Integral i =>
SparseSet i a -> SparseSet i b -> SparseSet i a
intersection SparseSet i a
as SparseSet i b
bs =
  let x :: SparseVector i
x = SparseVector i -> SparseVector i -> SparseVector i
forall a b. SparseVector a -> SparseVector b -> SparseVector a
SV.intersection (SparseSet i a -> SparseVector i
forall i a. SparseSet i a -> SparseVector i
sparse SparseSet i a
as) (SparseSet i b -> SparseVector i
forall i a. SparseSet i a -> SparseVector i
sparse SparseSet i b
bs)
      (i
_, SparseVector i
x') = (i -> i -> (i, i)) -> i -> SparseVector i -> (i, SparseVector i)
forall a b c.
(a -> b -> (a, c)) -> a -> SparseVector b -> (a, SparseVector c)
SV.mapAccum (\i
i i
_ -> (i
i i -> i -> i
forall a. Num a => a -> a -> a
+ i
1, i
i)) i
0 SparseVector i
x
      as' :: Vector a
as' = (i -> a) -> Vector i -> Vector a
forall a b. (a -> b) -> Vector a -> Vector b
V.map (Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.unsafeIndex (SparseSet i a -> Vector a
forall i a. SparseSet i a -> Vector a
dense SparseSet i a
as) (Int -> a) -> (i -> Int) -> i -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral) (SparseVector i -> Vector i
forall a. SparseVector a -> Vector a
SV.toVector SparseVector i
x)
   in SparseSet {dense :: Vector a
dense = Vector a
as', sparse :: SparseVector i
sparse = SparseVector i
x'}

intersectionVec ::
  (Integral i) =>
  SparseSet i a ->
  SparseSet i b ->
  Vector a
intersectionVec :: forall i a b.
Integral i =>
SparseSet i a -> SparseSet i b -> Vector a
intersectionVec SparseSet i a
as SparseSet i b
bs = (i -> i -> a) -> SparseVector i -> SparseVector i -> Vector a
forall a b c.
(a -> b -> c) -> SparseVector a -> SparseVector b -> Vector c
SV.intersectionVecWith i -> i -> a
forall {a} {p}. Integral a => a -> p -> a
go (SparseSet i a -> SparseVector i
forall i a. SparseSet i a -> SparseVector i
sparse SparseSet i a
as) (SparseVector i -> Vector a) -> SparseVector i -> Vector a
forall a b. (a -> b) -> a -> b
$ SparseSet i b -> SparseVector i
forall i a. SparseSet i a -> SparseVector i
sparse SparseSet i b
bs
  where
    go :: a -> p -> a
go a
a p
_ = Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.unsafeIndex (SparseSet i a -> Vector a
forall i a. SparseSet i a -> Vector a
dense SparseSet i a
as) (Int -> a) -> Int -> a
forall a b. (a -> b) -> a -> b
$ a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
a
{-# INLINE intersectionVec #-}

intersectionWith ::
  (Integral i) =>
  (a -> b -> c) ->
  SparseSet i a ->
  SparseSet i b ->
  SparseSet i c
intersectionWith :: forall i a b c.
Integral i =>
(a -> b -> c) -> SparseSet i a -> SparseSet i b -> SparseSet i c
intersectionWith a -> b -> c
f SparseSet i a
as SparseSet i b
bs =
  let x :: SparseVector i
x = SparseVector i -> SparseVector i -> SparseVector i
forall a b. SparseVector a -> SparseVector b -> SparseVector a
SV.intersection (SparseSet i a -> SparseVector i
forall i a. SparseSet i a -> SparseVector i
sparse SparseSet i a
as) (SparseSet i b -> SparseVector i
forall i a. SparseSet i a -> SparseVector i
sparse SparseSet i b
bs)
      (i
_, SparseVector i
x') = (i -> i -> (i, i)) -> i -> SparseVector i -> (i, SparseVector i)
forall a b c.
(a -> b -> (a, c)) -> a -> SparseVector b -> (a, SparseVector c)
SV.mapAccum (\i
i i
_ -> (i
i i -> i -> i
forall a. Num a => a -> a -> a
+ i
1, i
i)) i
0 SparseVector i
x
      as' :: Vector a
as' = (i -> a) -> Vector i -> Vector a
forall a b. (a -> b) -> Vector a -> Vector b
V.map (\i
i -> SparseSet i a -> Vector a
forall i a. SparseSet i a -> Vector a
dense SparseSet i a
as Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.! i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
i) (SparseVector i -> Vector i
forall a. SparseVector a -> Vector a
SV.toVector SparseVector i
x)
      bs' :: Vector b
bs' = (i -> b) -> Vector i -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
V.map (\i
i -> SparseSet i b -> Vector b
forall i a. SparseSet i a -> Vector a
dense SparseSet i b
bs Vector b -> Int -> b
forall a. Vector a -> Int -> a
V.! i -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral i
i) (SparseVector i -> Vector i
forall a. SparseVector a -> Vector a
SV.toVector SparseVector i
x)
      cs :: Vector c
cs = (a -> b -> c) -> Vector a -> Vector b -> Vector c
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith a -> b -> c
f Vector a
as' Vector b
bs'
   in SparseSet {dense :: Vector c
dense = Vector c
cs, sparse :: SparseVector i
sparse = SparseVector i
x'}
{-# INLINE intersectionWith #-}

toList :: (Integral i) => SparseSet i a -> [Maybe a]
toList :: forall i a. Integral i => SparseSet i a -> [Maybe a]
toList SparseSet i a
s = (Maybe i -> Maybe a) -> [Maybe i] -> [Maybe a]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Maybe i -> Maybe a
forall {a}. Integral a => Maybe a -> Maybe a
go ([Maybe i] -> [Maybe a]) -> [Maybe i] -> [Maybe a]
forall a b. (a -> b) -> a -> b
$ SparseVector i -> [Maybe i]
forall a. SparseVector a -> [Maybe a]
SV.toList (SparseVector i -> [Maybe i]) -> SparseVector i -> [Maybe i]
forall a b. (a -> b) -> a -> b
$ SparseSet i a -> SparseVector i
forall i a. SparseSet i a -> SparseVector i
sparse SparseSet i a
s
  where
    go :: Maybe a -> Maybe a
go (Just a
i) = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.unsafeIndex (SparseSet i a -> Vector a
forall i a. SparseSet i a -> Vector a
dense SparseSet i a
s) (a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
i)
    go Maybe a
Nothing = Maybe a
forall a. Maybe a
Nothing
{-# INLINE toList #-}

-- | Freeze a `MSparseSet` into a `SparseSet`.

freeze :: (PrimMonad m) => MSparseSet (PrimState m) i a -> m (SparseSet i a)
freeze :: forall (m :: * -> *) i a.
PrimMonad m =>
MSparseSet (PrimState m) i a -> m (SparseSet i a)
freeze (MSparseSet MVector (PrimState m) a
d MSparseVector (PrimState m) i
s) = do
  Vector a
d' <- MVector (PrimState m) a -> m (Vector a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.freeze MVector (PrimState m) a
d
  SparseVector i
s' <- MSparseVector (PrimState m) i -> m (SparseVector i)
forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> m (SparseVector a)
SV.freeze MSparseVector (PrimState m) i
s
  SparseSet i a -> m (SparseSet i a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (SparseSet i a -> m (SparseSet i a))
-> SparseSet i a -> m (SparseSet i a)
forall a b. (a -> b) -> a -> b
$ Vector a -> SparseVector i -> SparseSet i a
forall i a. Vector a -> SparseVector i -> SparseSet i a
SparseSet Vector a
d' SparseVector i
s'
{-# INLINE freeze #-}

unsafeFreeze :: (PrimMonad m) => MSparseSet (PrimState m) i a -> m (SparseSet i a)
unsafeFreeze :: forall (m :: * -> *) i a.
PrimMonad m =>
MSparseSet (PrimState m) i a -> m (SparseSet i a)
unsafeFreeze (MSparseSet MVector (PrimState m) a
d MSparseVector (PrimState m) i
s) = do
  Vector a
d' <- MVector (PrimState m) a -> m (Vector a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.unsafeFreeze MVector (PrimState m) a
d
  SparseVector i
s' <- MSparseVector (PrimState m) i -> m (SparseVector i)
forall (m :: * -> *) a.
PrimMonad m =>
MSparseVector (PrimState m) a -> m (SparseVector a)
SV.unsafeFreeze MSparseVector (PrimState m) i
s
  SparseSet i a -> m (SparseSet i a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (SparseSet i a -> m (SparseSet i a))
-> SparseSet i a -> m (SparseSet i a)
forall a b. (a -> b) -> a -> b
$ Vector a -> SparseVector i -> SparseSet i a
forall i a. Vector a -> SparseVector i -> SparseSet i a
SparseSet Vector a
d' SparseVector i
s'
{-# INLINE unsafeFreeze #-}

-- | Unfreeze a `SparseSet` into a `MSparseSet`.

thaw :: (PrimMonad m) => SparseSet i a -> m (MSparseSet (PrimState m) i a)
thaw :: forall (m :: * -> *) i a.
PrimMonad m =>
SparseSet i a -> m (MSparseSet (PrimState m) i a)
thaw (SparseSet Vector a
d SparseVector i
s) = do
  !MVector (PrimState m) a
d' <- Vector a -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Vector a -> m (MVector (PrimState m) a)
V.thaw Vector a
d
  !MSparseVector (PrimState m) i
s' <- SparseVector i -> m (MSparseVector (PrimState m) i)
forall (m :: * -> *) a.
PrimMonad m =>
SparseVector a -> m (MSparseVector (PrimState m) a)
SV.thaw SparseVector i
s
  MSparseSet (PrimState m) i a -> m (MSparseSet (PrimState m) i a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (MSparseSet (PrimState m) i a -> m (MSparseSet (PrimState m) i a))
-> MSparseSet (PrimState m) i a -> m (MSparseSet (PrimState m) i a)
forall a b. (a -> b) -> a -> b
$ MVector (PrimState m) a
-> MSparseVector (PrimState m) i -> MSparseSet (PrimState m) i a
forall s i a. MVector s a -> MSparseVector s i -> MSparseSet s i a
MSparseSet MVector (PrimState m) a
d' MSparseVector (PrimState m) i
s'
{-# INLINE thaw #-}

unsafeThaw :: (PrimMonad m) => SparseSet i a -> m (MSparseSet (PrimState m) i a)
unsafeThaw :: forall (m :: * -> *) i a.
PrimMonad m =>
SparseSet i a -> m (MSparseSet (PrimState m) i a)
unsafeThaw (SparseSet Vector a
d SparseVector i
s) = do
  !MVector (PrimState m) a
d' <- Vector a -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Vector a -> m (MVector (PrimState m) a)
V.unsafeThaw Vector a
d
  !MSparseVector (PrimState m) i
s' <- SparseVector i -> m (MSparseVector (PrimState m) i)
forall (m :: * -> *) a.
PrimMonad m =>
SparseVector a -> m (MSparseVector (PrimState m) a)
SV.unsafeThaw SparseVector i
s
  MSparseSet (PrimState m) i a -> m (MSparseSet (PrimState m) i a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (MSparseSet (PrimState m) i a -> m (MSparseSet (PrimState m) i a))
-> MSparseSet (PrimState m) i a -> m (MSparseSet (PrimState m) i a)
forall a b. (a -> b) -> a -> b
$ MVector (PrimState m) a
-> MSparseVector (PrimState m) i -> MSparseSet (PrimState m) i a
forall s i a. MVector s a -> MSparseVector s i -> MSparseSet s i a
MSparseSet MVector (PrimState m) a
d' MSparseVector (PrimState m) i
s'
{-# INLINE unsafeThaw #-}