{-|
    Code for manipulation equivalence classes on index types.  An 'Equivalence'
    is an equivalence relation.  The empty equivalence relation is constructed
    over a ranges of values using 'emptyEquivalence'.  Less discerning
    equivalence relations can be obtained with 'equate' and 'equateAll'.
    The relation can be tested with 'equiv' and 'equivalent', and canonical
    representatives can be chosen with 'repr'.

    An example follows:

    > import Data.Equivalence.Persistent
    >
    > rel = equateAll [1,3,5,7,9]
    >     . equate 5 6
    >     . equate 2 4
    >     $ emptyEquivalence (1,10)
    >
    > test1 = equiv rel 3 5 -- This is True
    > test2 = equiv rel 1 6 -- This is True
    > test3 = equiv rel 4 6 -- This is False
-}
module Data.Equivalence.Persistent (
    Equivalence,
    emptyEquivalence,
    repr,
    equiv,
    equivalent,
    equate,
    equateAll
    )
    where

import Control.Concurrent.MVar
import Control.Monad
import Data.Array.Diff
import Data.IORef
import Data.List
import System.IO.Unsafe

arrayFrom :: (IArray a e, Ix i) => (i,i) -> (i -> e) -> a i e
arrayFrom rng f = array rng [ (x, f x) | x <- range rng ]

{-|
    An 'Equivalence' is an equivalence relation on a range of values of some
    index type.
-}
data Equivalence i = Equivalence {
    ranks :: DiffArray i Int,
    parents :: IORef (DiffArray i i)
    }

{-|
    'emptyEquivalence' is an equivalence relation that equates two values
    only when they are equal to each other.  It is the most discerning such
    relation possible.
-}
emptyEquivalence :: Ix i => (i, i) -> Equivalence i
emptyEquivalence is = unsafePerformIO $ do
    v <- newIORef (arrayFrom is id)
    return $ Equivalence (arrayFrom is (const 0)) v

reprHelper :: Ix i => DiffArray i i -> i -> (DiffArray i i, i)
reprHelper ps i
    | pi == i   = (ps, i)
    | otherwise = let (ps', r) = reprHelper ps pi in (ps' // [(i,r)], r)
  where pi = ps ! i

{-|
    'repr' gives a canonical representative of the equivalence class
    containing @x@.  It is chosen arbitrarily, but is always the same for a
    given equivalence relation.

    This function is slightly unsafe.  In particular, it's possible to build
    the same equivalence relation by equating values in two different orders,
    and the choice of canonical representatives will differ.  You can either
    think of a value of type 'Equivalence' as an equivalence relation together
    with a choice of canonical representatives, or you can consider this not a
    pure function.  Since 'Equivalence' is not an instance of @Eq@ and equality
    is not observable, both perspectives are valid.
-}
repr :: Ix i => Equivalence i -> i -> i
repr (Equivalence rs vps) i = unsafePerformIO $ atomicModifyIORef vps f
  where f ps = reprHelper ps (ps ! i)

{-|
    Determines if two values are equivalent under the given equivalence
    relation.
-}
equiv :: Ix i => Equivalence i -> i -> i -> Bool
equiv eq x y = repr eq x == repr eq y

{-|
    Determines if all of the given values are equivalent under the given
    equivalence relation.
-}
equivalent :: Ix i => Equivalence i -> [i] -> Bool
equivalent eq []     = True
equivalent eq (x:xs) = all (== repr eq x) (map (repr eq) xs)

{-|
    Construct the equivalence relation obtained by equating the given two
    values.  This combines equivalence classes.
-}
equate :: Ix i => i -> i -> Equivalence i -> Equivalence i
equate x y (Equivalence rs vps) = unsafePerformIO $ do
    (px, py, ps) <- atomicModifyIORef vps $ \ ps ->
        let (ps',  px) = reprHelper ps  x
            (ps'', py) = reprHelper ps' y
        in  (ps'', (px, py, ps''))
    return (go px py ps)
  where
    go px py ps
        | px == py  = Equivalence rs vps
        | rx > ry   = let ps' = ps // [(py, px)]
                      in Equivalence rs (unsafePerformIO (newIORef ps'))
        | rx < ry   = let ps' = ps // [(px, py)]
                      in Equivalence rs (unsafePerformIO (newIORef ps'))
        | otherwise = let ps' = ps // [(py, px)]
                          rs' = rs // [(px, (rx + 1))]
                      in Equivalence rs (unsafePerformIO (newIORef ps'))
      where rx = rs ! px
            ry = rs ! py

{-|
    Construct the equivalence relation obtained by equating all of the given
    values.  This combines equivalence classes.
-}
equateAll :: Ix i => [i] -> Equivalence i -> Equivalence i
equateAll []     eq = eq
equateAll (x:xs) eq = foldl' (flip (equate x)) eq xs