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_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 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 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 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, 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 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_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_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_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))
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 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_from_bytecode_loop(self):
# for x in (1, 2, 3):
# if x == 2:
# break
# continue
if sys.version_info < (3, 8):
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)
else:
label_loop_start = Label()
label_loop_exit = Label()
code = Bytecode()
code.extend((
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("JUMP_ABSOLUTE", label_loop_exit, lineno=3),
Instr("JUMP_ABSOLUTE", label_loop_start, lineno=4),
Instr("JUMP_ABSOLUTE", label_loop_start, lineno=4),
label_loop_exit,
Instr("LOAD_CONST", None, lineno=4),
Instr("RETURN_VALUE", lineno=4),
))
blocks = ControlFlowGraph.from_bytecode(code)
expected = [
[
Instr("LOAD_CONST", (1, 2, 3), lineno=1),
Instr("GET_ITER", lineno=1)
],
[Instr("FOR_ITER", blocks[6], 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[1], lineno=2),
],
[Instr("JUMP_ABSOLUTE", blocks[6], lineno=3)],
[Instr("JUMP_ABSOLUTE", blocks[1], lineno=4)],
[Instr("JUMP_ABSOLUTE", blocks[1], lineno=4)],
[
Instr("LOAD_CONST", None, lineno=4),
Instr("RETURN_VALUE", lineno=4)
],
]
self.assertBlocksEqual(blocks, *expected)
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
enum_repr = "<Compare.EQ: 2>"
expected = textwrap.dedent(f"""
block1:
LOAD_FAST 'test'
LOAD_CONST 1
COMPARE_OP {enum_repr}
POP_JUMP_IF_FALSE <block3>
-> block2
block2:
LOAD_CONST 1
RETURN_VALUE
block3:
LOAD_FAST 'test'
LOAD_CONST 2
COMPARE_OP {enum_repr}
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(f"""
block1:
L. 2 0: LOAD_FAST 'test'
1: LOAD_CONST 1
2: COMPARE_OP {enum_repr}
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 {enum_repr}
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)