{-# 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 {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 {
      EntryData s c a -> Entry s c a
entryParent :: Entry s c a,
      EntryData s c a -> a
entryValue :: a
    }
                     | Root {
      EntryData s c a -> c
entryDesc :: c,
      EntryData s c a -> Int
entryWeight :: Int,
      entryValue :: 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.
      Equiv s c a -> Entries s c a
entries :: Entries s c a,
      -- | Constructs an equivalence class descriptor for a singleton class.
      Equiv s c a -> a -> c
singleDesc :: a -> c,
      -- | Combines the equivalence class descriptor of two classes
      --   which are meant to be combined.
      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 :: (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 <- 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
  Equiv s c a -> STT s m (Equiv s c a)
forall (m :: * -> *) a. Monad m => a -> m a
return Equiv :: forall s c a.
Entries s c a -> (a -> c) -> (c -> c -> c) -> Equiv s c a
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' :: 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 <- 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 EntryData s c a
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 (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
      (Maybe (Entry s c a)
mparent',Bool
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 Maybe (Entry s c a)
mparent' of
        Maybe (Entry s c a)
Nothing -> (Maybe (Entry s c a), Bool) -> STT s m (Maybe (Entry s c a), Bool)
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 :: Entry s c a
entryParent = Entry s c a
parent'} STT s m ()
-> STT s m (Maybe (Entry s c a), Bool)
-> STT s m (Maybe (Entry s c a), Bool)
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 (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 :: 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 <- 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 Maybe (Entry s c a)
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
      (Maybe (Entry s c a)
mrepr,Bool
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 Bool
del -- check whether equivalence class was deleted
        then 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 so, create a new entry
        else case Maybe (Entry s c a)
mrepr of
               Maybe (Entry s c a)
Nothing -> Entry s c a -> STT s m (Entry s c 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 (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 :: 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 <- 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
  (Maybe (Entry s c a)
mrepr,Bool
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 Bool
del -- check whether equivalence class was deleted
    then do a
v <- (EntryData s c a -> a) -> STT s m (EntryData s c a) -> STT s m a
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM EntryData s c a -> a
forall s c a. EntryData s c a -> a
entryValue (STT s m (EntryData s c a) -> STT s m a)
-> STT s m (EntryData s c a) -> STT s m a
forall a b. (a -> b) -> a -> b
$ 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
entry)
            Entry s c a
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 -- if so, create a new entry
            (Maybe (Entry s c a)
mrepr,Bool
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 Bool
del then do
                Entry s c a
en' <- Equiv s c a -> Entry 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 -> Entry s c a -> STT s m (Entry s c a)
mkEntry' Equiv s c a
eq Entry s c a
en
                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'
                Entry s c a -> STT s m (Entry s c a)
forall (m :: * -> *) a. Monad m => a -> m a
return Entry s c a
en'
              else Entry s c a -> STT s m (Entry s c a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Entry s c a -> Maybe (Entry s c a) -> Entry s c a
forall a. a -> Maybe a -> a
fromMaybe Entry s c a
en Maybe (Entry s c a)
mrepr)
    else Entry s c a -> STT s m (Entry s c a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Entry s c a -> Maybe (Entry s c a) -> Entry s c a
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' :: 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 (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 :: 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 <- 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 :: forall s c a. c -> Int -> a -> Bool -> EntryData s c a
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 = 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
  Map a (Entry s c a)
m <- 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
  Entries s c a -> Map a (Entry s c a) -> STT s m ()
forall (m :: * -> *) s a.
Applicative m =>
STRef s a -> a -> STT s m ()
writeSTRef Entries s c a
mref (a -> Entry s c a -> Map a (Entry s c a) -> Map a (Entry s c a)
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)
  Entry s c a -> STT s m (Entry s c a)
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 :: 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 <- (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)
  (STRef s (Entry s c a) -> Class s c a)
-> STT s m (STRef s (Entry s c a)) -> STT s m (Class s c a)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM STRef s (Entry s c a) -> Class s c a
forall s c a. STRef s (Entry s c a) -> Class s c a
Class (STT s m (STRef s (Entry s c a)) -> STT s m (Class s c a))
-> STT s m (STRef s (Entry s c a)) -> STT s m (Class s c a)
forall a b. (a -> b) -> a -> b
$ Entry s c a -> STT s m (STRef s (Entry s c a))
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' :: 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 <- 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 Maybe (Entry s c a)
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 (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 :: 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 <- 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
  case 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)
m of
    Maybe (Entry s c a)
Nothing -> Maybe (Entry s c a) -> STT s m (Maybe (Entry s c a))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Entry s c a)
forall a. Maybe a
Nothing
    Just Entry s c a
entry -> Maybe (Entry s c a) -> STT s m (Maybe (Entry s c a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Entry s c a) -> STT s m (Maybe (Entry s c a)))
-> Maybe (Entry s c a) -> STT s m (Maybe (Entry s c a))
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
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 :: 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
    EntryData s c a
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
    EntryData s c a
dy <- 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)
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 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 :: forall s c a. Entry s c a -> a -> EntryData s c a
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 :: Int
entryWeight = Int
wx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wy, entryDesc :: c
entryDesc = c -> c -> c
mkDesc c
chx c
chy}
          Entry s c a -> STT s m (Entry s c 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 :: forall s c a. Entry s c a -> a -> EntryData s c a
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 :: Int
entryWeight = Int
wx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
wy, entryDesc :: c
entryDesc = c -> c -> c
mkDesc c
chx c
chy}
          Entry s c a -> STT s m (Entry s c 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 (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 :: 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 (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' <- Equiv s c a -> Entry s c a -> Entry 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 -> 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]
_ = () -> STT s m ()
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 :: 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 :: 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 (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Class s c a -> STT s m (Class s c 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 :: 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 :: 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 :: 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 <- 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
  (EntryData s c a -> c) -> STT s m (EntryData s c a) -> STT s m c
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM EntryData s c a -> c
forall s c a. EntryData s c a -> c
entryDesc (STT s m (EntryData s c a) -> STT s m c)
-> STT s m (EntryData s c a) -> STT s m c
forall a b. (a -> b) -> a -> b
$ 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

{-| 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 :: 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 <- 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
  (EntryData s c a -> c) -> STT s m (EntryData s c a) -> STT s m c
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM EntryData s c a -> c
forall s c a. EntryData s c a -> c
entryDesc (STT s m (EntryData s c a) -> STT s m c)
-> STT s m (EntryData s c a) -> STT s m c
forall a b. (a -> b) -> a -> b
$ 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


{-| 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 :: 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) <- 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 STRef s (EntryData s c a)
r2) <- 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
c2
  Bool -> STT s m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (STRef s (EntryData s c a)
r1 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)
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 :: 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) <- 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 STRef s (EntryData s c a)
r2) <- 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
v2
  Bool -> STT s m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (STRef s (EntryData s c a)
r1 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)
r2)



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

modifySTRef :: (Monad m, Applicative m) => STRef s a -> (a -> a) -> STT s m ()
modifySTRef :: 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 (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 :: 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 :: 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 :: 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 <- 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
  (Maybe (Entry s c a)
mrepr,Bool
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 Bool
del then do
        a
v <- (EntryData s c a -> a) -> STT s m (EntryData s c a) -> STT s m a
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM EntryData s c a -> a
forall s c a. EntryData s c a -> a
entryValue (STT s m (EntryData s c a) -> STT s m a)
-> STT s m (EntryData s c a) -> STT s m a
forall a b. (a -> b) -> a -> b
$ 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
entry)
        Maybe (Entry s c a)
men <- 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 Maybe (Entry s c a)
men of
          Maybe (Entry s c a)
Nothing -> Bool -> STT s m Bool
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
            (Maybe (Entry s c a)
mentry,Bool
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 Bool
del
              then Bool -> STT s m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
              else Entry s c a -> STT s m ()
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Entry s c a -> STT s m ()
removeEntry (Entry s c a -> Maybe (Entry s c a) -> Entry s c a
forall a. a -> Maybe a -> a
fromMaybe Entry s c a
en Maybe (Entry s c a)
mentry)
                   STT s m () -> STT s m Bool -> STT s m Bool
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> STT s m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
    else Entry s c a -> STT s m ()
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Entry s c a -> STT s m ()
removeEntry (Entry s c a -> Maybe (Entry s c a) -> Entry s c a
forall a. a -> Maybe a -> a
fromMaybe Entry s c a
entry Maybe (Entry s c a)
mrepr)
         STT s m () -> STT s m Bool -> STT s m Bool
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> STT s m Bool
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 :: Equiv s c a -> a -> STT s m Bool
removeClass Equiv s c a
eq a
v = do
  Maybe (Entry s c a)
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 Maybe (Entry s c a)
mentry of
    Maybe (Entry s c a)
Nothing -> Bool -> STT s m Bool
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) <- 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 Bool
del
        then Bool -> STT s m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
        else Entry s c a -> STT s m ()
forall (m :: * -> *) a s c.
(Monad m, Applicative m, Ord a) =>
Entry s c a -> STT s m ()
removeEntry (Entry s c a -> Maybe (Entry s c a) -> Entry s c a
forall a. a -> Maybe a -> a
fromMaybe Entry s c a
entry Maybe (Entry s c a)
mentry)
             STT s m () -> STT s m Bool -> STT s m Bool
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Bool -> STT s m Bool
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 :: 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 :: 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 <- 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
  [Entry s c a]
rootEntries <- (Entry s c a -> STT s m Bool)
-> [Entry s c a] -> STT s m [Entry s c a]
forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM Entry s c a -> STT s m Bool
forall (m :: * -> *) s c a. Monad m => Entry s c a -> STT s m Bool
isRoot [Entry s c a]
allEntries
  (Entry s c a -> STT s m (Class s c a))
-> [Entry s c a] -> STT s m [Class s c a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((STRef s (Entry s c a) -> Class s c a)
-> STT s m (STRef s (Entry s c a)) -> STT s m (Class s c a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap STRef s (Entry s c a) -> Class s c a
forall s c a. STRef s (Entry s c a) -> Class s c a
Class (STT s m (STRef s (Entry s c a)) -> STT s m (Class s c a))
-> (Entry s c a -> STT s m (STRef s (Entry s c a)))
-> Entry s c a
-> STT s m (Class s c a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Entry s c a -> STT s m (STRef s (Entry s c a))
forall (m :: * -> *) a s. Applicative m => a -> STT s m (STRef s a)
newSTRef) ([Entry s c a] -> STT s m [Class s c a])
-> [Entry s c a] -> STT s m [Class s c a]
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 <- 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 EntryData s c a
x of
          Node {} -> Bool -> STT s m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
          Root {} -> Bool -> STT s m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True