UnionFind.hs

module Data.Struct.Internal.UnionFind where

import Control.Monad (when)
import Control.Monad.Primitive
import Data.Struct.Internal

-- | Union-Find
-- >>> a <- new
-- >>> b <- new
-- >>> c <- new
-- >>> find a b
-- False
-- >>> find a c
-- False
-- >>> unite a b
-- >>> find a b
-- True
-- >>> find a c
-- False
-- >>> unite b c
-- >>> find a c
-- True
newtype UnionFind s = UnionFind (Object s)

instance Struct UnionFind where
  struct _ = Dict

instance Eq (UnionFind s) where
  (==) = eqStruct

parent :: Slot UnionFind UnionFind
parent = slot 0

rank :: Field UnionFind Int
rank = field 1

new :: PrimMonad m => m (UnionFind (PrimState m))
new = st $ do
  this <- alloc 2
  set parent this this
  setField rank this 0
  return this
{-# INLINE new #-}

unite :: PrimMonad m => UnionFind (PrimState m) -> UnionFind (PrimState m) -> m ()
unite a b = st $ do
  aroot <- representative a
  broot <- representative b

  when (aroot /= broot) $ do
    arank <- getField rank aroot
    brank <- getField rank broot
    case compare arank brank of
      LT -> set parent aroot broot
      GT -> set parent broot aroot
      EQ -> do
        set parent aroot broot
        setField rank broot (brank + 1)

representative :: PrimMonad m => UnionFind (PrimState m) -> m (UnionFind (PrimState m))
representative this = st $ do
  p <- get parent this
  if p == this then
    return p
  else do
    r <- representative p
    set parent this r
    return r

find :: PrimMonad m => UnionFind (PrimState m) -> UnionFind (PrimState m) -> m Bool
find a b = (==) <$> representative a <*> representative b
{-# INLINE find #-}