{-# LANGUAGE MultiParamTypeClasses #-}

--------------------------------------------------------------------------------
-- |
-- Module      : Data.Equivalence.STT
-- Copyright   : Patrick Bahr, 2010
-- License     : BSD-3-Clause
--
-- Maintainer  :  Patrick Bahr, Andreas Abel
-- Stability   :  stable
-- Portability :  non-portable (MPTC)
--
-- This is an implementation of Tarjan's Union-Find algorithm (Robert
-- E. Tarjan. "Efficiency of a Good But Not Linear Set Union
-- Algorithm", JACM 22(2), 1975) in order to maintain an equivalence
-- relation.
--
-- This implementation is a port of the /union-find/ package using the
-- ST monad transformer (instead of the IO monad).
--
-- The implementation is based on mutable references.  Each
-- equivalence class has exactly one member that serves as its
-- representative element.  Every element either is the representative
-- element of its equivalence class or points to another element in
-- the same equivalence class.  Equivalence testing thus consists of
-- following the pointers to the representative elements and then
-- comparing these for identity.
--
-- The algorithm performs lazy path compression.  That is, whenever we
-- walk along a path greater than length 1 we automatically update the
-- pointers along the path to directly point to the representative
-- element.  Consequently future lookups will be have a path length of
-- at most 1.
--
-- Each equivalence class remains a descriptor, i.e. some piece of
-- data attached to an equivalence class which is combined when two
-- classes are unioned.
--
--------------------------------------------------------------------------------

module Data.Equivalence.STT
  (
   -- * Equivalence Relation
    Equiv
  , Class
  , leastEquiv
  -- * Operations on Equivalence Classes
  , getClass
  , combine
  , combineAll
  , same
  , desc
  , remove
  -- * Operations on Elements
  , equate
  , equateAll
  , equivalent
  , classDesc
  , removeClass
  -- Getting all represented items
  , values
  , classes
  ) where

import Control.Monad.ST.Trans
import Control.Monad

import Data.Maybe

import Data.Map (Map)
import qualified Data.Map as Map

{-| Abstract representation of an equivalence class. -}

newtype Class s c a = Class (STRef s (Entry s c a))

{-| This type represents a reference to an entry in the tree data
structure. An entry of type 'Entry' @s c a@ lives in the state space
indexed by @s@, contains equivalence class descriptors of type @c@ and
has elements of type @a@.-}

newtype Entry s c a = Entry {forall s c a. Entry s c a -> STRef s (EntryData s c a)
unentry :: STRef s (EntryData s c a)}

{-| This type represents entries (nodes) in the tree data
structure. Entry data of type 'EntryData' @s c a@ lives in the state space
indexed by @s@, contains equivalence class descriptors of type @c@ and
has elements of type @a@.  -}

data EntryData s c a
  = Node {
      forall s c a. EntryData s c a -> Entry s c a
entryParent :: Entry s c a,
      forall s c a. EntryData s c a -> a
entryValue :: a
    }
  | Root {
      forall s c a. EntryData s c a -> c
entryDesc :: c,
      forall s c a. EntryData s c a -> Int
entryWeight :: Int,
      entryValue :: a,
      forall s c a. EntryData s c a -> Bool
entryDeleted :: Bool
    }

type Entries s c a = STRef s (Map a (Entry s c a))

{-| This is the top-level data structure that represents an
equivalence relation. An equivalence relation of type 'Equiv' @s c a@
lives in the state space indexed by @s@, contains equivalence class
descriptors of type @c@ and has elements of type @a@. -}

data Equiv s c a = Equiv {
      -- | Maps elements to their entry in the tree data structure.
      forall s c a. Equiv s c a -> Entries s c a
entries :: Entries s c a,
      -- | Constructs an equivalence class descriptor for a singleton class.
      forall s c a. Equiv s c a -> a -> c
singleDesc :: a -> c,
      -- | Combines the equivalence class descriptor of two classes
      --   which are meant to be combined.
      forall s c a. Equiv s c a -> c -> c -> c
combDesc :: c -> c -> c
      }

{-| This function constructs the initial data structure for
maintaining an equivalence relation. That is, it represents the finest
(or least) equivalence class (of the set of all elements of type
@a@). The arguments are used to maintain equivalence class
descriptors. -}

leastEquiv
  :: (Monad m, Applicative m)
  => (a -> c)      -- ^ Used to construct an equivalence class descriptor for a singleton class.
  -> (c -> c -> c) -- ^ Used to combine the equivalence class descriptor of two classes
                   --   which are meant to be combined.
  -> STT s m (Equiv s c a)
leastEquiv :: forall (m :: * -> *) a c s.
(Monad m, Applicative m) =>
(a -> c) -> (c -> c -> c) -> STT s m (Equiv s c a)
leastEquiv a -> c
mk c -> c -> c
com = do
  es <- Map a (Entry s c a) -> STT s m (STRef s (Map a (Entry s c a)))
forall (m :: * -> *) a s. Applicative m => a -> STT s m (STRef s a)
newSTRef Map a (Entry s c a)
forall k a. Map k a
Map.empty
  return Equiv {entries = es, singleDesc = mk, combDesc = com}


{-| This function returns the representative entry of the argument's
equivalence class (i.e. the root of its tree) or @Nothing@ if it is
the representative itself.

This function performs path compression.  -}

representative' :: (Monad m, Applicative m) => Entry s c a -> STT s m (Maybe (Entry s c a),Bool)
representative' :: forall (m :: * -> *) s c a.
(Monad m, Applicative m) =>
Entry s c a -> STT s m (Maybe (Entry s c a), Bool)
representative' (Entry STRef s (EntryData s c a)
e) = do
  ed <- STRef s (EntryData s c a) -> STT s m (EntryData s c a)
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s (EntryData s c a)
e
  case ed of
    Root {entryDeleted :: forall s c a. EntryData s c a -> Bool
entryDeleted = Bool
del} -> do
      (Maybe (Entry s c a), Bool) -> STT s m (Maybe (Entry s c a), Bool)
