{-# 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
  STRef s (Map a (Entry s c a))
es <- forall (m :: * -> *) a s. Applicative m => a -> STT s m (STRef s a)
newSTRef forall k a. Map k a
Map.empty
  forall (m :: * -> *) a. Monad m => a -> m a
return Equiv {entries :: STRef s (Map a (Entry s c a))
entries = STRef s (Map a (Entry s c a))
es, singleDesc :: a -> c
singleDesc = a -> c
mk, combDesc :: c -> c -> c
combDesc = c -> c -> c
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
  EntryData s c a
ed <- forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s (EntryData s c a)
e
  case EntryData s c a
ed of
    Root {entryDeleted :: forall s c a. EntryData s c a -> Bool
entryDeleted = Bool
del} -> do
      forall (m :: * -> *) a. Monad m => a -> m a
return (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
      (Maybe (Entry s c a)
mparent',Bool
del) <- 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 Maybe (Entry s c a)
mparent' of
        Maybe (Entry s c a)
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ (forall a. a -> Maybe a
Just Entry s c a
parent, Bool
del)
        Just Entry s c a
parent' -> 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 :: Entry s c a
entryParent = Entry s c a
parent'} forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return (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
  Maybe (Entry s c a)
mentry <- 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 Maybe (Entry s c a)
mentry of -- check whether there is an entry
    Maybe (Entry s c a)
Nothing -> 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
      (Maybe (Entry s c a)
mrepr,Bool
del) <- 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 Bool
del -- check whether equivalence class was deleted
        then 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 so, create a new entry
        else case Maybe (Entry s c a)
mrepr of
               Maybe (Entry s c a)
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return Entry s c a
entry
               Just Entry s c a
repr -> 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 s c a
entry <- forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s (Entry s c a)
p
  (Maybe (Entry s c a)
mrepr,Bool
del) <- 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 Bool
del -- check whether equivalence class was deleted
    then do a
v <- forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall s c a. EntryData s c a -> a
entryValue forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef (forall s c a. Entry s c a -> STRef s (EntryData s c a)
unentry Entry s c a
entry)
            Entry s c a
en <- 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 -- if so, create a new entry
            (Maybe (Entry s c a)
mrepr,Bool
del) <- 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 Bool
del then do
                Entry s c a
en' <- 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 s c a
en
                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'
                forall (m :: * -> *) a. Monad m => a -> m a
return Entry s c a
en'
              else forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. a -> Maybe a -> a
fromMaybe Entry s c a
en Maybe (Entry s c a)
mrepr)
    else forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. a -> Maybe a -> a
fromMaybe Entry s c a
entry Maybe (Entry s c a)
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) = forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s (EntryData s c a)
e forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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 forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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
  STRef s (EntryData s c a)
e <- 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 :: Entry s c a
entry = forall s c a. STRef s (EntryData s c a) -> Entry s c a
Entry STRef s (EntryData s c a)
e
  Map a (Entry s c a)
m <- forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef Entries s c a
mref
  forall (m :: * -> *) s a.
Applicative m =>
STRef s a -> a -> STT s m ()
writeSTRef Entries s c a
mref (forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert a
val Entry s c a
entry Map a (Entry s c a)
m)
  forall (m :: * -> *) a. Monad m => a -> m a
return Entry s c a
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
  Entry s c a
en <- (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)
  forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall s c a. STRef s (Entry s c a) -> Class s c a
Class forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a s. Applicative m => a -> STT s m (STRef s a)
newSTRef Entry s c a
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
  Maybe (Entry s c a)
mentry <- 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 Maybe (Entry s c a)
mentry of
    Maybe (Entry s c a)
Nothing -> 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 -> 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
  Map a (Entry s c a)
m <- forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef Entries s c a
mref
  case forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
val Map a (Entry s c a)
m of
    Maybe (Entry s c a)
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
    Just Entry s c a
entry -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just Entry s c a
entry



{-| 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 forall a. Eq a => a -> a -> Bool
/= STRef s (EntryData s c a)
ry) then do
    EntryData s c a
dx <- forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s (EntryData s c a)
rx
    EntryData s c a
dy <- forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s (EntryData s c a)
ry
    case (EntryData s c a
dx, EntryData s c a
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 forall a. Ord a => a -> a -> Bool
>= Int
wy
        then do
          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}
          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 :: Int
entryWeight = Int
wx forall a. Num a => a -> a -> a
+ Int
wy, entryDesc :: c
entryDesc = c -> c -> c
mkDesc c
chx c
chy}
          forall (m :: * -> *) a. Monad m => a -> m a
return Entry s c a
repx
        else do
          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}
          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 :: Int
