-- | A union-find data structure.

module Test.QuickSpec.Reasoning.UnionFind(UF, Replacement((:>)), newSym, (=:=), rep, evalUF, execUF, runUF, S, isRep, initial) where

import Prelude hiding (min)
import Control.Monad
import Control.Monad.Trans.State.Strict
import Data.IntMap(IntMap)
import qualified Data.IntMap as IntMap

data Info = Rep Int | NonRep Int

defaultInfo :: Info
defaultInfo = Rep 1

data S = S {
  info :: IntMap Info,
  sym :: Int
  }

type UF = State S
data Replacement = Int :> Int

runUF :: S -> UF a -> (a, S)
runUF s m = runState m s

evalUF :: S -> UF a -> a
evalUF s m = fst (runUF s m)

execUF :: S -> UF a -> S
execUF s m = snd (runUF s m)

initial :: Int -> S
initial n = S IntMap.empty n

modifyInfo f = modify (\s -> s { info = f (info s) })
modifySym f = modify (\s -> s { sym = f (sym s) })
putInfo i = modifyInfo (const i)

newSym :: UF Int
newSym = do
  s <- get
  modifySym (+1)
  return (sym s)

(=:=) :: Int -> Int -> UF (Maybe Replacement)
s =:= t | s == t = return Nothing
s =:= t = do
  rs <- rep s
  rt <- rep t
  if (rs /= rt) then fmap Just (unifyRep rs rt) else return Nothing

unifyRep :: Int -> Int -> UF Replacement
unifyRep s t = do
  ss <- gets (IntMap.findWithDefault 0 s . reps)
  st <- gets (IntMap.findWithDefault 0 t . reps)
  if ss <= st then replace s t ss st else replace t s st ss

replace :: Int -> Int -> Int -> Int -> UF Replacement
replace s t ss st = do
  modifyReps (IntMap.delete s . IntMap.insert t (ss+st+1))
  modifyLinks (IntMap.insert s t)
  return (s :> t)

rep :: Int -> UF Int
rep t = do
  m <- fmap links get
  case IntMap.lookup t m of
    Nothing -> return t
    Just t' -> do
      r <- rep t'
      when (t' /= r) $ modifyLinks (IntMap.insert t r)
      return r

isRep :: Int -> UF Bool
isRep t = do
  t' <- rep t
  return (t == t')