module GHC.Data.UnionFind where
import GHC.Prelude
import Data.STRef
import Control.Monad.ST
import Control.Monad
newtype Point s a = Point (STRef s (Link s a))
deriving (Point s a -> Point s a -> Bool
(Point s a -> Point s a -> Bool)
-> (Point s a -> Point s a -> Bool) -> Eq (Point s a)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall s a. Point s a -> Point s a -> Bool
/= :: Point s a -> Point s a -> Bool
$c/= :: forall s a. Point s a -> Point s a -> Bool
== :: Point s a -> Point s a -> Bool
$c== :: forall s a. Point s a -> Point s a -> Bool
Eq)
writePoint :: Point s a -> Link s a -> ST s ()
writePoint :: forall s a. Point s a -> Link s a -> ST s ()
writePoint (Point STRef s (Link s a)
v) = STRef s (Link s a) -> Link s a -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s (Link s a)
v
readPoint :: Point s a -> ST s (Link s a)
readPoint :: forall s a. Point s a -> ST s (Link s a)
readPoint (Point STRef s (Link s a)
v) = STRef s (Link s a) -> ST s (Link s a)
forall s a. STRef s a -> ST s a
readSTRef STRef s (Link s a)
v
data Link s a
= Info {-# UNPACK #-} !(STRef s Int) {-# UNPACK #-} !(STRef s a)
| Link {-# UNPACK #-} !(Point s a)
fresh :: a -> ST s (Point s a)
fresh :: forall a s. a -> ST s (Point s a)
fresh a
desc = do
STRef s Int
weight <- Int -> ST s (STRef s Int)
forall a s. a -> ST s (STRef s a)
newSTRef Int
1
STRef s a
descriptor <- a -> ST s (STRef s a)
forall a s. a -> ST s (STRef s a)
newSTRef a
desc
STRef s (Link s a) -> Point s a
forall s a. STRef s (Link s a) -> Point s a
Point (STRef s (Link s a) -> Point s a)
-> ST s (STRef s (Link s a)) -> ST s (Point s a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` Link s a -> ST s (STRef s (Link s a))
forall a s. a -> ST s (STRef s a)
newSTRef (STRef s Int -> STRef s a -> Link s a
forall s a. STRef s Int -> STRef s a -> Link s a
Info STRef s Int
weight STRef s a
descriptor)
repr :: Point s a -> ST s (Point s a)
repr :: forall s a. Point s a -> ST s (Point s a)
repr Point s a
point = Point s a -> ST s (Link s a)
forall s a. Point s a -> ST s (Link s a)
readPoint Point s a
point ST s (Link s a)
-> (Link s a -> ST s (Point s a)) -> ST s (Point s a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Link s a
r ->
case Link s a
r of
Link Point s a
point' -> do
Point s a
point'' <- Point s a -> ST s (Point s a)
forall s a. Point s a -> ST s (Point s a)
repr Point s a
point'
Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Point s a
point'' Point s a -> Point s a -> Bool
forall a. Eq a => a -> a -> Bool
/= Point s a
point') (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
Point s a -> Link s a -> ST s ()
forall s a. Point s a -> Link s a -> ST s ()
writePoint Point s a
point (Link s a -> ST s ()) -> ST s (Link s a) -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Point s a -> ST s (Link s a)
forall s a. Point s a -> ST s (Link s a)
readPoint Point s a
point'
Point s a -> ST s (Point s a)
forall (m :: * -> *) a. Monad m => a -> m a
return Point s a
point''
Info STRef s Int
_ STRef s a
_ -> Point s a -> ST s (Point s a)
forall (m :: * -> *) a. Monad m => a -> m a
return Point s a
point
find :: Point s a -> ST s a
find :: forall s a. Point s a -> ST s a
find Point s a
point =
Point s a -> ST s (Link s a)
forall s a. Point s a -> ST s (Link s a)
readPoint Point s a
point ST s (Link s a) -> (Link s a -> ST s a) -> ST s a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Link s a
r ->
case Link s a
r of
Info STRef s Int
_ STRef s a
d_ref -> STRef s a -> ST s a
forall s a. STRef s a -> ST s a
readSTRef STRef s a
d_ref
Link Point s a
point' -> Point s a -> ST s (Link s a)
forall s a. Point s a -> ST s (Link s a)
readPoint Point s a
point' ST s (Link s a) -> (Link s a -> ST s a) -> ST s a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Link s a
r' ->
case Link s a
r' of
Info STRef s Int
_ STRef s a
d_ref -> STRef s a -> ST s a
forall s a. STRef s a -> ST s a
readSTRef STRef s a
d_ref
Link Point s a
_ -> Point s a -> ST s (Point s a)
forall s a. Point s a -> ST s (Point s a)
repr Point s a
point ST s (Point s a) -> (Point s a -> ST s a) -> ST s a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Point s a -> ST s a
forall s a. Point s a -> ST s a
find
union :: Point s a -> Point s a -> ST s ()
union :: forall s a. Point s a -> Point s a -> ST s ()
union Point s a
refpoint1 Point s a
refpoint2 = do
Point s a
point1 <- Point s a -> ST s (Point s a)
forall s a. Point s a -> ST s (Point s a)
repr Point s a
refpoint1
Point s a
point2 <- Point s a -> ST s (Point s a)
forall s a. Point s a -> ST s (Point s a)
repr Point s a
refpoint2
Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Point s a
point1 Point s a -> Point s a -> Bool
forall a. Eq a => a -> a -> Bool
/= Point s a
point2) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
Link s a
l1 <- Point s a -> ST s (Link s a)
forall s a. Point s a -> ST s (Link s a)
readPoint Point s a
point1
Link s a
l2 <- Point s a -> ST s (Link s a)
forall s a. Point s a -> ST s (Link s a)
readPoint Point s a
point2
case (Link s a
l1, Link s a
l2) of
(Info STRef s Int
wref1 STRef s a
dref1, Info STRef s Int
wref2 STRef s a
dref2) -> do
Int
weight1 <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef STRef s Int
wref1
Int
weight2 <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef STRef s Int
wref2
if Int
weight1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
weight2
then do
Point s a -> Link s a -> ST s ()
forall s a. Point s a -> Link s a -> ST s ()
writePoint Point s a
point2 (Point s a -> Link s a
forall s a. Point s a -> Link s a
Link Point s a
point1)
STRef s Int -> Int -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s Int
wref1 (Int
weight1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
weight2)
STRef s a -> a -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s a
dref1 (a -> ST s ()) -> ST s a -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< STRef s a -> ST s a
forall s a. STRef s a -> ST s a
readSTRef STRef s a
dref2
else do
Point s a -> Link s a -> ST s ()
forall s a. Point s a -> Link s a -> ST s ()
writePoint Point s a
point1 (Point s a -> Link s a
forall s a. Point s a -> Link s a
Link Point s a
point2)
STRef s Int -> Int -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s Int
wref2 (Int
weight1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
weight2)
(Link s a, Link s a)
_ -> [Char] -> ST s ()
forall a. HasCallStack => [Char] -> a
error [Char]
"UnionFind.union: repr invariant broken"
equivalent :: Point s a -> Point s a -> ST s Bool
equivalent :: forall s a. Point s a -> Point s a -> ST s Bool
equivalent Point s a
point1 Point s a
point2 = (Point s a -> Point s a -> Bool)
-> ST s (Point s a) -> ST s (Point s a) -> ST s Bool
forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 Point s a -> Point s a -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Point s a -> ST s (Point s a)
forall s a. Point s a -> ST s (Point s a)
repr Point s a
point1) (Point s a -> ST s (Point s a)
forall s a. Point s a -> ST s (Point s a)
repr Point s a
point2)