forall a. a -> STT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Entry s c a)
forall a. Maybe a
Nothing, Bool
del)
    Node {entryParent :: forall s c a. EntryData s c a -> Entry s c a
entryParent = Entry s c a
parent} -> do
      (mparent',del) <- Entry s c a -> STT s m (Maybe (Entry s c a), Bool)
forall (m :: * -> *) s c a.
(Monad m, Applicative m) =>
Entry s c a -> STT s m (Maybe (Entry s c a), Bool)
representative' Entry s c a
parent
      case mparent' of
        Maybe (Entry s c a)
Nothing -> (Maybe (Entry s c a), Bool) -> STT s m (Maybe (Entry s c a), Bool)
forall a. a -> STT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Maybe (Entry s c a), Bool)
 -> STT s m (Maybe (Entry s c a), Bool))
-> (Maybe (Entry s c a), Bool)
-> STT s m (Maybe (Entry s c a), Bool)
forall a b. (a -> b) -> a -> b
$ (Entry s c a -> Maybe (Entry s c a)
forall a. a -> Maybe a
Just Entry s c a
parent, Bool
del)
        Just Entry s c a
parent' -> STRef s (EntryData s c a) -> EntryData s c a -> STT s m ()
forall (m :: * -> *) s a.
Applicative m =>
STRef s a -> a -> STT s m ()
writeSTRef STRef s (EntryData s c a)
e EntryData s c a
ed{entryParent = parent'} STT s m ()
-> STT s m (Maybe (Entry s c a), Bool)
-> STT s m (Maybe (Entry s c a), Bool)
forall a b. STT s m a -> STT s m b -> STT s m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (Maybe (Entry s c a), Bool) -> STT s m (Maybe (Entry s c a), Bool)
forall a. a -> STT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Entry s c a -> Maybe (Entry s c a)
forall a. a -> Maybe a
Just Entry s c a
parent', Bool
del)


{-| This function returns the representative entry of the argument's
equivalence class (i.e. the root of its tree).

This function performs path compression.  -}

representative :: (Monad m, Applicative m, Ord a) => Equiv s c a -> a -> STT s m (Entry s c a)
representative :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m (Entry s c a)
representative Equiv s c a
eq a
v = do
  mentry <- Equiv s c a -> a -> STT s m (Maybe (Entry s c a))
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m (Maybe (Entry s c a))
getEntry Equiv s c a
eq a
v
  case mentry of -- check whether there is an entry
    Maybe (Entry s c a)
Nothing -> Equiv s c a -> a -> STT s m (Entry s c a)
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m (Entry s c a)
mkEntry Equiv s c a
eq a
v -- if not, create a new one
    Just Entry s c a
entry -> do
      (mrepr,del) <- Entry s c a -> STT s m (Maybe (Entry s c a), Bool)
