-- | Union-Find algorithm.
module Util.UnionFind(
   -- NB.  The functions in this module are not guaranteed thread-safe.


   UnionFind, -- :: type with parameter.  Instance of Eq.
   newElement, -- :: a -> IO (UnionFind a)
   toValue, -- :: UnionFind a -> a

   union, -- :: UnionFind a -> UnionFind a -> IO ()
   isSame, -- :: UnionFind a -> UnionFind a -> IO Bool
   sameElements, -- :: UnionFind a -> IO [UnionFind a]
   ) where

import Data.IORef

import Util.Computation(done)
import Util.ExtendedPrelude

-- -------------------------------------------------------------------
-- The datatype and instance of Eq
-- -------------------------------------------------------------------

data UnionFind a = UnionFind {
   UnionFind a -> a
value :: a,
   UnionFind a -> IORef [UnionFind a]
contentsRef :: IORef [UnionFind a],
      -- All items union'd with this one.
      -- Thus these contents form a tree-structure.
   UnionFind a -> IORef (Maybe (UnionFind a))
headRef :: IORef (Maybe (UnionFind a))
      -- If Just, an item with which this one is union'd, possibly
      -- indirectly.
      --
      -- To avoid spending lots of time chasing up long chains of
      -- head pointers, we in each case replace the head with the eventual
      -- parent.  I think this is Tarjan's algorithm and makes the operations
      -- almost linear (amortized time), but can't be bothered to chase up
      -- the reference.
   }

instance Eq (UnionFind a) where
   == :: UnionFind a -> UnionFind a -> Bool
(==) = (UnionFind a -> IORef [UnionFind a])
-> UnionFind a -> UnionFind a -> Bool
forall a b. Eq a => (b -> a) -> b -> b -> Bool
mapEq UnionFind a -> IORef [UnionFind a]
forall a. UnionFind a -> IORef [UnionFind a]
contentsRef

-- -------------------------------------------------------------------
-- The external functions
-- -------------------------------------------------------------------

newElement :: a -> IO (UnionFind a)
newElement :: a -> IO (UnionFind a)
newElement a
value =
   do
      IORef [UnionFind a]
contentsRef <- [UnionFind a] -> IO (IORef [UnionFind a])
forall a. a -> IO (IORef a)
newIORef []
      IORef (Maybe (UnionFind a))
headRef <- Maybe (UnionFind a) -> IO (IORef (Maybe (UnionFind a)))
forall a. a -> IO (IORef a)
newIORef Maybe (UnionFind a)
forall a. Maybe a
Nothing
      let
         unionFind :: UnionFind a
unionFind = UnionFind :: forall a.
a
-> IORef [UnionFind a]
-> IORef (Maybe (UnionFind a))
-> UnionFind a
UnionFind {
            value :: a
value = a
value,
            contentsRef :: IORef [UnionFind a]
contentsRef = IORef [UnionFind a]
contentsRef,
            headRef :: IORef (Maybe (UnionFind a))
headRef = IORef (Maybe (UnionFind a))
headRef
            }
      UnionFind a -> IO (UnionFind a)
forall (m :: * -> *) a. Monad m => a -> m a
return UnionFind a
unionFind

toValue :: UnionFind a -> a
toValue :: UnionFind a -> a
toValue = UnionFind a -> a
forall a. UnionFind a -> a
value

union :: UnionFind a -> UnionFind a -> IO ()
union :: UnionFind a -> UnionFind a -> IO ()
union UnionFind a
uf1 UnionFind a
uf2 =
   do
      UnionFind a
head1 <- UnionFind a -> IO (UnionFind a)
forall a. UnionFind a -> IO (UnionFind a)
getHead UnionFind a
uf1
      UnionFind a
head2 <- UnionFind a -> IO (UnionFind a)
forall a. UnionFind a -> IO (UnionFind a)
getHead UnionFind a
uf2
      if UnionFind a
head1 UnionFind a -> UnionFind a -> Bool
forall a. Eq a => a -> a -> Bool
== UnionFind a
head2
         then
            IO ()
forall (m :: * -> *). Monad m => m ()
done
         else
            do
               IORef (Maybe (UnionFind a)) -> Maybe (UnionFind a) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (UnionFind a -> IORef (Maybe (UnionFind a))
forall a. UnionFind a -> IORef (Maybe (UnionFind a))
headRef UnionFind a
head2) (UnionFind a -> Maybe (UnionFind a)
forall a. a -> Maybe a
Just UnionFind a
head1)

               [UnionFind a]