entryWeight = Int
wx forall a. Num a => a -> a -> a
+ Int
wy, entryDesc :: c
entryDesc = c -> c -> c
mkDesc c
chx c
chy}
          forall (m :: * -> *) a. Monad m => a -> m a
return Entry s c a
repy

      (EntryData s c a, EntryData 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 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)
_ = 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
  Entry s c a
er <- b -> STT s m (Entry s c a)
rep b
e
  Entry s c a -> [b] -> STT s m ()
run Entry s c a
er [b]
es
    where run :: Entry s c a -> [b] -> STT s m ()
run Entry s c a
er (b
f:[b]
r) = do
            Entry s c a
fr <- b -> STT s m (Entry s c a)
rep b
f
            Entry s c a
er' <- 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 s c a
eq Entry s c a
er Entry s c a
fr
            Entry s c a -> [b] -> STT s m ()
run Entry s c a
er' [b]
r
          run Entry s c a
_ [b]
_ = 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 = 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 (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 = 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] forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> 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 = 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 (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 = 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 STRef s (EntryData s c a)
e <- 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
  forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall s c a. EntryData s c a -> c
entryDesc forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s (EntryData s c a)
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 STRef s (EntryData s c a)
e <- 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
  forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall s c a. EntryData s c a -> c
entryDesc forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s (EntryData s c a)
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 STRef s (EntryData s c a)
r1) <- 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 STRef s (EntryData s c a)
r2) <- 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
c2
  forall (m :: * -> *) a. Monad m => a -> m a
return (STRef s (EntryData s c a)
r1 forall a. Eq a => a -> a -> Bool
== STRef s (EntryData s c a)
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 STRef s (EntryData s c a)
r1) <- 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 STRef s (EntryData s c a)
r2) <- 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
v2
  forall (m :: * -> *) a. Monad m => a -> m a
return (STRef s (EntryData s c a)
r1 forall a. Eq a => a -> a -> Bool
== STRef s (EntryData s c a)
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 = forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s a
r forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (forall (m :: * -> *) s a.
Applicative m =>
STRef s a -> a -> STT s m ()
writeSTRef STRef s a
r 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) = forall (m :: * -> *) s a.
(Monad m, Applicative m) =>
STRef s a -> (a -> a) -> STT s m ()
modifySTRef STRef s (EntryData s c a)
r 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 :: Bool
entryDeleted = Bool
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 s c a
entry <- forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s (Entry s c a)
p
  (Maybe (Entry s c a)
mrepr,Bool
del) <- 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 Bool
del then do
        a
v <- forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall s c a. EntryData s c a -> a
entryValue forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef (forall s c a. Entry s c a -> STRef s (EntryData s c a)
unentry Entry s c a
entry)
        Maybe (Entry s c a)
men <- 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 Maybe (Entry s c a)
men of
          Maybe (Entry s c a)
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
          Just Entry s c a
en -> do
            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
            (Maybe (Entry s c a)
mentry,Bool
del) <- 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 Bool
del
              then forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
              else forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Entry s c a -> STT s m ()
removeEntry (forall a. a -> Maybe a -> a
fromMaybe Entry s c a
en Maybe (Entry s c a)
mentry)
                   forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
    else forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Entry s c a -> STT s m ()
removeEntry (forall a. a -> Maybe a -> a
fromMaybe Entry s c a
entry Maybe (Entry s c a)
mrepr)
         forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
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
  Maybe (Entry s c a)
mentry <- 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 Maybe (Entry s c a)
mentry of
    Maybe (Entry s c a)
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
    Just Entry s c a
entry -> do
      (Maybe (Entry s c a)
mentry, Bool
del) <- 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 Bool
del
        then forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
        else forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Entry s c a -> STT s m ()
removeEntry (forall a. a -> Maybe a -> a
fromMaybe Entry s c a
entry Maybe (Entry s c a)
mentry)
             forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
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} = forall k a. Map k a -> [k]
Map.keys forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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
  [Entry s c a]
allEntries <- forall k a. Map k a -> [a]
Map.elems forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef Entries s c a
mref
  [Entry s c a]
rootEntries <- forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM forall {m :: * -> *} {s} {c} {a}.
Monad m =>
Entry s c a -> STT s m Bool
isRoot [Entry s c a]
allEntries
  forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall s c a. STRef s (Entry s c a) -> Class s c a
Class forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a s. Applicative m => a -> STT s m (STRef s a)
newSTRef) forall a b. (a -> b) -> a -> b
$ [Entry s c a]
rootEntries
    where
      isRoot :: Entry s c a -> STT s m Bool
isRoot Entry s c a
e = do
        EntryData s c a
x <- forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef (forall s c a. Entry s c a -> STRef s (EntryData s c a)
unentry Entry s c a
e)
        case EntryData s c a
x of
          Node {} -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
          Root {} -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True