module Data.Equivalence.STT
(
Equiv
, Class
, leastEquiv
, getClass
, combine
, combineAll
, same
, desc
, remove
, equate
, equateAll
, equivalent
, classDesc
, removeClass
) where
import Control.Monad.ST.Trans
import Control.Monad
import Data.Maybe
import Data.Map (Map)
import qualified Data.Map as Map
newtype Class s c a = Class (STRef s (Entry s c a))
newtype Entry s c a = Entry {unentry :: STRef s (EntryData s c a)}
data EntryData s c a = Node {
entryParent :: Entry s c a,
entryValue :: a
}
| Root {
entryDesc :: c,
entryWeight :: Int,
entryValue :: a,
entryDeleted :: Bool
}
type Entries s c a = STRef s (Map a (Entry s c a))
data Equiv s c a = Equiv {
entries :: Entries s c a,
singleDesc :: a -> c,
combDesc :: c -> c -> c
}
leastEquiv :: Monad m
=> (a -> c)
-> (c -> c -> c)
-> STT s m (Equiv s c a)
leastEquiv mk com = do
es <- newSTRef Map.empty
return Equiv {entries = es, singleDesc = mk, combDesc = com}
representative' :: Monad m => Entry s c a -> STT s m (Maybe (Entry s c a),Bool)
representative' (Entry e) = do
ed <- readSTRef e
case ed of
Root {entryDeleted = del} -> do
return (Nothing, del)
Node {entryParent = parent} -> do
(mparent',del) <- representative' parent
case mparent' of
Nothing -> return $ (Just parent, del)
Just parent' -> writeSTRef e ed{entryParent = parent'} >> return (Just parent', del)
representative :: (Monad m, Ord a) => Equiv s c a -> a -> STT s m (Entry s c a)
representative eq v = do
mentry <- getEntry eq v
case mentry of
Nothing -> mkEntry eq v
Just entry -> do
(mrepr,del) <- representative' entry
if del
then mkEntry eq v
else case mrepr of
Nothing -> return entry
Just repr -> return repr
classRep :: (Monad m, Ord a) => Equiv s c a -> Class s c a -> STT s m (Entry s c a)
classRep eq (Class p) = do
entry <- readSTRef p
(mrepr,del) <- representative' entry
if del
then do v <- liftM entryValue $ readSTRef (unentry entry)
en <- getEntry' eq v
(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)
mkEntry' :: (Monad m, Ord a)
=> Equiv s c a -> Entry s c a
-> STT s m (Entry s c a)
mkEntry' eq (Entry e) = readSTRef e >>= mkEntry eq . entryValue
mkEntry :: (Monad m, Ord a)
=> Equiv s c a -> a
-> STT s m (Entry s c a)
mkEntry Equiv {entries = mref, singleDesc = mkDesc} val = do
e <- newSTRef Root
{ entryDesc = mkDesc val,
entryWeight = 1,
entryValue = val,
entryDeleted = False
}
let entry = Entry e
m <- readSTRef mref
writeSTRef mref (Map.insert val entry m)
return entry
getClass :: (Monad m, Ord a) => Equiv s c a -> a -> STT s m (Class s c a)
getClass eq v = do
en <- (getEntry' eq v)
liftM Class $ newSTRef en
getEntry' :: (Monad m, Ord a) => Equiv s c a -> a -> STT s m (Entry s c a)
getEntry' eq v = do
mentry <- getEntry eq v
case mentry of
Nothing -> mkEntry eq v
Just entry -> return entry
getEntry :: (Monad m, Ord a) => Equiv s c a -> a -> STT s m (Maybe (Entry s c a))
getEntry Equiv { entries = mref} val = do
m <- readSTRef mref
case Map.lookup val m of
Nothing -> return Nothing
Just entry -> return $ Just entry
equateEntry :: (Monad 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 = mkDesc} repx@(Entry rx) repy@(Entry ry) =
if (rx /= ry) then do
dx@Root{entryWeight = wx, entryDesc = chx, entryValue = vx} <- readSTRef rx
dy@Root{entryWeight = wy, entryDesc = chy, entryValue = vy} <- readSTRef ry
if wx >= wy
then do
writeSTRef ry Node {entryParent = repx, entryValue = vy}
writeSTRef rx dx{entryWeight = wx + wy, entryDesc = mkDesc chx chy}
return repx
else do
writeSTRef rx Node {entryParent = repy, entryValue = vx}
writeSTRef ry dy{entryWeight = wx + wy, entryDesc = mkDesc chx chy}
return repy
else return repx
combineEntries :: (Monad m, Ord a)
=> Equiv s c a -> [b] -> (b -> STT s m (Entry s c a)) -> STT s m ()
combineEntries _ [] _ = return ()
combineEntries eq (e:es) rep = do
er <- rep e
run er es
where run er (f:r) = do
fr <- rep f
er' <- equateEntry eq er fr
run er' r
run _ _ = return ()
combineAll :: (Monad m, Ord a) => Equiv s c a -> [Class s c a] -> STT s m ()
combineAll eq cls = combineEntries eq cls (classRep eq)
combine :: (Monad m, Ord a) => Equiv s c a -> Class s c a -> Class s c a -> STT s m (Class s c a)
combine eq x y = combineAll eq [x,y] >> return x
equateAll :: (Monad m, Ord a) => Equiv s c a -> [a] -> STT s m ()
equateAll eq cls = combineEntries eq cls (representative eq)
equate :: (Monad m, Ord a) => Equiv s c a -> a -> a -> STT s m ()
equate eq x y = equateAll eq [x,y]
desc :: (Monad m, Ord a) => Equiv s c a -> Class s c a -> STT s m c
desc eq cl = do
Entry e <- classRep eq cl
liftM entryDesc $ readSTRef e
classDesc :: (Monad m, Ord a) => Equiv s c a -> a -> STT s m c
classDesc eq val = do
Entry e <- representative eq val
liftM entryDesc $ readSTRef e
same :: (Monad m, Ord a) => Equiv s c a -> Class s c a -> Class s c a -> STT s m Bool
same eq c1 c2 = do
(Entry r1) <- classRep eq c1
(Entry r2) <- classRep eq c2
return (r1 == r2)
equivalent :: (Monad m, Ord a) => Equiv s c a -> a -> a -> STT s m Bool
equivalent eq v1 v2 = do
(Entry r1) <- representative eq v1
(Entry r2) <- representative eq v2
return (r1 == r2)
modifySTRef :: (Monad m) => STRef s a -> (a -> a) -> STT s m ()
modifySTRef r f = readSTRef r >>= (writeSTRef r . f)
removeEntry :: (Monad m, Ord a) => Entry s c a -> STT s m ()
removeEntry (Entry r) = modifySTRef r change
where change e = e {entryDeleted = True}
remove :: (Monad m, Ord a) => Equiv s c a -> Class s c a -> STT s m Bool
remove eq (Class p) = do
entry <- readSTRef p
(mrepr,del) <- representative' entry
if del then do
v <- liftM entryValue $ readSTRef (unentry entry)
men <- getEntry eq v
case men of
Nothing -> return False
Just en -> do
writeSTRef p en
(mentry,del) <- representative' en
if del
then return False
else removeEntry (fromMaybe en mentry)
>> return True
else removeEntry (fromMaybe entry mrepr)
>> return True
removeClass :: (Monad m, Ord a) => Equiv s c a -> a -> STT s m Bool
removeClass eq v = do
mentry <- getEntry eq v
case mentry of
Nothing -> return False
Just entry -> do
(mentry, del) <- representative' entry
if del
then return False
else removeEntry (fromMaybe entry mentry)
>> return True