PythonBuddy/pypy/rpython/translator/test/test_stackcheck.py at restrictedpython · BoragoCode/PythonBuddy · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
from rpython.conftest import option
from rpython.translator.translator import TranslationContext, graphof
from rpython.translator.backendopt.all import backend_optimizations
from rpython.translator.transform import insert_ll_stackcheck
from rpython.memory.gctransform import shadowstack

def _follow_path_naive(block, cur_path, accum):
    cur_path = (cur_path, block)
    if not block.exits:
        ops = []
        while cur_path:
            block = cur_path[1]
            ops.extend(reversed(block.operations))
            cur_path = cur_path[0]
        accum.append(list(reversed(ops)))
        return
    for link in block.exits:
        _follow_path_naive(link.target, cur_path, accum)

# explodes on loops!
def paths_naive(g):
    accum = []
    _follow_path_naive(g.startblock, None, accum)
    return accum

def direct_target(spaceop):
    return spaceop.args[0].value._obj.graph.name

def direct_calls(p):
    names = []
    for spaceop in p:
        if spaceop.opname == 'direct_call':
            names.append(direct_target(spaceop))
    return names


def check(g, funcname, ignore=None):
    paths = paths_naive(g)
    relevant = []
    for p in paths:
        funcs_called = direct_calls(p)
        if funcname in funcs_called and ignore not in funcs_called:
            assert 'stack_check___' in funcs_called
            assert (funcs_called.index(funcname) >
                    funcs_called.index('stack_check___'))
            relevant.append(p)
    return relevant


class A(object):
    def __init__(self, n):
        self.n = n

def f(a):
    x = A(a.n+1)
    if x.n == 10:
        return
    f(x)

def g(n):
    f(A(n))
    return 0

def test_simple():
    t = TranslationContext()
    a = t.buildannotator()
    a.build_types(g, [int])
    a.simplify()
    t.buildrtyper().specialize()
    backend_optimizations(t)
    t.checkgraphs()
    n = insert_ll_stackcheck(t)
    t.checkgraphs()
    assert n == 1
    if option.view:
        t.view()
    check(graphof(t, f), 'f')

def test_gctransformed():
    t = TranslationContext()
    a = t.buildannotator()
    a.build_types(g, [int])
    a.simplify()
    t.buildrtyper().specialize()
    backend_optimizations(t)
    t.checkgraphs()
    n = insert_ll_stackcheck(t)
    t.checkgraphs()
    assert n == 1
    exctransf = t.getexceptiontransformer()
    f_graph = graphof(t, f)
    exctransf.create_exception_handling(f_graph)
    if option.view:
        f_graph.show()
    check(f_graph, 'f')

    class GCTransform(shadowstack.ShadowStackFrameworkGCTransformer):
        from rpython.memory.gc.generation import GenerationGC as \
                                                          GCClass
        GC_PARAMS = {}

    gctransf = GCTransform(t)
    gctransf.transform_graph(f_graph)
    if option.view:
        f_graph.show()
    relevant = check(f_graph, 'f')
    for p in relevant:
        in_between = False
        reload = 0
        for spaceop in p:
            if spaceop.opname == 'direct_call':
                target = direct_target(spaceop)
                if target == 'f':
                    in_between = False
                elif target == 'stack_check___':
                    in_between = True
            if in_between and spaceop.opname == 'gc_reload_possibly_moved':
                reload += 1

        assert reload == 0