{-# LANGUAGE RecordWildCards #-}
module PureSAT.SparseSet (
SparseSet (..),
sizeofSparseSet,
indexSparseSet,
newSparseSet,
memberSparseSet,
insertSparseSet,
deleteSparseSet,
popSparseSet,
popSparseSet_,
elemsSparseSet,
clearSparseSet,
) where
import Data.Primitive.PrimVar
import PureSAT.Prim
import PureSAT.Base
data SparseSet s = SS
{ forall s. SparseSet s -> PrimVar s Int
size :: {-# UNPACK #-} !(PrimVar s Int)
, forall s. SparseSet s -> MutablePrimArray s Int
dense :: {-# UNPACK #-} !(MutablePrimArray s Int)
, forall s. SparseSet s -> MutablePrimArray s Int
sparse :: {-# UNPACK #-} !(MutablePrimArray s Int)
}
_invariant :: SparseSet s -> ST s ()
_invariant :: forall s. SparseSet s -> ST s ()
_invariant SS {PrimVar s Int
MutablePrimArray s Int
size :: forall s. SparseSet s -> PrimVar s Int
dense :: forall s. SparseSet s -> MutablePrimArray s Int
sparse :: forall s. SparseSet s -> MutablePrimArray s Int
size :: PrimVar s Int
dense :: MutablePrimArray s Int
sparse :: MutablePrimArray s Int
..} = do
n <- PrimVar (PrimState (ST s)) Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
size
capacity <- getSizeofMutablePrimArray dense
capacity' <- getSizeofMutablePrimArray sparse
assertST "capacities" (n <= capacity && capacity == capacity')
go capacity n 0
where
go :: Int -> Int -> Int -> ST s ()
go Int
capacity Int
n Int
i =
if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n
then () -> ST s ()
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
else do
x <- MutablePrimArray s Int -> Int -> ST s Int
forall a s.
(HasCallStack, Prim a) =>
MutablePrimArray s a -> Int -> ST s a
readPrimArray MutablePrimArray s Int
dense Int
i
assertST "x < capacity" $ x < capacity
j <- readPrimArray sparse x
assertST "i == j" $ i == j
go capacity n (i + 1)
checkInvariant :: SparseSet s -> ST s ()
checkInvariant :: forall s. SparseSet s -> ST s ()
checkInvariant SparseSet s
_ = () -> ST s ()
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
newSparseSet
:: Int
-> ST s (SparseSet s)
newSparseSet :: forall s. Int -> ST s (SparseSet s)
newSparseSet Int
capacity = do
size <- Int -> ST s (PrimVar (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
a -> m (PrimVar (PrimState m) a)
newPrimVar Int
0
dense <- newPrimArray capacity
sparse <- newPrimArray capacity
return SS {..}
indexSparseSet :: SparseSet s -> Int -> ST s Int
indexSparseSet :: forall s. SparseSet s -> Int -> ST s Int
indexSparseSet SS {PrimVar s Int
MutablePrimArray s Int
size :: forall s. SparseSet s -> PrimVar s Int
dense :: forall s. SparseSet s -> MutablePrimArray s Int
sparse :: forall s. SparseSet s -> MutablePrimArray s Int
size :: PrimVar s Int
dense :: MutablePrimArray s Int
sparse :: MutablePrimArray s Int
..} Int
i = MutablePrimArray s Int -> Int -> ST s Int
forall a s.
(HasCallStack, Prim a) =>
MutablePrimArray s a -> Int -> ST s a
readPrimArray MutablePrimArray s Int
dense Int
i
sizeofSparseSet :: SparseSet s -> ST s Int
sizeofSparseSet :: forall s. SparseSet s -> ST s Int
sizeofSparseSet SS {PrimVar s Int
MutablePrimArray s Int
size :: forall s. SparseSet s -> PrimVar s Int
dense :: forall s. SparseSet s -> MutablePrimArray s Int
sparse :: forall s. SparseSet s -> MutablePrimArray s Int
size :: PrimVar s Int
dense :: MutablePrimArray s Int
sparse :: MutablePrimArray s Int
..} = PrimVar (PrimState (ST s)) Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
size
memberSparseSet :: SparseSet s -> Int -> ST s Bool
memberSparseSet :: forall s. SparseSet s -> Int -> ST s Bool
memberSparseSet set :: SparseSet s
set@SS {PrimVar s Int
MutablePrimArray s Int
size :: forall s. SparseSet s -> PrimVar s Int
dense :: forall s. SparseSet s -> MutablePrimArray s Int
sparse :: forall s. SparseSet s -> MutablePrimArray s Int
size :: PrimVar s Int
dense :: MutablePrimArray s Int
sparse :: MutablePrimArray s Int
..} Int
x = do
SparseSet s -> ST s ()
forall s. SparseSet s -> ST s ()
checkInvariant SparseSet s
set
n <- PrimVar (PrimState (ST s)) Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
size
i <- readPrimArray sparse x
if 0 <= i && i < n
then do
x' <- readPrimArray dense i
return (x' == x)
else return False
insertSparseSet :: SparseSet s -> Int -> ST s ()
insertSparseSet :: forall s. SparseSet s -> Int -> ST s ()
insertSparseSet set :: SparseSet s
set@SS {PrimVar s Int
MutablePrimArray s Int
size :: forall s. SparseSet s -> PrimVar s Int
dense :: forall s. SparseSet s -> MutablePrimArray s Int
sparse :: forall s. SparseSet s -> MutablePrimArray s Int
size :: PrimVar s Int
dense :: MutablePrimArray s Int
sparse :: MutablePrimArray s Int
..} Int
x = do
SparseSet s -> ST s ()
forall s. SparseSet s -> ST s ()
checkInvariant SparseSet s
set
n <- PrimVar (PrimState (ST s)) Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
size
i <- readPrimArray sparse x
if 0 <= i && i < n
then do
x' <- readPrimArray dense i
if x == x' then return () else insert n
else insert n
where
{-# INLINE insert #-}
insert :: Int -> ST s ()
insert Int
n = do
MutablePrimArray s Int -> Int -> Int -> ST s ()
forall a s.
(HasCallStack, Prim a) =>
MutablePrimArray s a -> Int -> a -> ST s ()
writePrimArray MutablePrimArray s Int
dense Int
n Int
x
MutablePrimArray s Int -> Int -> Int -> ST s ()
forall a s.
(HasCallStack, Prim a) =>
MutablePrimArray s a -> Int -> a -> ST s ()
writePrimArray MutablePrimArray s Int
sparse Int
x Int
n
PrimVar (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> a -> m ()
writePrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
size (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
deleteSparseSet :: SparseSet s -> Int -> ST s ()
deleteSparseSet :: forall s. SparseSet s -> Int -> ST s ()
deleteSparseSet set :: SparseSet s
set@SS {PrimVar s Int
MutablePrimArray s Int
size :: forall s. SparseSet s -> PrimVar s Int
dense :: forall s. SparseSet s -> MutablePrimArray s Int
sparse :: forall s. SparseSet s -> MutablePrimArray s Int
size :: PrimVar s Int
dense :: MutablePrimArray s Int
sparse :: MutablePrimArray s Int
..} Int
x = do
SparseSet s -> ST s ()
forall s. SparseSet s -> ST s ()
checkInvariant SparseSet s
set
n <- PrimVar (PrimState (ST s)) Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
size
i <- readPrimArray sparse x
if 0 <= i && i < n
then do
x' <- readPrimArray dense i
if x == x' then delete i n else return ()
else return ()
where
{-# INLINE delete #-}
delete :: Int -> Int -> ST s ()
delete Int
i Int
n = do
PrimVar (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> a -> m ()
writePrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
size (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
MutablePrimArray s Int
-> MutablePrimArray s Int -> Int -> Int -> Int -> ST s ()
forall s.
MutablePrimArray s Int
-> MutablePrimArray s Int -> Int -> Int -> Int -> ST s ()
swap MutablePrimArray s Int
dense MutablePrimArray s Int
sparse Int
i Int
x (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
{-# INLINE swap #-}
swap :: MutablePrimArray s Int -> MutablePrimArray s Int -> Int -> Int -> Int -> ST s ()
swap :: forall s.
MutablePrimArray s Int
-> MutablePrimArray s Int -> Int -> Int -> Int -> ST s ()
swap !MutablePrimArray s Int
dense !MutablePrimArray s Int
sparse !Int
i !Int
x !Int
j
| Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j
= () -> ST s ()
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
| Bool
otherwise = do
y <- MutablePrimArray s Int -> Int -> ST s Int
forall a s.
(HasCallStack, Prim a) =>
MutablePrimArray s a -> Int -> ST s a
readPrimArray MutablePrimArray s Int
dense Int
j
writePrimArray dense j x
writePrimArray dense i y
writePrimArray sparse x j
writePrimArray sparse y i
popSparseSet :: SparseSet s -> ST s (Maybe Int)
popSparseSet :: forall s. SparseSet s -> ST s (Maybe Int)
popSparseSet SparseSet s
set = SparseSet s
-> ST s (Maybe Int)
-> (Int -> ST s (Maybe Int))
-> ST s (Maybe Int)
forall s r. SparseSet s -> ST s r -> (Int -> ST s r) -> ST s r
popSparseSet_ SparseSet s
set (Maybe Int -> ST s (Maybe Int)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Int
forall a. Maybe a
Nothing) (Maybe Int -> ST s (Maybe Int)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Int -> ST s (Maybe Int))
-> (Int -> Maybe Int) -> Int -> ST s (Maybe Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Maybe Int
forall a. a -> Maybe a
Just)
{-# INLINE popSparseSet_ #-}
popSparseSet_ :: SparseSet s -> ST s r -> (Int -> ST s r) -> ST s r
popSparseSet_ :: forall s r. SparseSet s -> ST s r -> (Int -> ST s r) -> ST s r
popSparseSet_ set :: SparseSet s
set@SS {PrimVar s Int
MutablePrimArray s Int
size :: forall s. SparseSet s -> PrimVar s Int
dense :: forall s. SparseSet s -> MutablePrimArray s Int
sparse :: forall s. SparseSet s -> MutablePrimArray s Int
size :: PrimVar s Int
dense :: MutablePrimArray s Int
sparse :: MutablePrimArray s Int
..} ST s r
no Int -> ST s r
yes = do
SparseSet s -> ST s ()
forall s. SparseSet s -> ST s ()
checkInvariant SparseSet s
set
n <- PrimVar (PrimState (ST s)) Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
size
if n <= 0
then no
else do
let !n' = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
x <- readPrimArray dense n'
writePrimVar size n'
yes x
clearSparseSet :: SparseSet s -> ST s ()
clearSparseSet :: forall s. SparseSet s -> ST s ()
clearSparseSet SS {PrimVar s Int
MutablePrimArray s Int
size :: forall s. SparseSet s -> PrimVar s Int
dense :: forall s. SparseSet s -> MutablePrimArray s Int
sparse :: forall s. SparseSet s -> MutablePrimArray s Int
size :: PrimVar s Int
dense :: MutablePrimArray s Int
sparse :: MutablePrimArray s Int
..} = do
PrimVar (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> a -> m ()
writePrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
size Int
0
elemsSparseSet :: SparseSet s -> ST s [Int]
elemsSparseSet :: forall s. SparseSet s -> ST s [Int]
elemsSparseSet SS {PrimVar s Int
MutablePrimArray s Int
size :: forall s. SparseSet s -> PrimVar s Int
dense :: forall s. SparseSet s -> MutablePrimArray s Int
sparse :: forall s. SparseSet s -> MutablePrimArray s Int
size :: PrimVar s Int
dense :: MutablePrimArray s Int
sparse :: MutablePrimArray s Int
..} = do
n <- PrimVar (PrimState (ST s)) Int -> ST s Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
PrimVar (PrimState m) a -> m a
readPrimVar PrimVar s Int
PrimVar (PrimState (ST s)) Int
size
go [] 0 n
where
go :: [Int] -> Int -> Int -> ST s [Int]
go ![Int]
acc !Int
i !Int
n
| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
= do
x <- MutablePrimArray s Int -> Int -> ST s Int
forall a s.
(HasCallStack, Prim a) =>
MutablePrimArray s a -> Int -> ST s a
readPrimArray MutablePrimArray s Int
dense Int
i
go (x : acc) (i + 1) n
| Bool
otherwise
= [Int] -> ST s [Int]
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Int] -> [Int]
forall a. [a] -> [a]
reverse [Int]
acc)