module Jukebox.UnionFind(UF, Replacement((:>)), (=:=), rep, evalUF, execUF, runUF, S, isRep, initial, reps) where

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

type S a = Map a a
type UF a = State (S a)
data Replacement a = a :> a

runUF :: S a -> UF a b -> (b, S a)
runUF :: S a -> UF a b -> (b, S a)
runUF S a
s UF a b
m = UF a b -> S a -> (b, S a)
forall s a. State s a -> s -> (a, s)
runState UF a b
m S a
s

evalUF :: S a -> UF a b -> b
evalUF :: S a -> UF a b -> b
evalUF S a
s UF a b
m = (b, S a) -> b
forall a b. (a, b) -> a
fst (S a -> UF a b -> (b, S a)
forall a b. S a -> UF a b -> (b, S a)
runUF S a
s UF a b
m)

execUF :: S a -> UF a b -> S a
execUF :: S a -> UF a b -> S a
execUF S a
s UF a b
m = (b, S a) -> S a
forall a b. (a, b) -> b
snd (S a -> UF a b -> (b, S a)
forall a b. S a -> UF a b -> (b, S a)
runUF S a
s UF a b
m)

initial :: S a
initial :: S a
initial = S a
forall k a. Map k a
Map.empty

(=:=) :: Ord a => a -> a -> UF a (Maybe (Replacement a))
a
s =:= :: a -> a -> UF a (Maybe (Replacement a))
=:= a
t | a
s a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
t = Maybe (Replacement a) -> UF a (Maybe (Replacement a))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Replacement a)
forall a. Maybe a
Nothing
a
s =:= a
t = do
  a
rs <- a -> UF a a
forall a. Ord a => a -> UF a a
rep a
s
  a
rt <- a -> UF a a
forall a. Ord a => a -> UF a a
rep a
t
  case a
rs a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` a
rt of
    Ordering
EQ -> Maybe (Replacement a) -> UF a (Maybe (Replacement a))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Replacement a)
forall a. Maybe a
Nothing
    Ordering
LT -> do
      (Map a a -> Map a a) -> StateT (Map a a) Identity ()
forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
modify (a -> a -> Map a a -> Map a a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert a
rt a
rs)
      Maybe (Replacement a) -> UF a (Maybe (Replacement a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Replacement a -> Maybe (Replacement a)
forall a. a -> Maybe a
Just (a
rt a -> a -> Replacement a
forall a. a -> a -> Replacement a
:> a
rs))
    Ordering
GT -> do
      (Map a a -> Map a a) -> StateT (Map a a) Identity ()
forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
modify (a -> a -> Map a a -> Map a a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert a
rs a
rt)
      Maybe (Replacement a) -> UF a (Maybe (Replacement a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Replacement a -> Maybe (Replacement a)
forall a. a -> Maybe a
Just (a
rs a -> a -> Replacement a
forall a. a -> a -> Replacement a
:> a
rt))

{-# INLINE rep #-}
rep :: Ord a => a -> UF a a
rep :: a -> UF a a
rep a
s = do
  S a
m <- StateT (S a) Identity (S a)
forall (m :: * -> *) s. Monad m => StateT s m s
get
  case a -> S a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
s S a
m of
    Maybe a
Nothing -> a -> UF a a
forall (m :: * -> *) a. Monad m => a -> m a
return a
s
    Just a
t -> do
      a
u <- a -> UF a a
forall a. Ord a => a -> UF a a
rep a
t
      Bool -> StateT (S a) Identity () -> StateT (S a) Identity ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (a
t a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
u) (StateT (S a) Identity () -> StateT (S a) Identity ())
-> StateT (S a) Identity () -> StateT (S a) Identity ()
forall a b. (a -> b) -> a -> b
$ (S a -> S a) -> StateT (S a) Identity ()
forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
modify (a -> a -> S a -> S a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert a
s a
u)
      a -> UF a a
forall (m :: * -> *) a. Monad m => a -> m a
return a
u
      -- case Map.lookup t m of
      --   Nothing -> return t
      --   Just u -> do
      --     v <- rep' t u
      --     modify (Map.insert s v)
      --     return v

reps :: Ord a => UF a (a -> a)
reps :: UF a (a -> a)
reps = do
  S a
s <- StateT (S a) Identity (S a)
forall (m :: * -> *) s. Monad m => StateT s m s
get
  (a -> a) -> UF a (a -> a)
forall (m :: * -> *) a. Monad m => a -> m a
return (\a
x -> S a -> UF a a -> a
forall a b. S a -> UF a b -> b
evalUF S a
s (a -> UF a a
forall a. Ord a => a -> UF a a
rep a
x))

-- rep' :: Ord a => a -> a -> UF a a
-- rep' s t = do
--   m <- get
--   case Map.lookup t m of
--     Nothing -> do
--       modify (Map.insert s t)
--       return t
--     Just u -> do
--       v <- rep' t u
--       modify (Map.insert s v)
--       return v

isRep :: Ord a => a -> UF a Bool
isRep :: a -> UF a Bool
isRep a
t = do
  a
t' <- a -> UF a a
forall a. Ord a => a -> UF a a
rep a
t
  Bool -> UF a Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (a
t a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
t')