{-# OPTIONS_GHC -funbox-strict-fields #-}
module Data.UnionFind.ST
( Point, fresh, repr, union, union', equivalent, redundant,
descriptor, setDescriptor, modifyDescriptor )
where
import Control.Applicative
import Control.Monad ( when )
import Control.Monad.ST
import Data.STRef
newtype Point s a = Pt (STRef s (Link s a)) deriving Point s a -> Point s a -> Bool
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
data Link s a
= Info {-# UNPACK #-} !(STRef s (Info a))
| Link {-# UNPACK #-} !(Point s a)
deriving Link s a -> Link s a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall s a. Link s a -> Link s a -> Bool
/= :: Link s a -> Link s a -> Bool
$c/= :: forall s a. Link s a -> Link s a -> Bool
== :: Link s a -> Link s a -> Bool
$c== :: forall s a. Link s a -> Link s a -> Bool
Eq
data Info a = MkInfo
{ forall a. Info a -> Int
weight :: {-# UNPACK #-} !Int
, forall a. Info a -> a
descr :: a
} deriving Info a -> Info a -> Bool
forall a. Eq a => Info a -> Info a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Info a -> Info a -> Bool
$c/= :: forall a. Eq a => Info a -> Info a -> Bool
== :: Info a -> Info a -> Bool
$c== :: forall a. Eq a => Info a -> Info a -> Bool
Eq
fresh :: a -> ST s (Point s a)
fresh :: forall a s. a -> ST s (Point s a)
fresh a
desc = do
STRef s (Info a)
info <- forall a s. a -> ST s (STRef s a)
newSTRef (MkInfo { weight :: Int
weight = Int
1, descr :: a
descr = a
desc })
STRef s (Link s a)
l <- forall a s. a -> ST s (STRef s a)
newSTRef (forall s a. STRef s (Info a) -> Link s a
Info STRef s (Info a)
info)
forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. STRef s (Link s a) -> Point s a
Pt STRef s (Link s a)
l)
repr :: Point s a -> ST s (Point s a)
repr :: forall s a. Point s a -> ST s (Point s a)
repr point :: Point s a
point@(Pt STRef s (Link s a)
l) = do
Link s a
link <- forall s a. STRef s a -> ST s a
readSTRef STRef s (Link s a)
l
case Link s a
link of
Info STRef s (Info a)
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return Point s a
point
Link pt' :: Point s a
pt'@(Pt STRef s (Link s a)
l') -> do
Point s a
pt'' <- forall s a. Point s a -> ST s (Point s a)
repr Point s a
pt'
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Point s a
pt'' forall a. Eq a => a -> a -> Bool
/= Point s a
pt') forall a b. (a -> b) -> a -> b
$ do
Link s a
link' <- forall s a. STRef s a -> ST s a
readSTRef STRef s (Link s a)
l'
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s (Link s a)
l Link s a
link'
forall (m :: * -> *) a. Monad m => a -> m a
return Point s a
pt''
descrRef :: Point s a -> ST s (STRef s (Info a))
descrRef :: forall s a. Point s a -> ST s (STRef s (Info a))
descrRef point :: Point s a
point@(Pt STRef s (Link s a)
link_ref) = do
Link s a
link <- forall s a. STRef s a -> ST s a
readSTRef STRef s (Link s a)
link_ref
case Link s a
link of
Info STRef s (Info a)
info -> forall (m :: * -> *) a. Monad m => a -> m a
return STRef s (Info a)
info
Link (Pt STRef s (Link s a)
link'_ref) -> do
Link s a
link' <- forall s a. STRef s a -> ST s a
readSTRef STRef s (Link s a)
link'_ref
case Link s a
link' of
Info STRef s (Info a)
info -> forall (m :: * -> *) a. Monad m => a -> m a
return STRef s (Info a)
info
Link s a
_ -> forall s a. Point s a -> ST s (STRef s (Info a))
descrRef forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall s a. Point s a -> ST s (Point s a)
repr Point s a
point
descriptor :: Point s a -> ST s a
descriptor :: forall s a. Point s a -> ST s a
descriptor Point s a
point = do
forall a. Info a -> a
descr forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall s a. STRef s a -> ST s a
readSTRef forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall s a. Point s a -> ST s (STRef s (Info a))
descrRef Point s a
point)
setDescriptor :: Point s a -> a -> ST s ()
setDescriptor :: forall s a. Point s a -> a -> ST s ()
setDescriptor Point s a
point a
new_descr = do
STRef s (Info a)
r <- forall s a. Point s a -> ST s (STRef s (Info a))
descrRef Point s a
point
forall s a. STRef s a -> (a -> a) -> ST s ()
modifySTRef STRef s (Info a)
r forall a b. (a -> b) -> a -> b
$ \Info a
i -> Info a
i { descr :: a
descr = a
new_descr }
modifyDescriptor :: Point s a -> (a -> a) -> ST s ()
modifyDescriptor :: forall s a. Point s a -> (a -> a) -> ST s ()
modifyDescriptor Point s a
point a -> a
f = do
STRef s (Info a)
r <- forall s a. Point s a -> ST s (STRef s (Info a))
descrRef Point s a
point
forall s a. STRef s a -> (a -> a) -> ST s ()
modifySTRef STRef s (Info a)
r forall a b. (a -> b) -> a -> b
$ \Info a
i -> Info a
i { descr :: a
descr = a -> a
f (forall a. Info a -> a
descr Info a
i) }
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
p1 Point s a
p2 = forall s a. Point s a -> Point s a -> (a -> a -> ST s a) -> ST s ()
union' Point s a
p1 Point s a
p2 (\a
_ a
d2 -> forall (m :: * -> *) a. Monad m => a -> m a
return a
d2)
union' :: Point s a -> Point s a -> (a -> a -> ST s a) -> ST s ()
union' :: forall s a. Point s a -> Point s a -> (a -> a -> ST s a) -> ST s ()
union' Point s a
p1 Point s a
p2 a -> a -> ST s a
update = do
point1 :: Point s a
point1@(Pt STRef s (Link s a)
link_ref1) <- forall s a. Point s a -> ST s (Point s a)
repr Point s a
p1
point2 :: Point s a
point2@(Pt STRef s (Link s a)
link_ref2) <- forall s a. Point s a -> ST s (Point s a)
repr Point s a
p2
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Point s a
point1 forall a. Eq a => a -> a -> Bool
/= Point s a
point2) forall a b. (a -> b) -> a -> b
$ do
Info STRef s (Info a)
info_ref1 <- forall s a. STRef s a -> ST s a
readSTRef STRef s (Link s a)
link_ref1
Info STRef s (Info a)
info_ref2 <- forall s a. STRef s a -> ST s a
readSTRef STRef s (Link s a)
link_ref2
MkInfo Int
w1 a
d1 <- forall s a. STRef s a -> ST s a
readSTRef STRef s (Info a)
info_ref1
MkInfo Int
w2 a
d2 <- forall s a. STRef s a -> ST s a
readSTRef STRef s (Info a)
info_ref2
a
d2' <- a -> a -> ST s a
update a
d1 a
d2
if Int
w1 forall a. Ord a => a -> a -> Bool
>= Int
w2 then do
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s (Link s a)
link_ref2 (forall s a. Point s a -> Link s a
Link Point s a
point1)
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s (Info a)
info_ref1 (forall a. Int -> a -> Info a
MkInfo (Int
w1 forall a. Num a => a -> a -> a
+ Int
w2) a
d2')
else do
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s (Link s a)
link_ref1 (forall s a. Point s a -> Link s a
Link Point s a
point2)
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef s (Info a)
info_ref2 (forall a. Int -> a -> Info a
MkInfo (Int
w1 forall a. Num a => a -> a -> a
+ Int
w2) a
d2')
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
p1 Point s a
p2 = forall a. Eq a => a -> a -> Bool
(==) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s a. Point s a -> ST s (Point s a)
repr Point s a
p1 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall s a. Point s a -> ST s (Point s a)
repr Point s a
p2
redundant :: Point s a -> ST s Bool
redundant :: forall s a. Point s a -> ST s Bool
redundant (Pt STRef s (Link s a)
link_r) = do
Link s a
link <- forall s a. STRef s a -> ST s a
readSTRef STRef s (Link s a)
link_r
case Link s a
link of
Info STRef s (Info a)
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
Link Point s a
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True