Major Call Arity rework

This patch improves the call arity analysis in various ways.

Most importantly, it enriches the analysis result information so that
when looking at a call, we do not have to make a random choice about
what side we want to take the information from. Instead we can combine
the results in a way that does not lose valuable information.

To do so, besides the incoming arities, we store remember "what can be
called with what", i.e. an undirected graph between the (interesting)
free variables of an expression. Of course it makes combining the
results a bit more tricky (especially mutual recursion), but still
doable.

The actually implemation of the graph structure is abstractly put away
in a module of its own (UnVarGraph.hs)

The implementation is geared towards efficiently representing the graphs
that we need (which can contain large complete and large complete
bipartite graphs, which would be huge in other representations). If
someone feels like designing data structures: There is surely some
speed-up to be obtained by improving that data structure.

Additionally, the analysis now takes into account that if a RHS stays a
thunk, then its calls happen only once, even if the variables the RHS is
bound to is evaluated multiple times, or is part of a recursive group.

2 files changed, 140 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
diff --git a/compiler/utils/UniqFM.lhs b/compiler/utils/UniqFM.lhs
index 52cd3dd791..a13a17c412 100644
--- a/compiler/utils/UniqFM.lhs
+++ b/compiler/utils/UniqFM.lhs
@@ -58,6 +58,7 @@ module UniqFM (
         lookupUFM, lookupUFM_Directly,
         lookupWithDefaultUFM, lookupWithDefaultUFM_Directly,
         eltsUFM, keysUFM, splitUFM,
+        ufmToSet_Directly,
         ufmToList,
         joinUFM
     ) where
@@ -69,6 +70,7 @@ import Compiler.Hoopl   hiding (Unique)
 
 import Data.Function (on)
 import qualified Data.IntMap as M
+import qualified Data.IntSet as S
 import qualified Data.Foldable as Foldable
 import qualified Data.Traversable as Traversable
 import Data.Typeable
@@ -180,6 +182,7 @@ lookupWithDefaultUFM_Directly
                 :: UniqFM elt -> elt -> Unique -> elt
 keysUFM         :: UniqFM elt -> [Unique]       -- Get the keys
 eltsUFM         :: UniqFM elt -> [elt]
+ufmToSet_Directly :: UniqFM elt -> S.IntSet
 ufmToList       :: UniqFM elt -> [(Unique, elt)]
 
 \end{code}
@@ -293,6 +296,7 @@ lookupWithDefaultUFM (UFM m) v k = M.findWithDefault v (getKey $ getUnique k) m
 lookupWithDefaultUFM_Directly (UFM m) v u = M.findWithDefault v (getKey u) m
 keysUFM (UFM m) = map getUnique $ M.keys m
 eltsUFM (UFM m) = M.elems m
+ufmToSet_Directly (UFM m) = M.keysSet m
 ufmToList (UFM m) = map (\(k, v) -> (getUnique k, v)) $ M.toList m
 
 -- Hoopl