module Util.UnionFind(
UnionFind,
newElement,
toValue,
union,
isSame,
sameElements,
) where
import Data.IORef
import Util.Computation(done)
import Util.ExtendedPrelude
data UnionFind a = UnionFind {
value :: a,
contentsRef :: IORef [UnionFind a],
headRef :: IORef (Maybe (UnionFind a))
}
instance Eq (UnionFind a) where
(==) = mapEq contentsRef
newElement :: a -> IO (UnionFind a)
newElement value =
do
contentsRef <- newIORef []
headRef <- newIORef Nothing
let
unionFind = UnionFind {
value = value,
contentsRef = contentsRef,
headRef = headRef
}
return unionFind
toValue :: UnionFind a -> a
toValue = value
union :: UnionFind a -> UnionFind a -> IO ()
union uf1 uf2 =
do
head1 <- getHead uf1
head2 <- getHead uf2
if head1 == head2
then
done
else
do
writeIORef (headRef head2) (Just head1)
contents0 <- readIORef (contentsRef head1)
writeIORef (contentsRef head1) (head2 : contents0)
isSame :: UnionFind a -> UnionFind a -> IO Bool
isSame uf1 uf2 =
do
head1 <- getHead uf1
head2 <- getHead uf2
return (head1 == head2)
sameElements :: UnionFind a -> IO [UnionFind a]
sameElements uf =
do
head <- getHead uf
allContents head
where
allContents :: UnionFind a -> IO [UnionFind a]
allContents uf =
do
contents <- readIORef (contentsRef uf)
innerContents <- mapM allContents contents
return (uf : concat innerContents)
getHead :: UnionFind a -> IO (UnionFind a)
getHead unionFind =
do
thisHeadOpt <- readIORef (headRef unionFind)
case thisHeadOpt of
Nothing -> return unionFind
Just unionFind2 ->
do
thisHead <- getHead unionFind2
writeIORef (headRef unionFind) (Just thisHead)
return thisHead