{-# LANGUAGE DeriveDataTypeable #-}
-- | A persistent, but not so efficient union-find structure.
module Data.UnionFind (
    UnionFind
  , unionFind
  , map
  , empty
  , size
  , equate
  , fromList
  , equateList
  , toList
  , union
  , find
  , findWithDefault
  , equiv
  ) where

import Prelude hiding (map)

import Control.Arrow
import Control.Applicative hiding (empty)

import Data.Maybe
import Data.Monoid
import qualified Data.List as L
import qualified Data.Map as M
import Data.Data


--  We store a map from every element to the representative of its equivalence
--  class which is the minimal element of the class.
newtype UnionFind a = UnionFind { unUnionFind :: M.Map a a }
  deriving( Eq, Ord, Show, Typeable )

instance Ord a => Monoid (UnionFind a) where
    mempty  = empty
    mappend = union

instance (Data a, Ord a) => Data (UnionFind a) where
     gfoldl k z (UnionFind a) = z (unionFind) `k` a

     gunfold k z c = case constrIndex c of
                         1 -> k (z (unionFind))
                         _ -> error "Data (UnionFind a): impossible"

     toConstr (UnionFind _) = con_UnionFind

     dataTypeOf _ = ty_T

con_UnionFind :: Constr
con_UnionFind = mkConstr ty_T "UnionFind" [] Prefix

ty_T :: DataType
ty_T = mkDataType "Data.UnionFind.UnionFind" [con_UnionFind]

-- | Smart constructor from a 'M.Map' to a union-find structure.
unionFind :: Ord a => M.Map a a -> UnionFind a
unionFind = fromList . M.toList

map :: (Ord a, Ord b) => (a -> b) -> UnionFind a -> UnionFind b
map f = fromList . fmap (f *** f) . toList

-- | @empty@ is the syntactic identity equivalence relation.
empty :: UnionFind a
empty = UnionFind M.empty

-- | @size uf@ returns the number of stored equalities.
size :: UnionFind a -> Int
size = M.size . unUnionFind

-- | @equate x y uf@ inserts the equality @x = y@ into @uf@.
equate :: Ord a => a -> a -> UnionFind a -> UnionFind a
equate x y (UnionFind uf)
  | x == y    = UnionFind uf
  | otherwise = UnionFind $
      case (M.lookup x uf, M.lookup y uf) of
         (Nothing, Nothing)
           | x <= y    -> M.insert y x uf
           | otherwise -> M.insert x y uf

         (Just xr, Nothing)
           | xr <= x   -> M.insert y xr uf
           | otherwise -> M.insert y x $ update xr x uf

         (Nothing, Just yr)
           | yr <= y   -> M.insert x yr uf
           | otherwise -> M.insert x y $ update yr y uf

         (Just xr, Just yr)
           | xr <= yr  -> update yr xr uf
           | otherwise -> update xr yr uf
  where
    update old_rep new_rep = M.map upd
      where
        upd rep | rep == old_rep = new_rep
                | otherwise      = rep

fromList :: Ord a => [(a,a)] -> UnionFind a
fromList = equateList empty

equateList :: Ord a => UnionFind a -> [(a,a)] -> UnionFind a
equateList = L.foldl' (flip $ uncurry equate)

toList :: UnionFind a -> [(a,a)]
toList = M.toList . unUnionFind

union :: Ord a => UnionFind a -> UnionFind a -> UnionFind a
union uf1 uf2
  | size uf1 < size uf2 = equateList uf2 $ toList uf1
  | otherwise           = equateList uf1 $ toList uf2

-- | @find x uf@ returns the representative of the equivalence class that @x@
-- belongs to in @uf@, if there is any.
find :: Ord a => a -> UnionFind a -> Maybe a
find x = M.lookup x . unUnionFind

-- | @findWithDefault def x uf@ returns the representative of the equivalence
-- class that @x@ belongs to in @uf@ or @def@, if there is no representative.
findWithDefault :: Ord a => a -> a -> UnionFind a -> a
findWithDefault def x = M.findWithDefault def x . unUnionFind

-- | @(x,y) `equiv` uf@ iff @x@ and @y@ are in the same equivalence class with
-- respect to @uf@.
equiv :: Ord a => (a,a) -> UnionFind a -> Bool
equiv (x, y) uf =
    x == y || (fromMaybe False $ (==) <$> find x uf <*> find y uf)