[pypy-commit] pypy shadowstack-perf-2: find_successors()

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: shadowstack-perf-2
Changeset: r84377:1fbf31cc03dc
Date: 2016-05-11 16:55 +0200
http://bitbucket.org/pypy/pypy/changeset/1fbf31cc03dc/

Log:	find_successors()

diff --git a/rpython/memory/gctransform/shadowcolor.py b/rpython/memory/gctransform/shadowcolor.py
--- a/rpython/memory/gctransform/shadowcolor.py
+++ b/rpython/memory/gctransform/shadowcolor.py
@@ -35,6 +35,30 @@
     return pred
 
 
+def find_successors(graph, pending_succ):
+    """Return the set of variables where one of the 'pending_succ' can
+    end up.  'block_succ' is a list of (block, var) tuples.
+    """
+    succ = set([v for block, v in pending_succ])
+
+    def add(block, v):
+        if isinstance(v, Variable):
+            if v not in succ:
+                pending_succ.append((block, v))
+                succ.add(v)
+
+    while pending_succ:
+        block, v = pending_succ.pop()
+        for op in block.operations:
+            if op.args and v is op.args[0] and is_trivial_rewrite(op):
+                add(block, op.result)
+        for link in block.exits:
+            for i, v1 in enumerate(link.args):
+                if v1 is v:
+                    add(link.target, link.target.inputargs[i])
+    return succ
+
+
 def find_interesting_variables(graph):
     # Decide which variables are "interesting" or not.  Interesting
     # variables contain at least the ones that appear in gc_push_roots
diff --git a/rpython/memory/gctransform/test/test_shadowcolor.py b/rpython/memory/gctransform/test/test_shadowcolor.py
--- a/rpython/memory/gctransform/test/test_shadowcolor.py
+++ b/rpython/memory/gctransform/test/test_shadowcolor.py
@@ -65,6 +65,36 @@
     pred = find_precessors(graph, [(graph.returnblock, graph.getreturnvar())])
     assert summary(pred) == {'a': 4, 'c': 1, 'v': 1}
 
+def test_find_successors_1():
+    def f(a, b):
+        return a + b
+    graph = make_graph(f, [int, int])
+    succ = find_successors(graph, [(graph.startblock, graph.getargs()[0])])
+    assert summary(succ) == {'a': 1}
+
+def test_find_successors_2():
+    def f(a, b):
+        if b > 10:
+            return a + b
+        else:
+            return a - b
+    graph = make_graph(f, [int, int])
+    succ = find_successors(graph, [(graph.startblock, graph.getargs()[0])])
+    assert summary(succ) == {'a': 3}
+
+def test_find_successors_3():
+    def f(a, b):
+        if b > 10:      # 'a' condition block
+            a = a + b   # 'a' input
+            while b > 100:
+                b -= 2
+        while b > 5:    # 'a' in loop header
+            b -= 2      # 'a' in loop body
+        return a * b    # 'a' in product
+    graph = make_graph(f, [int, int])
+    succ = find_successors(graph, [(graph.startblock, graph.getargs()[0])])
+    assert summary(succ) == {'a': 5}
+
 
 def test_interesting_vars_0():
     def f(a, b):