def test_return_value(self):
# return+return: remove second return
#
# def func():
# return 4
# return 5
code = Bytecode([
Instr('LOAD_CONST', 4, lineno=2),
Instr('RETURN_VALUE', lineno=2),
Instr('LOAD_CONST', 5, lineno=3),
Instr('RETURN_VALUE', lineno=3)
])
code = ControlFlowGraph.from_bytecode(code)
self.check(code, Instr('LOAD_CONST', 4, lineno=2),
Instr('RETURN_VALUE', lineno=2))
# return+return + return+return: remove second and fourth return
#
# def func():
# return 4
# return 5
# return 6
# return 7
code = Bytecode([
Instr('LOAD_CONST', 4, lineno=2),
Instr('RETURN_VALUE', lineno=2),
Instr('LOAD_CONST', 5, lineno=3),
Instr('RETURN_VALUE', lineno=3),
Instr('LOAD_CONST', 6, lineno=4),
Instr('RETURN_VALUE', lineno=4),
Instr('LOAD_CONST', 7, lineno=5),
Instr('RETURN_VALUE', lineno=5)
])
code = ControlFlowGraph.from_bytecode(code)
self.check(code, Instr('LOAD_CONST', 4, lineno=2),
Instr('RETURN_VALUE', lineno=2))
# return + JUMP_ABSOLUTE: remove JUMP_ABSOLUTE
# while 1:
# return 7
setup_loop = Label()
return_label = Label()
code = Bytecode([
setup_loop,
Instr('SETUP_LOOP', return_label, lineno=2),
Instr('LOAD_CONST', 7, lineno=3),
Instr('RETURN_VALUE', lineno=3),
Instr('JUMP_ABSOLUTE', setup_loop, lineno=3),
Instr('POP_BLOCK', lineno=3), return_label,
Instr('LOAD_CONST', None, lineno=3),
Instr('RETURN_VALUE', lineno=3)
])
code = ControlFlowGraph.from_bytecode(code)
end_loop = Label()
self.check(code, Instr('SETUP_LOOP', end_loop, lineno=2),
Instr('LOAD_CONST', 7, lineno=3),
Instr('RETURN_VALUE', lineno=3), end_loop,
Instr('LOAD_CONST', None, lineno=3),
Instr('RETURN_VALUE', lineno=3))
def test_return_value(self):
# return+return: remove second return
#
# def func():
# return 4
# return 5
code = Bytecode([Instr('LOAD_CONST', 4, lineno=2),
Instr('RETURN_VALUE', lineno=2),
Instr('LOAD_CONST', 5, lineno=3),
Instr('RETURN_VALUE', lineno=3)])
code = ControlFlowGraph.from_bytecode(code)
self.check(code,
Instr('LOAD_CONST', 4, lineno=2),
Instr('RETURN_VALUE', lineno=2))
# return+return + return+return: remove second and fourth return
#
# def func():
# return 4
# return 5
# return 6
# return 7
code = Bytecode([Instr('LOAD_CONST', 4, lineno=2),
Instr('RETURN_VALUE', lineno=2),
Instr('LOAD_CONST', 5, lineno=3),
Instr('RETURN_VALUE', lineno=3),
Instr('LOAD_CONST', 6, lineno=4),
Instr('RETURN_VALUE', lineno=4),
Instr('LOAD_CONST', 7, lineno=5),
Instr('RETURN_VALUE', lineno=5)])
code = ControlFlowGraph.from_bytecode(code)
self.check(code,
Instr('LOAD_CONST', 4, lineno=2),
Instr('RETURN_VALUE', lineno=2))
# return + JUMP_ABSOLUTE: remove JUMP_ABSOLUTE
# while 1:
# return 7
setup_loop = Label()
return_label = Label()
code = Bytecode([setup_loop,
Instr('SETUP_LOOP', return_label, lineno=2),
Instr('LOAD_CONST', 7, lineno=3),
Instr('RETURN_VALUE', lineno=3),
Instr('JUMP_ABSOLUTE', setup_loop, lineno=3),
Instr('POP_BLOCK', lineno=3),
return_label,
Instr('LOAD_CONST', None, lineno=3),
Instr('RETURN_VALUE', lineno=3)])
code = ControlFlowGraph.from_bytecode(code)
end_loop = Label()
self.check(code,
Instr('SETUP_LOOP', end_loop, lineno=2),
Instr('LOAD_CONST', 7, lineno=3),
Instr('RETURN_VALUE', lineno=3),
end_loop,
Instr('LOAD_CONST', None, lineno=3),
Instr('RETURN_VALUE', lineno=3))
def code_info(cls, code: Bytecode, *, debug_passes=()) -> PyCodeInfo[Repr]:
cfg = ControlFlowGraph.from_bytecode(code)
current = cls.empty()
run_machine(Interpreter(code.first_lineno).abs_i_cfg(cfg), current)
glob_deps = tuple(current.globals)
instrs = current.instrs
instrs = cls.pass_push_pop_inline(instrs)
instrs = list(relabel(instrs))
if Options.get('log-stack-vm'):
print('DEBUG: stack-vm'.center(20, '='))
show_instrs(instrs)
phi_pass_name = Options['phi-pass']
e = None
try:
phi_pass = {
'phi-elim-by-move': phi_elim,
'keep-phi': phi_keep
}[Options['phi-pass']]
except KeyError as ke:
e = Exception("Unknown phi pass {!r}".format(phi_pass_name))
if e is not None:
raise e
instrs = list(phi_pass(instrs))
if Options.get('log-phi'):
print('DEBUG: phi'.center(20, '='))
show_instrs(instrs)
return PyCodeInfo(code.name, tuple(glob_deps), code.argnames,
code.freevars, code.cellvars, code.filename,
code.first_lineno, code.argcount,
code.kwonlyargcount,
bool(code.flags & CompilerFlags.GENERATOR),
bool(code.flags & CompilerFlags.VARKEYWORDS),
bool(code.flags & CompilerFlags.VARARGS), instrs)
def test_from_bytecode_loop(self):
# for x in (1, 2, 3):
# if x == 2:
# break
# continue
label_loop_start = Label()
label_loop_exit = Label()
label_loop_end = Label()
code = Bytecode()
code.extend((Instr('SETUP_LOOP', label_loop_end, lineno=1),
Instr('LOAD_CONST', (1, 2, 3), lineno=1),
Instr('GET_ITER', lineno=1),
label_loop_start,
Instr('FOR_ITER', label_loop_exit, lineno=1),
Instr('STORE_NAME', 'x', lineno=1),
Instr('LOAD_NAME', 'x', lineno=2),
Instr('LOAD_CONST', 2, lineno=2),
Instr('COMPARE_OP', Compare.EQ, lineno=2),
Instr('POP_JUMP_IF_FALSE', label_loop_start, lineno=2),
Instr('BREAK_LOOP', lineno=3),
Instr('JUMP_ABSOLUTE', label_loop_start, lineno=4),
Instr('JUMP_ABSOLUTE', label_loop_start, lineno=4),
label_loop_exit,
Instr('POP_BLOCK', lineno=4),
label_loop_end,
Instr('LOAD_CONST', None, lineno=4),
Instr('RETURN_VALUE', lineno=4),
))
blocks = ControlFlowGraph.from_bytecode(code)
expected = [[Instr('SETUP_LOOP', blocks[8], lineno=1)],
[Instr('LOAD_CONST', (1, 2, 3), lineno=1),
Instr('GET_ITER', lineno=1)],
[Instr('FOR_ITER', blocks[7], lineno=1)],
[Instr('STORE_NAME', 'x', lineno=1),
Instr('LOAD_NAME', 'x', lineno=2),
Instr('LOAD_CONST', 2, lineno=2),
Instr('COMPARE_OP', Compare.EQ, lineno=2),
Instr('POP_JUMP_IF_FALSE', blocks[2], lineno=2)],
[Instr('BREAK_LOOP', lineno=3)],
[Instr('JUMP_ABSOLUTE', blocks[2], lineno=4)],
[Instr('JUMP_ABSOLUTE', blocks[2], lineno=4)],
[Instr('POP_BLOCK', lineno=4)],
[Instr('LOAD_CONST', None, lineno=4),
Instr('RETURN_VALUE', lineno=4)]]
self.assertBlocksEqual(blocks, *expected)
def check_dont_optimize(self, code):
code = ControlFlowGraph.from_bytecode(code)
noopt = code.to_bytecode()
optim = self.optimize_blocks(code)
optim = optim.to_bytecode()
self.assertEqual(optim, noopt)
def check_dont_optimize(self, code):
code = ControlFlowGraph.from_bytecode(code)
noopt = code.to_bytecode()
optim = self.optimize_blocks(code)
optim = optim.to_bytecode()
self.assertEqual(optim, noopt)
def test_setlineno(self):
# x = 7
# y = 8
# z = 9
code = Bytecode()
code.first_lineno = 3
code.extend([Instr("LOAD_CONST", 7),
Instr("STORE_NAME", 'x'),
SetLineno(4),
Instr("LOAD_CONST", 8),
Instr("STORE_NAME", 'y'),
SetLineno(5),
Instr("LOAD_CONST", 9),
Instr("STORE_NAME", 'z')])
blocks = ControlFlowGraph.from_bytecode(code)
self.assertBlocksEqual(blocks,
[Instr("LOAD_CONST", 7),
Instr("STORE_NAME", 'x'),
SetLineno(4),
Instr("LOAD_CONST", 8),
Instr("STORE_NAME", 'y'),
SetLineno(5),
Instr("LOAD_CONST", 9),
Instr("STORE_NAME", 'z')])
def test_from_bytecode(self):
bytecode = Bytecode()
label = Label()
bytecode.extend([Instr('LOAD_NAME', 'test', lineno=1),
Instr('POP_JUMP_IF_FALSE', label, lineno=1),
Instr('LOAD_CONST', 5, lineno=2),
Instr('STORE_NAME', 'x', lineno=2),
Instr('JUMP_FORWARD', label, lineno=2),
# dead code!
Instr('LOAD_CONST', 7, lineno=4),
Instr('STORE_NAME', 'x', lineno=4),
Label(), # unused label
label,
Label(), # unused label
Instr('LOAD_CONST', None, lineno=4),
Instr('RETURN_VALUE', lineno=4)])
blocks = ControlFlowGraph.from_bytecode(bytecode)
label2 = blocks[3]
self.assertBlocksEqual(blocks,
[Instr('LOAD_NAME', 'test', lineno=1),
Instr('POP_JUMP_IF_FALSE', label2, lineno=1)],
[Instr('LOAD_CONST', 5, lineno=2),
Instr('STORE_NAME', 'x', lineno=2),
Instr('JUMP_FORWARD', label2, lineno=2)],
[Instr('LOAD_CONST', 7, lineno=4),
Instr('STORE_NAME', 'x', lineno=4)],
[Instr('LOAD_CONST', None, lineno=4),
Instr('RETURN_VALUE', lineno=4)])
def test_from_bytecode(self):
bytecode = Bytecode()
label = Label()
bytecode.extend([Instr('LOAD_NAME', 'test', lineno=1),
Instr('POP_JUMP_IF_FALSE', label, lineno=1),
Instr('LOAD_CONST', 5, lineno=2),
Instr('STORE_NAME', 'x', lineno=2),
Instr('JUMP_FORWARD', label, lineno=2),
# dead code!
Instr('LOAD_CONST', 7, lineno=4),
Instr('STORE_NAME', 'x', lineno=4),
Label(), # unused label
label,
Label(), # unused label
Instr('LOAD_CONST', None, lineno=4),
Instr('RETURN_VALUE', lineno=4)])
blocks = ControlFlowGraph.from_bytecode(bytecode)
label2 = blocks[3]
self.assertBlocksEqual(blocks,
[Instr('LOAD_NAME', 'test', lineno=1),
Instr('POP_JUMP_IF_FALSE', label2, lineno=1)],
[Instr('LOAD_CONST', 5, lineno=2),
Instr('STORE_NAME', 'x', lineno=2),
Instr('JUMP_FORWARD', label2, lineno=2)],
[Instr('LOAD_CONST', 7, lineno=4),
Instr('STORE_NAME', 'x', lineno=4)],
[Instr('LOAD_CONST', None, lineno=4),
Instr('RETURN_VALUE', lineno=4)])
def test_setlineno(self):
# x = 7
# y = 8
# z = 9
code = Bytecode()
code.first_lineno = 3
code.extend([Instr("LOAD_CONST", 7),
Instr("STORE_NAME", 'x'),
SetLineno(4),
Instr("LOAD_CONST", 8),
Instr("STORE_NAME", 'y'),
SetLineno(5),
Instr("LOAD_CONST", 9),
Instr("STORE_NAME", 'z')])
blocks = ControlFlowGraph.from_bytecode(code)
self.assertBlocksEqual(blocks,
[Instr("LOAD_CONST", 7),
Instr("STORE_NAME", 'x'),
SetLineno(4),
Instr("LOAD_CONST", 8),
Instr("STORE_NAME", 'y'),
SetLineno(5),
Instr("LOAD_CONST", 9),
Instr("STORE_NAME", 'z')])
def from_bytecode(bytecode: Bytecode) -> CFG:
"""Generates a new control-flow graph from a bytecode segment.
Besides generating a node for each block in the bytecode segment, as returned by
`bytecode`'s `ControlFlowGraph` implementation, we add two artificial nodes to
the generated CFG:
- an artificial entry node, having index -1, that is guaranteed to fulfill the
property of an entry node, i.e., there is no incoming edge, and
- an artificial exit node, having index `sys.maxsize`, that is guaranteed to
fulfill the property of an exit node, i.e., there is no outgoing edge, and
that is the only such node in the graph, which is important, e.g., for graph
reversal.
The index values are chosen that they do not appear in regular graphs, thus one
can easily distinguish them from the normal nodes in the graph by checking for
their index-property's value.
:param bytecode: The bytecode segment
:return: The control-flow graph for the segment
"""
blocks = ControlFlowGraph.from_bytecode(bytecode)
cfg = CFG(blocks)
# Create the nodes and a mapping of all edges to generate
edges, nodes = CFG._create_nodes(blocks)
# Insert all edges between the previously generated nodes
CFG._create_graph(cfg, edges, nodes)
# Filter all dead-code nodes
cfg = CFG._filter_dead_code_nodes(cfg)
# Insert dummy exit and entry nodes
cfg = CFG._insert_dummy_exit_node(cfg)
cfg = CFG._insert_dummy_entry_node(cfg)
return cfg
def test_legalize(self):
code = Bytecode()
code.first_lineno = 3
code.extend(
[
Instr("LOAD_CONST", 7),
Instr("STORE_NAME", "x"),
Instr("LOAD_CONST", 8, lineno=4),
Instr("STORE_NAME", "y"),
SetLineno(5),
Instr("LOAD_CONST", 9, lineno=6),
Instr("STORE_NAME", "z"),
]
)
blocks = ControlFlowGraph.from_bytecode(code)
blocks.legalize()
self.assertBlocksEqual(
blocks,
[
Instr("LOAD_CONST", 7, lineno=3),
Instr("STORE_NAME", "x", lineno=3),
Instr("LOAD_CONST", 8, lineno=4),
Instr("STORE_NAME", "y", lineno=4),
Instr("LOAD_CONST", 9, lineno=5),
Instr("STORE_NAME", "z", lineno=5),
],
)
def optimize(self, code_obj):
bytecode = Bytecode.from_code(code_obj)
cfg = ControlFlowGraph.from_bytecode(bytecode)
self.optimize_cfg(cfg)
bytecode = cfg.to_bytecode()
code = bytecode.to_code()
return code
def test_from_bytecode_loop(self):
# for x in (1, 2, 3):
# if x == 2:
# break
# continue
label_loop_start = Label()
label_loop_exit = Label()
label_loop_end = Label()
code = Bytecode()
code.extend((
Instr("SETUP_LOOP", label_loop_end, lineno=1),
Instr("LOAD_CONST", (1, 2, 3), lineno=1),
Instr("GET_ITER", lineno=1),
label_loop_start,
Instr("FOR_ITER", label_loop_exit, lineno=1),
Instr("STORE_NAME", "x", lineno=1),
Instr("LOAD_NAME", "x", lineno=2),
Instr("LOAD_CONST", 2, lineno=2),
Instr("COMPARE_OP", Compare.EQ, lineno=2),
Instr("POP_JUMP_IF_FALSE", label_loop_start, lineno=2),
Instr("BREAK_LOOP", lineno=3),
Instr("JUMP_ABSOLUTE", label_loop_start, lineno=4),
Instr("JUMP_ABSOLUTE", label_loop_start, lineno=4),
label_loop_exit,
Instr("POP_BLOCK", lineno=4),
label_loop_end,
Instr("LOAD_CONST", None, lineno=4),
Instr("RETURN_VALUE", lineno=4),
))
blocks = ControlFlowGraph.from_bytecode(code)
expected = [
[Instr("SETUP_LOOP", blocks[8], lineno=1)],
[
Instr("LOAD_CONST", (1, 2, 3), lineno=1),
Instr("GET_ITER", lineno=1)
],
[Instr("FOR_ITER", blocks[7], lineno=1)],
[
Instr("STORE_NAME", "x", lineno=1),
Instr("LOAD_NAME", "x", lineno=2),
Instr("LOAD_CONST", 2, lineno=2),
Instr("COMPARE_OP", Compare.EQ, lineno=2),
Instr("POP_JUMP_IF_FALSE", blocks[2], lineno=2),
],
[Instr("BREAK_LOOP", lineno=3)],
[Instr("JUMP_ABSOLUTE", blocks[2], lineno=4)],
[Instr("JUMP_ABSOLUTE", blocks[2], lineno=4)],
[Instr("POP_BLOCK", lineno=4)],
[
Instr("LOAD_CONST", None, lineno=4),
Instr("RETURN_VALUE", lineno=4)
],
]
self.assertBlocksEqual(blocks, *expected)
def optimize(self, code_obj):
bytecode = Bytecode.from_code(code_obj)
cfg = ControlFlowGraph.from_bytecode(bytecode)
self._optimize(cfg)
bytecode = cfg.to_bytecode()
code = bytecode.to_code()
return code
def test_from_bytecode_loop(self):
# for x in (1, 2, 3):
# if x == 2:
# break
# continue
label_loop_start = Label()
label_loop_exit = Label()
label_loop_end = Label()
code = Bytecode()
code.extend((
Instr('SETUP_LOOP', label_loop_end, lineno=1),
Instr('LOAD_CONST', (1, 2, 3), lineno=1),
Instr('GET_ITER', lineno=1),
label_loop_start,
Instr('FOR_ITER', label_loop_exit, lineno=1),
Instr('STORE_NAME', 'x', lineno=1),
Instr('LOAD_NAME', 'x', lineno=2),
Instr('LOAD_CONST', 2, lineno=2),
Instr('COMPARE_OP', Compare.EQ, lineno=2),
Instr('POP_JUMP_IF_FALSE', label_loop_start, lineno=2),
Instr('BREAK_LOOP', lineno=3),
Instr('JUMP_ABSOLUTE', label_loop_start, lineno=4),
Instr('JUMP_ABSOLUTE', label_loop_start, lineno=4),
label_loop_exit,
Instr('POP_BLOCK', lineno=4),
label_loop_end,
Instr('LOAD_CONST', None, lineno=4),
Instr('RETURN_VALUE', lineno=4),
))
blocks = ControlFlowGraph.from_bytecode(code)
expected = [[Instr('SETUP_LOOP', blocks[8], lineno=1)],
[
Instr('LOAD_CONST', (1, 2, 3), lineno=1),
Instr('GET_ITER', lineno=1)
], [Instr('FOR_ITER', blocks[7], lineno=1)],
[
Instr('STORE_NAME', 'x', lineno=1),
Instr('LOAD_NAME', 'x', lineno=2),
Instr('LOAD_CONST', 2, lineno=2),
Instr('COMPARE_OP', Compare.EQ, lineno=2),
Instr('POP_JUMP_IF_FALSE', blocks[2], lineno=2)
], [Instr('BREAK_LOOP', lineno=3)],
[Instr('JUMP_ABSOLUTE', blocks[2], lineno=4)],
[Instr('JUMP_ABSOLUTE', blocks[2], lineno=4)],
[Instr('POP_BLOCK', lineno=4)],
[
Instr('LOAD_CONST', None, lineno=4),
Instr('RETURN_VALUE', lineno=4)
]]
self.assertBlocksEqual(blocks, *expected)
def check(self, source, function=False):
ref_code = get_code(source, function=function)
code = ConcreteBytecode.from_code(ref_code).to_code()
self.assertEqual(code, ref_code)
code = Bytecode.from_code(ref_code).to_code()
self.assertEqual(code, ref_code)
bytecode = Bytecode.from_code(ref_code)
blocks = ControlFlowGraph.from_bytecode(bytecode)
code = blocks.to_bytecode().to_code()
self.assertEqual(code, ref_code)
def check(self, source, function=False):
ref_code = get_code(source, function=function)
code = ConcreteBytecode.from_code(ref_code).to_code()
self.assertEqual(code, ref_code)
code = Bytecode.from_code(ref_code).to_code()
self.assertEqual(code, ref_code)
bytecode = Bytecode.from_code(ref_code)
blocks = ControlFlowGraph.from_bytecode(bytecode)
code = blocks.to_bytecode().to_code()
self.assertEqual(code, ref_code)
def code_info(cls, code: Bytecode) -> PyCodeInfo[Repr]:
cfg = ControlFlowGraph.from_bytecode(code)
current = cls.empty()
run_machine(Interpreter(code.first_lineno).abs_i_cfg(cfg), current)
glob_deps = tuple(current.globals)
instrs = current.instrs
instrs = current.pass_push_pop_inline(instrs)
return PyCodeInfo(code.name, tuple(glob_deps), code.argnames,
code.freevars, code.cellvars, code.filename,
code.first_lineno, code.argcount,
code.kwonlyargcount,
bool(code.flags & CompilerFlags.GENERATOR),
bool(code.flags & CompilerFlags.VARKEYWORDS),
bool(code.flags & CompilerFlags.VARARGS), instrs)
def disassemble(source, *, filename="<string>", function=False,
remove_last_return_none=False):
code = _disassemble(source, filename=filename, function=function)
blocks = ControlFlowGraph.from_bytecode(code)
if remove_last_return_none:
# drop LOAD_CONST+RETURN_VALUE to only keep 2 instructions,
# to make unit tests shorter
block = blocks[-1]
test = (block[-2].name == "LOAD_CONST"
and block[-2].arg is None
and block[-1].name == "RETURN_VALUE")
if not test:
raise ValueError("unable to find implicit RETURN_VALUE <None>: %s"
% block[-2:])
del block[-2:]
return blocks
def __init__(self, path: str, module_name='', to_import=['*']):
self.path = path
self.to_import = to_import
self.module_name = module_name
self._local_methods = []
source = open(path, 'rb')
compiled_source = compile(source.read(), path, 'exec')
self.bc = Bytecode.from_code(compiled_source)
self.cfg = ControlFlowGraph.from_bytecode(self.bc)
source.close()
self.build()
def disassemble(source, *, filename="<string>", function=False,
remove_last_return_none=False):
code = _disassemble(source, filename=filename, function=function)
blocks = ControlFlowGraph.from_bytecode(code)
if remove_last_return_none:
# drop LOAD_CONST+RETURN_VALUE to only keep 2 instructions,
# to make unit tests shorter
block = blocks[-1]
test = (block[-2].name == "LOAD_CONST"
and block[-2].arg is None
and block[-1].name == "RETURN_VALUE")
if not test:
raise ValueError("unable to find implicit RETURN_VALUE <None>: %s"
% block[-2:])
del block[-2:]
return blocks
def test_label_at_the_end(self):
label = Label()
code = Bytecode([Instr('LOAD_NAME', 'x'),
Instr('UNARY_NOT'),
Instr('POP_JUMP_IF_FALSE', label),
Instr('LOAD_CONST', 9),
Instr('STORE_NAME', 'y'),
label])
cfg = ControlFlowGraph.from_bytecode(code)
self.assertBlocksEqual(cfg,
[Instr('LOAD_NAME', 'x'),
Instr('UNARY_NOT'),
Instr('POP_JUMP_IF_FALSE', cfg[2])],
[Instr('LOAD_CONST', 9),
Instr('STORE_NAME', 'y')],
[])
def test_label_at_the_end(self):
label = Label()
code = Bytecode([Instr('LOAD_NAME', 'x'),
Instr('UNARY_NOT'),
Instr('POP_JUMP_IF_FALSE', label),
Instr('LOAD_CONST', 9),
Instr('STORE_NAME', 'y'),
label])
cfg = ControlFlowGraph.from_bytecode(code)
self.assertBlocksEqual(cfg,
[Instr('LOAD_NAME', 'x'),
Instr('UNARY_NOT'),
Instr('POP_JUMP_IF_FALSE', cfg[2])],
[Instr('LOAD_CONST', 9),
Instr('STORE_NAME', 'y')],
[])
def check(self, code, *expected):
if isinstance(code, Bytecode):
code = ControlFlowGraph.from_bytecode(code)
optimizer = peephole_opt.PeepholeOptimizer()
optimizer._optimize(code)
code = code.to_bytecode()
try:
self.assertEqual(code, expected)
except AssertionError:
print("Optimized code:")
dump_bytecode(code)
print("Expected code:")
for instr in expected:
print(instr)
raise
def check(self, code, *expected):
if isinstance(code, Bytecode):
code = ControlFlowGraph.from_bytecode(code)
optimizer = peephole_opt.PeepholeOptimizer()
optimizer.optimize_cfg(code)
code = code.to_bytecode()
try:
self.assertEqual(code, expected)
except AssertionError:
print("Optimized code:")
dump_bytecode(code)
print("Expected code:")
for instr in expected:
print(instr)
raise
def test_from_bytecode(self):
bytecode = Bytecode()
label = Label()
bytecode.extend([
Instr("LOAD_NAME", "test", lineno=1),
Instr("POP_JUMP_IF_FALSE", label, lineno=1),
Instr("LOAD_CONST", 5, lineno=2),
Instr("STORE_NAME", "x", lineno=2),
Instr("JUMP_FORWARD", label, lineno=2),
# dead code!
Instr("LOAD_CONST", 7, lineno=4),
Instr("STORE_NAME", "x", lineno=4),
Label(), # unused label
label,
Label(), # unused label
Instr("LOAD_CONST", None, lineno=4),
Instr("RETURN_VALUE", lineno=4),
])
blocks = ControlFlowGraph.from_bytecode(bytecode)
label2 = blocks[3]
self.assertBlocksEqual(
blocks,
[
Instr("LOAD_NAME", "test", lineno=1),
Instr("POP_JUMP_IF_FALSE", label2, lineno=1),
],
[
Instr("LOAD_CONST", 5, lineno=2),
Instr("STORE_NAME", "x", lineno=2),
Instr("JUMP_FORWARD", label2, lineno=2),
],
[
Instr("LOAD_CONST", 7, lineno=4),
Instr("STORE_NAME", "x", lineno=4)
],
[
Instr("LOAD_CONST", None, lineno=4),
Instr("RETURN_VALUE", lineno=4)
],
)
def __init__(self, path: str, module_name='', to_import=['*']):
self.path = path
self.to_import = to_import
self.module_name = module_name
self._local_methods = []
self.abi = None
source = open(path, 'rb')
source_src = source.read()
compiled_source = compile(source_src, path, 'exec')
ast_tree = ast.parse(source_src)
if module_name == '':
self.abi = ABI()
self.abi.visit(ast_tree)
self.bc = Bytecode.from_code(compiled_source)
self.cfg = ControlFlowGraph.from_bytecode(self.bc)
source.close()
self.build()
def test_label_at_the_end(self):
label = Label()
code = Bytecode([
Instr("LOAD_NAME", "x"),
Instr("UNARY_NOT"),
Instr("POP_JUMP_IF_FALSE", label),
Instr("LOAD_CONST", 9),
Instr("STORE_NAME", "y"),
label,
])
cfg = ControlFlowGraph.from_bytecode(code)
self.assertBlocksEqual(
cfg,
[
Instr("LOAD_NAME", "x"),
Instr("UNARY_NOT"),
Instr("POP_JUMP_IF_FALSE", cfg[2]),
],
[Instr("LOAD_CONST", 9),
Instr("STORE_NAME", "y")],
[],
)
def optimize_blocks(self, code):
if isinstance(code, Bytecode):
code = ControlFlowGraph.from_bytecode(code)
optimizer = peephole_opt.PeepholeOptimizer()
optimizer._optimize(code)
return code
def check_stack_size(self, func):
code = func.__code__
bytecode = Bytecode.from_code(code)
cfg = ControlFlowGraph.from_bytecode(bytecode)
self.assertEqual(code.co_stacksize, cfg.compute_stacksize())
def test_bytecode_blocks(self):
source = """
def func(test):
if test == 1:
return 1
elif test == 2:
return 2
return 3
"""
code = disassemble(source, function=True)
code = ControlFlowGraph.from_bytecode(code)
# without line numbers
expected = textwrap.dedent("""
block1:
LOAD_FAST 'test'
LOAD_CONST 1
COMPARE_OP <Compare.EQ: 2>
POP_JUMP_IF_FALSE <block3>
-> block2
block2:
LOAD_CONST 1
RETURN_VALUE
block3:
LOAD_FAST 'test'
LOAD_CONST 2
COMPARE_OP <Compare.EQ: 2>
POP_JUMP_IF_FALSE <block5>
-> block4
block4:
LOAD_CONST 2
RETURN_VALUE
block5:
LOAD_CONST 3
RETURN_VALUE
""").lstrip()
self.check_dump_bytecode(code, expected)
# with line numbers
expected = textwrap.dedent("""
block1:
L. 2 0: LOAD_FAST 'test'
1: LOAD_CONST 1
2: COMPARE_OP <Compare.EQ: 2>
3: POP_JUMP_IF_FALSE <block3>
-> block2
block2:
L. 3 0: LOAD_CONST 1
1: RETURN_VALUE
block3:
L. 4 0: LOAD_FAST 'test'
1: LOAD_CONST 2
2: COMPARE_OP <Compare.EQ: 2>
3: POP_JUMP_IF_FALSE <block5>
-> block4
block4:
L. 5 0: LOAD_CONST 2
1: RETURN_VALUE
block5:
L. 6 0: LOAD_CONST 3
1: RETURN_VALUE
""").lstrip()
self.check_dump_bytecode(code, expected, lineno=True)
def optimize_blocks(self, code):
if isinstance(code, Bytecode):
code = ControlFlowGraph.from_bytecode(code)
optimizer = peephole_opt.PeepholeOptimizer()
optimizer.optimize_cfg(code)
return code
def test_bytecode_blocks(self):
source = """
def func(test):
if test == 1:
return 1
elif test == 2:
return 2
return 3
"""
code = disassemble(source, function=True)
code = ControlFlowGraph.from_bytecode(code)
# without line numbers
expected = textwrap.dedent("""
block1:
LOAD_FAST 'test'
LOAD_CONST 1
COMPARE_OP <Compare.EQ: 2>
POP_JUMP_IF_FALSE <block3>
-> block2
block2:
LOAD_CONST 1
RETURN_VALUE
block3:
LOAD_FAST 'test'
LOAD_CONST 2
COMPARE_OP <Compare.EQ: 2>
POP_JUMP_IF_FALSE <block5>
-> block4
block4:
LOAD_CONST 2
RETURN_VALUE
block5:
LOAD_CONST 3
RETURN_VALUE
""").lstrip()
self.check_dump_bytecode(code, expected)
# with line numbers
expected = textwrap.dedent("""
block1:
L. 2 0: LOAD_FAST 'test'
1: LOAD_CONST 1
2: COMPARE_OP <Compare.EQ: 2>
3: POP_JUMP_IF_FALSE <block3>
-> block2
block2:
L. 3 0: LOAD_CONST 1
1: RETURN_VALUE
block3:
L. 4 0: LOAD_FAST 'test'
1: LOAD_CONST 2
2: COMPARE_OP <Compare.EQ: 2>
3: POP_JUMP_IF_FALSE <block5>
-> block4
block4:
L. 5 0: LOAD_CONST 2
1: RETURN_VALUE
block5:
L. 6 0: LOAD_CONST 3
1: RETURN_VALUE
""").lstrip()
self.check_dump_bytecode(code, expected, lineno=True)
def test_return_value(self):
# return+return: remove second return
#
# def func():
# return 4
# return 5
code = Bytecode([
Instr("LOAD_CONST", 4, lineno=2),
Instr("RETURN_VALUE", lineno=2),
Instr("LOAD_CONST", 5, lineno=3),
Instr("RETURN_VALUE", lineno=3),
])
code = ControlFlowGraph.from_bytecode(code)
self.check(code, Instr("LOAD_CONST", 4, lineno=2),
Instr("RETURN_VALUE", lineno=2))
# return+return + return+return: remove second and fourth return
#
# def func():
# return 4
# return 5
# return 6
# return 7
code = Bytecode([
Instr("LOAD_CONST", 4, lineno=2),
Instr("RETURN_VALUE", lineno=2),
Instr("LOAD_CONST", 5, lineno=3),
Instr("RETURN_VALUE", lineno=3),
Instr("LOAD_CONST", 6, lineno=4),
Instr("RETURN_VALUE", lineno=4),
Instr("LOAD_CONST", 7, lineno=5),
Instr("RETURN_VALUE", lineno=5),
])
code = ControlFlowGraph.from_bytecode(code)
self.check(code, Instr("LOAD_CONST", 4, lineno=2),
Instr("RETURN_VALUE", lineno=2))
# return + JUMP_ABSOLUTE: remove JUMP_ABSOLUTE
# while 1:
# return 7
if sys.version_info < (3, 8):
setup_loop = Label()
return_label = Label()
code = Bytecode([
setup_loop,
Instr("SETUP_LOOP", return_label, lineno=2),
Instr("LOAD_CONST", 7, lineno=3),
Instr("RETURN_VALUE", lineno=3),
Instr("JUMP_ABSOLUTE", setup_loop, lineno=3),
Instr("POP_BLOCK", lineno=3),
return_label,
Instr("LOAD_CONST", None, lineno=3),
Instr("RETURN_VALUE", lineno=3),
])
code = ControlFlowGraph.from_bytecode(code)
end_loop = Label()
self.check(
code,
Instr("SETUP_LOOP", end_loop, lineno=2),
Instr("LOAD_CONST", 7, lineno=3),
Instr("RETURN_VALUE", lineno=3),
end_loop,
Instr("LOAD_CONST", None, lineno=3),
Instr("RETURN_VALUE", lineno=3),
)
else:
setup_loop = Label()
return_label = Label()
code = Bytecode([
setup_loop,
Instr("LOAD_CONST", 7, lineno=3),
Instr("RETURN_VALUE", lineno=3),
Instr("JUMP_ABSOLUTE", setup_loop, lineno=3),
Instr("LOAD_CONST", None, lineno=3),
Instr("RETURN_VALUE", lineno=3),
])
code = ControlFlowGraph.from_bytecode(code)
self.check(code, Instr("LOAD_CONST", 7, lineno=3),
Instr("RETURN_VALUE", lineno=3))
