{-

Copyright (c) 2014 Joachim Breitner

A data structure for undirected graphs of variables
(or in plain terms: Sets of unordered pairs of numbers)


This is very specifically tailored for the use in CallArity. In particular it
stores the graph as a union of complete and complete bipartite graph, which
would be very expensive to store as sets of edges or as adjanceny lists.

It does not normalize the graphs. This means that g `unionUnVarGraph` g is
equal to g, but twice as expensive and large.

-}
module GHC.Data.Graph.UnVar
    ( UnVarSet
    , emptyUnVarSet, mkUnVarSet, varEnvDom, unionUnVarSet, unionUnVarSets
    , delUnVarSet
    , elemUnVarSet, isEmptyUnVarSet
    , UnVarGraph
    , emptyUnVarGraph
    , unionUnVarGraph, unionUnVarGraphs
    , completeGraph, completeBipartiteGraph
    , neighbors
    , hasLoopAt
    , delNode
    ) where

import GHC.Prelude

import GHC.Types.Id
import GHC.Types.Var.Env
import GHC.Types.Unique.FM
import GHC.Utils.Outputable
import GHC.Data.Bag
import GHC.Types.Unique

import qualified Data.IntSet as S

-- We need a type for sets of variables (UnVarSet).
-- We do not use VarSet, because for that we need to have the actual variable
-- at hand, and we do not have that when we turn the domain of a VarEnv into a UnVarSet.
-- Therefore, use a IntSet directly (which is likely also a bit more efficient).

-- Set of uniques, i.e. for adjancet nodes
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]


-- The graph type. A list of complete bipartite graphs
data Gen = CBPG UnVarSet UnVarSet -- complete bipartite
         | CG   UnVarSet          -- complete
newtype UnVarGraph = UnVarGraph (Bag Gen)

emptyUnVarGraph :: UnVarGraph
emptyUnVarGraph = UnVarGraph emptyBag

unionUnVarGraph :: UnVarGraph -> UnVarGraph -> UnVarGraph
{-
Premature optimisation, it seems.
unionUnVarGraph (UnVarGraph [CBPG s1 s2]) (UnVarGraph [CG s3, CG s4])
    | s1 == s3 && s2 == s4
    = pprTrace "unionUnVarGraph fired" empty $
      completeGraph (s1 `unionUnVarSet` s2)
unionUnVarGraph (UnVarGraph [CBPG s1 s2]) (UnVarGraph [CG s3, CG s4])
    | s2 == s3 && s1 == s4
    = pprTrace "unionUnVarGraph fired2" empty $
      completeGraph (s1 `unionUnVarSet` s2)
-}
unionUnVarGraph (UnVarGraph g1) (UnVarGraph g2)
    = -- pprTrace "unionUnVarGraph" (ppr (length g1, length g2)) $
      UnVarGraph (g1 `unionBags` g2)

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

-- completeBipartiteGraph A B = { {a,b} | a ∈ A, b ∈ B }
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 [])

-- hasLoopAt G v <=> v--v ∈ G
hasLoopAt :: UnVarGraph -> Var -> Bool
hasLoopAt (UnVarGraph g) v = any go $ bagToList g
  where go (CG s)       = v `elemUnVarSet` s
        go (CBPG s1 s2) = v `elemUnVarSet` s1 && v `elemUnVarSet` s2


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