forall (m :: * -> *) s c a.
(Monad m, Applicative m) =>
Entry s c a -> STT s m (Maybe (Entry s c a), Bool)
representative' Entry s c a
entry
      if del -- check whether equivalence class was deleted
        then mkEntry eq v -- if so, create a new entry
        else case mrepr of
               Maybe (Entry s c a)
Nothing -> Entry s c a -> STT s m (Entry s c a)
forall a. a -> STT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return Entry s c a
entry
               Just Entry s c a
repr -> Entry s c a -> STT s m (Entry s c a)
forall a. a -> STT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return Entry s c a
repr

{-| This function provides the representative entry of the given
equivalence class. This function performs path compression. -}

classRep :: (Monad m, Applicative m, Ord a) => Equiv s c a -> Class s c a -> STT s m (Entry s c a)
classRep :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> Class s c a -> STT s m (Entry s c a)
classRep Equiv s c a
eq (Class STRef s (Entry s c a)
p) = do
  entry <- STRef s (Entry s c a) -> STT s m (Entry s c a)
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s (Entry s c a)
p
  (mrepr,del) <- representative' entry
  if del -- check whether equivalence class was deleted
    then do v <- liftM entryValue $ readSTRef (unentry entry)
            en <- getEntry' eq v -- if so, create a new entry
            (mrepr,del) <- representative' en
            if del then do
                en' <- mkEntry' eq en
                writeSTRef p en'
                return en'
              else return (fromMaybe en mrepr)
    else return (fromMaybe entry mrepr)


{-| This function constructs a new (root) entry containing the given
entry's value, inserts it into the lookup table (thereby removing any
existing entry). -}

mkEntry' :: (Monad m, Applicative m, Ord a)
        => Equiv s c a -> Entry s c a
        -> STT s m (Entry s c a)  -- ^ the constructed entry