contents0 <- IORef [UnionFind a] -> IO [UnionFind a]
forall a. IORef a -> IO a
readIORef (UnionFind a -> IORef [UnionFind a]
forall a. UnionFind a -> IORef [UnionFind a]
contentsRef UnionFind a
head1)
               IORef [UnionFind a] -> [UnionFind a] -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (UnionFind a -> IORef [UnionFind a]
forall a. UnionFind a -> IORef [UnionFind a]
contentsRef UnionFind a
head1) (UnionFind a
head2 UnionFind a -> [UnionFind a] -> [UnionFind a]
forall a. a -> [a] -> [a]
: [UnionFind a]
contents0)

isSame :: UnionFind a -> UnionFind a -> IO Bool
isSame :: UnionFind a -> UnionFind a -> IO Bool
isSame UnionFind a
uf1 UnionFind a
uf2 =
   do
      UnionFind a
head1 <- UnionFind a -> IO (UnionFind a)
forall a. UnionFind a -> IO (UnionFind a)
getHead UnionFind a
uf1
      UnionFind a
head2 <- UnionFind a -> IO (UnionFind a)
forall a. UnionFind a -> IO (UnionFind a)
getHead UnionFind a
uf2
      Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (UnionFind a
head1 UnionFind a -> UnionFind a -> Bool
forall a. Eq a => a -> a -> Bool
== UnionFind a
head2)

sameElements :: UnionFind a -> IO [UnionFind a]
sameElements :: UnionFind a -> IO [UnionFind a]
sameElements UnionFind a
uf =
   do
      UnionFind a
head <- UnionFind a -> IO (UnionFind a)
forall a. UnionFind a -> IO (UnionFind a)
getHead UnionFind a
uf
      UnionFind a -> IO [UnionFind a]
forall a. UnionFind a -> IO [UnionFind a]
allContents UnionFind a
head
   where
      allContents :: UnionFind a -> IO [UnionFind a]
      allContents :: UnionFind a -> IO [UnionFind a]
allContents UnionFind a
uf =
         do
            [UnionFind a]
contents <- IORef [UnionFind a] -> IO [UnionFind a]
forall a. IORef a -> IO a
readIORef (UnionFind a -> IORef [UnionFind a]
forall a. UnionFind a -> IORef [UnionFind a]
contentsRef UnionFind a
uf)
            [[UnionFind a]]
innerContents <- (UnionFind a -> IO [UnionFind a])
-> [UnionFind a] -> IO [[UnionFind a]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM UnionFind a -> IO [UnionFind a]
forall a. UnionFind a -> IO [UnionFind a]
allContents [UnionFind a]
contents
            [UnionFind a] -> IO [UnionFind a]
forall (m :: * -> *) a. Monad m => a -> m a
return (UnionFind a
uf UnionFind a -> [UnionFind a] -> [UnionFind a]
forall a. a -> [a] -> [a]
: [[UnionFind a]] -> [UnionFind a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[UnionFind a]]
innerContents)

-- -------------------------------------------------------------------
-- Retrieving the head (the most important operation).
-- -------------------------------------------------------------------

getHead :: UnionFind a -> IO (UnionFind a)
getHead :: UnionFind a -> IO (UnionFind a)
getHead UnionFind a
unionFind =
   do
      Maybe (UnionFind a)
thisHeadOpt <- IORef (Maybe (UnionFind a)) -> IO (Maybe (UnionFind a))
forall a. IORef a -> IO a
readIORef (UnionFind a -> IORef (Maybe (UnionFind a))
forall a. UnionFind a -> IORef (Maybe (UnionFind a))
headRef UnionFind a
unionFind)
      case Maybe (UnionFind a)
thisHeadOpt of
         Maybe (UnionFind a)
Nothing -> UnionFind a -> IO (UnionFind a)
forall (m :: * -> *) a. Monad m => a -> m a
return UnionFind a
unionFind
         Just UnionFind a
unionFind2 ->
            do
               UnionFind a
thisHead <- UnionFind a -> IO (UnionFind a)
forall a. UnionFind a -> IO (UnionFind a)
getHead UnionFind a
unionFind2
               IORef (Maybe (UnionFind a)) -> Maybe (UnionFind a) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (UnionFind a -> IORef (Maybe (UnionFind a))
forall a. UnionFind a -> IORef (Maybe (UnionFind a))
headRef UnionFind a
unionFind) (UnionFind a -> Maybe (UnionFind a)
forall a. a -> Maybe a
Just UnionFind a
thisHead)
               UnionFind a -> IO (UnionFind a)
forall (m :: * -> *) a. Monad m => a -> m a
return UnionFind a
thisHead