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 :: forall a b. 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 :: forall a b. 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 :: forall a b. 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 :: forall a. S a
initial = Map a a
forall k a. Map k a
Map.empty
(=:=) :: Ord a => a -> a -> UF a (Maybe (Replacement a))
a
s =:= :: forall a. Ord a => 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)
-> StateT (S a) Identity (Maybe (Replacement a))
forall a. a -> StateT (S a) Identity 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)
-> StateT (S a) Identity (Maybe (Replacement a))
forall a. a -> StateT (S a) Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Replacement a)
forall a. Maybe a
Nothing
Ordering
LT -> do
(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
rt a
rs)
Maybe (Replacement a)
-> StateT (S a) Identity (Maybe (Replacement a))
forall a. a -> StateT (S a) Identity 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
(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
rs a
rt)
Maybe (Replacement a)
-> StateT (S a) Identity (Maybe (Replacement a))
forall a. a -> StateT (S a) Identity 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 :: forall a. Ord a => 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 a. a -> StateT (S a) Identity 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 a. a -> StateT (S a) Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return a
u
reps :: Ord a => UF a (a -> a)
reps :: forall a. Ord a => 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 a. a -> StateT (S a) Identity 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))
isRep :: Ord a => a -> UF a Bool
isRep :: forall a. Ord a => 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 a. a -> StateT (S a) Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (a
t a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
t')