\begin{code}
module UnVarGraph
    (  UnVarSet
    ,  emptyUnVarSet, mkUnVarSet, varEnvDom,
    ,  unionUnVarSet, unionUnVarSets
    ,  delUnVarSet
    ,  elemUnVarSet, isEmptyUnVarSet
    ,  UnVarGraph
    ,  emptyUnVarGraph
    ,  unionUnVarGraph, unionUnVarGraphs
    ,  completeGraph, completeBipartiteGraph
    ,  neighbors
    ,  delNode
    ) where

\end{code}
\\
\begin{code}

import Id
import VarEnv
import UniqFM
import Outputable
import Data.List
import Bag
import Unique

import qualified Data.IntSet as S

newtype UnVarSet = UnVarSet (S.IntSet)
    deriving Eq

k :: Var -> Int
k v = getKey (getUnique v)

emptyUnVarSet :: UnVarSet
emptyUnVarSet = UnVarSet S.empty

elemUnVarSet :: Var -> UnVarSet -> Bool
elemUnVarSet v (UnVarSet s) = k v `S.member` s

isEmptyUnVarSet :: UnVarSet -> Bool
isEmptyUnVarSet (UnVarSet s) = S.null s

delUnVarSet :: UnVarSet -> Var -> UnVarSet
delUnVarSet (UnVarSet s) v = UnVarSet $ k v `S.delete` s

mkUnVarSet :: [Var] -> UnVarSet
mkUnVarSet vs = UnVarSet $ S.fromList $ map k vs

varEnvDom :: VarEnv a -> UnVarSet
varEnvDom ae = UnVarSet $ ufmToSet_Directly ae

unionUnVarSet :: UnVarSet -> UnVarSet -> UnVarSet
unionUnVarSet (UnVarSet set1) (UnVarSet set2)
    = UnVarSet (set1 `S.union` set2)

unionUnVarSets :: [UnVarSet] -> UnVarSet
unionUnVarSets = foldr unionUnVarSet emptyUnVarSet

instance Outputable UnVarSet where
    ppr (UnVarSet s) = braces $
        hcat $ punctuate comma [ ppr (getUnique i) | i <- S.toList s]


data Gen  =  CBPG  UnVarSet UnVarSet
          |  CG    UnVarSet
newtype UnVarGraph = UnVarGraph (Bag Gen)

emptyUnVarGraph :: UnVarGraph
emptyUnVarGraph = UnVarGraph emptyBag

unionUnVarGraph :: UnVarGraph -> UnVarGraph -> UnVarGraph
unionUnVarGraph (UnVarGraph g1) (UnVarGraph g2)
    = UnVarGraph (g1 `unionBags` g2)

unionUnVarGraphs :: [UnVarGraph] -> UnVarGraph
unionUnVarGraphs = foldl' unionUnVarGraph emptyUnVarGraph

completeBipartiteGraph :: UnVarSet -> UnVarSet -> UnVarGraph
completeBipartiteGraph s1 s2
    = prune $ UnVarGraph $ unitBag $ CBPG s1 s2

completeGraph :: UnVarSet -> UnVarGraph
completeGraph s = prune $ UnVarGraph $ unitBag $ CG s

neighbors :: UnVarGraph -> Var -> UnVarSet
neighbors (UnVarGraph g) v
    = unionUnVarSets $ concatMap go $ bagToList g
  where
      go (CG s)        =  (if v `elemUnVarSet` s   then [s]   else [])
      go (CBPG s1 s2)  =  (if v `elemUnVarSet` s1  then [s2]  else []) ++
                          (if v `elemUnVarSet` s2  then [s1]  else [])

delNode :: UnVarGraph -> Var -> UnVarGraph
delNode (UnVarGraph g) v = prune $ UnVarGraph $ mapBag go g
  where
      go (CG s)        =  CG (s `delUnVarSet` v)
      go (CBPG s1 s2)  =  CBPG (s1 `delUnVarSet` v) (s2 `delUnVarSet` v)

prune :: UnVarGraph -> UnVarGraph
prune (UnVarGraph g) = UnVarGraph $ filterBag go g
  where
      go (CG s)        =  not (isEmptyUnVarSet s)
      go (CBPG s1 s2)  =  not (isEmptyUnVarSet s1) &&
                          not (isEmptyUnVarSet s2)

instance Outputable Gen where
    ppr (CG s)          = ppr s  <> char '²'
    ppr (CBPG s1 s2)    = ppr s1 <+> char 'x' <+> ppr s2
instance Outputable UnVarGraph where
    ppr (UnVarGraph g)  = ppr g
\end{code}