1 files changed, 136 insertions, 0 deletions
diff --git a/compiler/utils/UnVarGraph.hs b/compiler/utils/UnVarGraph.hs
new file mode 100644
index 0000000000..228f3b5220
--- /dev/null
+++ b/compiler/utils/UnVarGraph.hs
@@ -0,0 +1,136 @@
+{-
+
+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 UnVarGraph
+    ( UnVarSet
+    , emptyUnVarSet, mkUnVarSet, varEnvDom, unionUnVarSet, unionUnVarSets
+    , delUnVarSet
+    , elemUnVarSet, isEmptyUnVarSet
+    , UnVarGraph
+    , emptyUnVarGraph
+    , unionUnVarGraph, unionUnVarGraphs
+    , completeGraph, completeBipartiteGraph
+    , neighbors
+    , delNode
+    ) where
+
+import Id
+import VarEnv
+import UniqFM
+import Outputable
+import Data.List
+import Bag
+import 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 [])
+
+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