mkEntry' :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> Entry s c a -> STT s m (Entry s c a)
mkEntry' Equiv s c a
eq (Entry STRef s (EntryData s c a)
e) = STRef s (EntryData s c a) -> STT s m (EntryData s c a)
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s (EntryData s c a)
e STT s m (EntryData s c a)
-> (EntryData s c a -> STT s m (Entry s c a))
-> STT s m (Entry s c a)
forall a b. STT s m a -> (a -> STT s m b) -> STT s m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Equiv s c a -> a -> STT s m (Entry s c a)
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m (Entry s c a)
mkEntry Equiv s c a
eq (a -> STT s m (Entry s c a))
-> (EntryData s c a -> a)
-> EntryData s c a
-> STT s m (Entry s c a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EntryData s c a -> a
forall s c a. EntryData s c a -> a
entryValue

{-| This function constructs a new (root) entry containing the given
value, inserts it into the lookup table (thereby removing any existing
entry). -}

mkEntry :: (Monad m, Applicative m, Ord a)
        => Equiv s c a -> a
        -> STT s m (Entry s c a)  -- ^ the constructed entry
mkEntry :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m (Entry s c a)
mkEntry Equiv {entries :: forall s c a. Equiv s c a -> Entries s c a
entries = Entries s c a
mref, singleDesc :: forall s c a. Equiv s c a -> a -> c
singleDesc = a -> c
mkDesc} a
val = do
  e <- EntryData s c a -> STT s m (STRef s (EntryData s c a))
forall (m :: * -> *) a s. Applicative m => a -> STT s m (STRef s a)
newSTRef Root
       { entryDesc :: c
entryDesc = a -> c
mkDesc a
val,
         entryWeight :: Int
entryWeight = Int
1,
         entryValue :: a
entryValue = a
val,
         entryDeleted :: Bool
entryDeleted = Bool
False
       }
  let entry = STRef s (EntryData s c a) -> Entry s c a
forall s c a. STRef s (EntryData s c a) -> Entry s c a
Entry STRef s (EntryData s c a)
e
  m <- readSTRef mref
  writeSTRef mref (Map.insert val entry m)
  return entry

{-| This function provides the equivalence class the given element is
contained in. -}

getClass :: (Monad m, Applicative m, Ord a) => Equiv s c a -> a -> STT s m (Class s c a)
getClass :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m (Class s c a)
getClass Equiv s c a
eq a
v = do
  en <- (Equiv s c a -> a -> STT s m (Entry s c a)
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m (Entry s c a)
getEntry' Equiv s c a
eq a
v)
  liftM Class $ newSTRef en


getEntry' :: (Monad m, Applicative m, Ord a) => Equiv s c a -> a -> STT s m (Entry s c a)
getEntry' :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m (Entry s c a)
getEntry' Equiv s c a
eq a
v = do
  mentry <- Equiv s c a -> a -> STT s m (Maybe (Entry s c a))
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m (Maybe (Entry s c a))
getEntry Equiv s c a
eq a
v
  case mentry of
    Maybe (Entry s c a)
Nothing -> Equiv s c a -> a -> STT s m (Entry s c a)
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m (Entry s c a)
mkEntry Equiv s c a
eq a
v
    Just Entry s c a
entry -> Entry s c a -> STT s m (Entry s c a)
forall a. a -> STT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return Entry s c a
entry

{-| This function looks up the entry of the given element in the given
equivalence relation representation or @Nothing@ if there is none,
yet.  -}

getEntry :: (Monad m, Applicative m, Ord a) => Equiv s c a -> a -> STT s m (Maybe (Entry s c a))
getEntry :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m (Maybe (Entry s c a))
getEntry Equiv{ entries :: forall s c a. Equiv s c a -> Entries s c a
entries = Entries s c a
mref } a
val = do
  a -> Map a (Entry s c a) -> Maybe (Entry s c a)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
val (Map a (Entry s c a) -> Maybe (Entry s c a))
-> STT s m (Map a (Entry s c a)) -> STT s m (Maybe (Entry s c a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Entries s c a -> STT s m (Map a (Entry s c a))
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef Entries s c a
mref

{-| This function equates the two given (representative) elements. That
is, it unions the equivalence classes of the two elements and combines
their descriptor. The returned entry is the representative of the new
equivalence class -}

equateEntry :: (Monad m, Applicative m, Ord a) => Equiv s c a -> Entry s c a -> Entry s c a -> STT s m (Entry s c a)
equateEntry :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> Entry s c a -> Entry s c a -> STT s m (Entry s c a)
equateEntry Equiv {combDesc :: forall s c a. Equiv s c a -> c -> c -> c
combDesc = c -> c -> c
mkDesc} repx :: Entry s c a
repx@(Entry STRef s (EntryData s c a)
rx) repy :: Entry s c a
repy@(Entry STRef s (EntryData s c a)
ry) =
  if (STRef s (EntryData s c a)
rx STRef s (EntryData s c a) -> STRef s (EntryData s c a) -> Bool
forall a. Eq a => a -> a -> Bool
/= STRef s (EntryData s c a)
ry) then do
    dx <- STRef s (EntryData s c a) -> STT s m (EntryData s c a)
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s (EntryData s c a)
rx
    dy <- readSTRef ry
    case (dx, dy) of
      ( Root{entryWeight :: forall s c a. EntryData s c a -> Int
entryWeight = Int
wx, entryDesc :: forall s c a. EntryData s c a -> c
entryDesc = c
chx, entryValue :: forall s c a. EntryData s c a -> a
entryValue = a
vx}
        , Root{entryWeight :: forall s c a. EntryData s c a -> Int
entryWeight = Int
wy, entryDesc :: forall s c a. EntryData s c a -> c
entryDesc = c
chy, entryValue :: forall s c a. EntryData s c a -> a
entryValue = a
vy} ) ->
        if  Int
wx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
wy
        then do
          STRef s (EntryData s c a) -> EntryData s c a -> STT s m ()
forall (m :: * -> *) s a.
Applicative m =>
STRef s a -> a -> STT s m ()
writeSTRef STRef s (EntryData s c a)
ry Node {entryParent :: Entry s c a
entryParent = Entry s c a
repx, entryValue :: a
entryValue = a
vy}
          STRef s (EntryData s c a) -> EntryData s c a -> STT s m ()
forall (m :: * -> *) s a.
Applicative m =>
STRef s a -> a -> STT s m ()
writeSTRef STRef s (EntryData s c a)
rx EntryData s c a
dx{entryWeight = wx + wy, entryDesc = mkDesc chx chy}
          Entry s c a -> STT s m (Entry s c a)
forall a. a -> STT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return Entry s c a
repx
        else do
          STRef s (EntryData s c a) -> EntryData s c a -> STT s m ()
forall (m :: * -> *) s a.
Applicative m =>
STRef s a -> a -> STT s m ()
writeSTRef STRef s (EntryData s c a)
rx Node {entryParent :: Entry s c a
entryParent = Entry s c a
repy, entryValue :: a
entryValue = a
vx}
          STRef s (EntryData s c a) -> EntryData s c a -> STT s m ()
forall (m :: * -> *) s a.
Applicative m =>
STRef s a -> a -> STT s m ()
writeSTRef STRef s (EntryData s c a)
ry EntryData s c a
dy{entryWeight = wx + wy, entryDesc = mkDesc chx chy}
          Entry s c a -> STT s m (Entry s c a)
forall a. a -> STT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return Entry s c a
repy

      (EntryData s c a, EntryData s c a)
_ -> [Char] -> STT s m (Entry s c a)
forall a. HasCallStack => [Char] -> a
error [Char]
"error on `equateEntry`"
      -- this should not happen as this function is only called by
      -- 'combineEntries', which always uses representative entries
  else Entry s c a -> STT s m (Entry s c a)
forall a. a -> STT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return  Entry s c a
repx

combineEntries :: (Monad m, Applicative m, Ord a)
               => Equiv s c a -> [b] -> (b -> STT s m (Entry s c a)) -> STT s m ()
combineEntries :: forall (m :: * -> *) a s c b.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> [b] -> (b -> STT s m (Entry s c a)) -> STT s m ()
combineEntries  Equiv s c a
_ [] b -> STT s m (Entry s c a)
_ = () -> STT s m ()
forall a. a -> STT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
combineEntries Equiv s c a
eq (b
e:[b]
es) b -> STT s m (Entry s c a)
rep = do
  er <- b -> STT s m (Entry s c a)
rep b
e
  run er es
    where run :: Entry s c a -> [b] -> STT s m ()
run Entry s c a
er (b
f:[b]
r) = do
            fr <- b -> STT s m (Entry s c a)
rep b
f
            er' <- equateEntry eq er fr
            run er' r
          run Entry s c a
_ [b]
_ = () -> STT s m ()
forall a. a -> STT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()


{-| This function combines all equivalence classes in the given
list. Afterwards all elements in the argument list represent the same
equivalence class! -}

combineAll :: (Monad m, Applicative m, Ord a) => Equiv s c a -> [Class s c a] -> STT s m ()
combineAll :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> [Class s c a] -> STT s m ()
combineAll Equiv s c a
eq [Class s c a]
cls = Equiv s c a
-> [Class s c a]
-> (Class s c a -> STT s m (Entry s c a))
-> STT s m ()
forall (m :: * -> *) a s c b.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> [b] -> (b -> STT s m (Entry s c a)) -> STT s m ()
combineEntries Equiv s c a
eq [Class s c a]
cls (Equiv s c a -> Class s c a -> STT s m (Entry s c a)
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> Class s c a -> STT s m (Entry s c a)
classRep Equiv s c a
eq)


{-| This function combines the two given equivalence
classes. Afterwards both arguments represent the same equivalence
class! One of it is returned in order to represent the new combined
equivalence class. -}

combine :: (Monad m, Applicative m, Ord a) => Equiv s c a -> Class s c a -> Class s c a -> STT s m (Class s c a)
combine :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> Class s c a -> Class s c a -> STT s m (Class s c a)
combine Equiv s c a
eq Class s c a
x Class s c a
y = Equiv s c a -> [Class s c a] -> STT s m ()
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> [Class s c a] -> STT s m ()
combineAll Equiv s c a
eq [Class s c a
x,Class s c a
y] STT s m () -> STT s m (Class s c a) -> STT s m (Class s c a)
forall a b. STT s m a -> STT s m b -> STT s m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Class s c a -> STT s m (Class s c a)
forall a. a -> STT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return Class s c a
x


{-| This function equates the element in the given list. That is, it
unions the equivalence classes of the elements and combines their
descriptor. -}

equateAll :: (Monad m, Applicative m, Ord a) => Equiv s c a -> [a] -> STT s m ()
equateAll :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> [a] -> STT s m ()
equateAll Equiv s c a
eq [a]
cls = Equiv s c a -> [a] -> (a -> STT s m (Entry s c a)) -> STT s m ()
forall (m :: * -> *) a s c b.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> [b] -> (b -> STT s m (Entry s c a)) -> STT s m ()
combineEntries Equiv s c a
eq [a]
cls (Equiv s c a -> a -> STT s m (Entry s c a)
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m (Entry s c a)
representative Equiv s c a
eq)

{-| This function equates the two given elements. That is, it unions
the equivalence classes of the two elements and combines their
descriptor. -}

equate :: (Monad m, Applicative m, Ord a) => Equiv s c a -> a -> a -> STT s m ()
equate :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> a -> STT s m ()
equate Equiv s c a
eq a
x a
y = Equiv s c a -> [a] -> STT s m ()
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> [a] -> STT s m ()
equateAll Equiv s c a
eq [a
x,a
y]


{-| This function returns the descriptor of the given
equivalence class. -}

desc :: (Monad m, Applicative m, Ord a) => Equiv s c a -> Class s c a -> STT s m c
desc :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> Class s c a -> STT s m c
desc Equiv s c a
eq Class s c a
cl = do
  Entry e <- Equiv s c a -> Class s c a -> STT s m (Entry s c a)
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> Class s c a -> STT s m (Entry s c a)
classRep Equiv s c a
eq Class s c a
cl
  liftM entryDesc $ readSTRef e

{-| This function returns the descriptor of the given element's
equivalence class. -}

classDesc :: (Monad m, Applicative m, Ord a) => Equiv s c a -> a -> STT s m c
classDesc :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m c
classDesc Equiv s c a
eq a
val = do
  Entry e <- Equiv s c a -> a -> STT s m (Entry s c a)
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m (Entry s c a)
representative Equiv s c a
eq a
val
  liftM entryDesc $ readSTRef e


{-| This function decides whether the two given equivalence classes
are the same. -}

same :: (Monad m, Applicative m, Ord a) => Equiv s c a -> Class s c a -> Class s c a -> STT s m Bool
same :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> Class s c a -> Class s c a -> STT s m Bool
same Equiv s c a
eq Class s c a
c1 Class s c a
c2 = do
  (Entry r1) <- Equiv s c a -> Class s c a -> STT s m (Entry s c a)
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> Class s c a -> STT s m (Entry s c a)
classRep Equiv s c a
eq Class s c a
c1
  (Entry r2) <- classRep eq c2
  return (r1 == r2)

{-| This function decides whether the two given elements are in the
same equivalence class according to the given equivalence relation
representation. -}

equivalent :: (Monad m, Applicative m, Ord a) => Equiv s c a -> a -> a -> STT s m Bool
equivalent :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> a -> STT s m Bool
equivalent Equiv s c a
eq a
v1 a
v2 = do
  (Entry r1) <- Equiv s c a -> a -> STT s m (Entry s c a)
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m (Entry s c a)
representative Equiv s c a
eq a
v1
  (Entry r2) <- representative eq v2
  return (r1 == r2)



{-|
  This function modifies the content of a reference cell.
 -}

modifySTRef :: (Monad m, Applicative m) => STRef s a -> (a -> a) -> STT s m ()
modifySTRef :: forall (m :: * -> *) s a.
(Monad m, Applicative m) =>
STRef s a -> (a -> a) -> STT s m ()
modifySTRef STRef s a
r a -> a
f = STRef s a -> STT s m a
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s a
r STT s m a -> (a -> STT s m ()) -> STT s m ()
forall a b. STT s m a -> (a -> STT s m b) -> STT s m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (STRef s a -> a -> STT s m ()
forall (m :: * -> *) s a.
Applicative m =>
STRef s a -> a -> STT s m ()
writeSTRef STRef s a
r (a -> STT s m ()) -> (a -> a) -> a -> STT s m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f)


{-| This function marks the given root entry as deleted.  -}

removeEntry :: (Monad m, Applicative m, Ord a) => Entry s c a -> STT s m ()
removeEntry :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Entry s c a -> STT s m ()
removeEntry (Entry STRef s (EntryData s c a)
r) = STRef s (EntryData s c a)
-> (EntryData s c a -> EntryData s c a) -> STT s m ()
forall (m :: * -> *) s a.
(Monad m, Applicative m) =>
STRef s a -> (a -> a) -> STT s m ()
modifySTRef STRef s (EntryData s c a)
r EntryData s c a -> EntryData s c a
forall {s} {c} {a} {s}. EntryData s c a -> EntryData s c a
change
    where change :: EntryData s c a -> EntryData s c a
change EntryData s c a
e = EntryData s c a
e {entryDeleted = True}


{-| This function removes the given equivalence class. If the
equivalence class does not exist anymore, @False@ is returned;
otherwise @True@. -}

remove :: (Monad m, Applicative m, Ord a) => Equiv s c a -> Class s c a -> STT s m Bool
remove :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> Class s c a -> STT s m Bool
remove Equiv s c a
eq (Class STRef s (Entry s c a)
p) = do
  entry <- STRef s (Entry s c a) -> STT s m (Entry s c a)
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s (Entry s c a)
p
  (mrepr,del) <- representative' entry
  if del then do
        v <- liftM entryValue $ readSTRef (unentry entry)
        men <- getEntry eq v
        case men of
          Maybe (Entry s c a)
Nothing -> Bool -> STT s m Bool
forall a. a -> STT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
          Just Entry s c a
en -> do
            STRef s (Entry s c a) -> Entry s c a -> STT s m ()
forall (m :: * -> *) s a.
Applicative m =>
STRef s a -> a -> STT s m ()
writeSTRef STRef s (Entry s c a)
p Entry s c a
en
            (mentry,del) <- Entry s c a -> STT s m (Maybe (Entry s c a), Bool)
forall (m :: * -> *) s c a.
(Monad m, Applicative m) =>
Entry s c a -> STT s m (Maybe (Entry s c a), Bool)
representative' Entry s c a
en
            if del
              then return False
              else removeEntry (fromMaybe en mentry)
                   >> return True
    else removeEntry (fromMaybe entry mrepr)
         >> return True

{-| This function removes the equivalence class of the given
element. If there is no corresponding equivalence class, @False@ is
returned; otherwise @True@. -}

removeClass :: (Monad m, Applicative m, Ord a) => Equiv s c a -> a -> STT s m Bool
removeClass :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m Bool
removeClass Equiv s c a
eq a
v = do
  mentry <- Equiv s c a -> a -> STT s m (Maybe (Entry s c a))
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> a -> STT s m (Maybe (Entry s c a))
getEntry Equiv s c a
eq a
v
  case mentry of
    Maybe (Entry s c a)
Nothing -> Bool -> STT s m Bool
forall a. a -> STT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
    Just Entry s c a
entry -> do
      (mentry, del) <- Entry s c a -> STT s m (Maybe (Entry s c a), Bool)
forall (m :: * -> *) s c a.
(Monad m, Applicative m) =>
Entry s c a -> STT s m (Maybe (Entry s c a), Bool)
representative' Entry s c a
entry
      if del
        then return False
        else removeEntry (fromMaybe entry mentry)
             >> return True

{-| This function returns all values represented by
   some equivalence class. -}

values :: (Monad m, Applicative m, Ord a) => Equiv s c a -> STT s m [a]
values :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> STT s m [a]
values Equiv {entries :: forall s c a. Equiv s c a -> Entries s c a
entries = Entries s c a
mref} = Map a (Entry s c a) -> [a]
forall k a. Map k a -> [k]
Map.keys (Map a (Entry s c a) -> [a])
-> STT s m (Map a (Entry s c a)) -> STT s m [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Entries s c a -> STT s m (Map a (Entry s c a))
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef Entries s c a
mref

{-| This function returns the list of
   all equivalence classes. -}

classes :: (Monad m, Applicative m, Ord a) => Equiv s c a -> STT s m [Class s c a]
classes :: forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Equiv s c a -> STT s m [Class s c a]
classes Equiv {entries :: forall s c a. Equiv s c a -> Entries s c a
entries = Entries s c a
mref} = do
  allEntries <- Map a (Entry s c a) -> [Entry s c a]
forall k a. Map k a -> [a]
Map.elems (Map a (Entry s c a) -> [Entry s c a])
-> STT s m (Map a (Entry s c a)) -> STT s m [Entry s c a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Entries s c a -> STT s m (Map a (Entry s c a))
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef Entries s c a
mref
  rootEntries <- filterM isRoot allEntries
  mapM (fmap Class . newSTRef) $ rootEntries
    where
      isRoot :: Entry s c a -> STT s m Bool
isRoot Entry s c a
e = do
        x <- STRef s (EntryData s c a) -> STT s m (EntryData s c a)
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef (Entry s c a -> STRef s (EntryData s c a)
forall s c a. Entry s c a -> STRef s (EntryData s c a)
unentry Entry s c a
e)
        case x of
          Node {} -> Bool -> STT s m Bool
forall a. a -> STT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
          Root {} -> Bool -> STT s m Bool
forall a. a -> STT s m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True