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