def test_get_block_index(self):
blocks = ControlFlowGraph()
block0 = blocks[0]
block1 = blocks.add_block()
block2 = blocks.add_block()
self.assertEqual(blocks.get_block_index(block0), 0)
self.assertEqual(blocks.get_block_index(block1), 1)
self.assertEqual(blocks.get_block_index(block2), 2)
other_block = BasicBlock()
self.assertRaises(ValueError, blocks.get_block_index, other_block)
def test_to_bytecode(self):
# if test:
# x = 2
# x = 5
blocks = ControlFlowGraph()
blocks.add_block()
blocks.add_block()
blocks[0].extend(
[
Instr("LOAD_NAME", "test", lineno=1),
Instr("POP_JUMP_IF_FALSE", blocks[2], lineno=1),
]
)
blocks[1].extend(
[
Instr("LOAD_CONST", 5, lineno=2),
Instr("STORE_NAME", "x", lineno=2),
Instr("JUMP_FORWARD", blocks[2], lineno=2),
]
)
blocks[2].extend(
[
Instr("LOAD_CONST", 7, lineno=3),
Instr("STORE_NAME", "x", lineno=3),
Instr("LOAD_CONST", None, lineno=3),
Instr("RETURN_VALUE", lineno=3),
]
)
bytecode = blocks.to_bytecode()
label = Label()
self.assertEqual(
bytecode,
[
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),
label,
Instr("LOAD_CONST", 7, lineno=3),
Instr("STORE_NAME", "x", lineno=3),
Instr("LOAD_CONST", None, lineno=3),
Instr("RETURN_VALUE", lineno=3),
],
)
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 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_add_del_block(self):
code = ControlFlowGraph()
code[0].append(Instr("LOAD_CONST", 0))
block = code.add_block()
self.assertEqual(len(code), 2)
self.assertIs(block, code[1])
code[1].append(Instr("LOAD_CONST", 2))
self.assertBlocksEqual(code, [Instr("LOAD_CONST", 0)], [Instr("LOAD_CONST", 2)])
del code[0]
self.assertBlocksEqual(code, [Instr("LOAD_CONST", 2)])
del code[0]
self.assertEqual(len(code), 0)
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_blocks_broken_jump(self):
block = BasicBlock()
code = ControlFlowGraph()
code[0].append(Instr("JUMP_ABSOLUTE", block))
expected = textwrap.dedent("""
block1:
JUMP_ABSOLUTE <error: unknown block>
""").lstrip("\n")
self.check_dump_bytecode(code, 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, 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_eq(self):
# compare codes with multiple blocks and labels,
# Code.__eq__() renumbers labels to get equal labels
source = "x = 1 if test else 2"
code1 = disassemble(source)
code2 = disassemble(source)
self.assertEqual(code1, code2)
# Type mismatch
self.assertFalse(code1 == 1)
# argnames mismatch
cfg = ControlFlowGraph()
cfg.argnames = 10
self.assertFalse(code1 == cfg)
# instr mismatch
cfg = ControlFlowGraph()
cfg.argnames = code1.argnames
self.assertFalse(code1 == cfg)
def test_flag_inference(self):
# Check no loss of non-infered flags
code = ControlFlowGraph()
code.flags |= (CompilerFlags.NEWLOCALS
| CompilerFlags.VARARGS
| CompilerFlags.VARKEYWORDS
| CompilerFlags.NESTED
| CompilerFlags.FUTURE_GENERATOR_STOP)
code.update_flags()
for f in (
CompilerFlags.NEWLOCALS,
CompilerFlags.VARARGS,
CompilerFlags.VARKEYWORDS,
CompilerFlags.NESTED,
CompilerFlags.NOFREE,
CompilerFlags.OPTIMIZED,
CompilerFlags.FUTURE_GENERATOR_STOP,
):
self.assertTrue(bool(code.flags & f))
# Infer optimized and nofree
code = Bytecode()
flags = infer_flags(code)
self.assertTrue(bool(flags & CompilerFlags.OPTIMIZED))
self.assertTrue(bool(flags & CompilerFlags.NOFREE))
code.append(ConcreteInstr("STORE_NAME", 1))
flags = infer_flags(code)
self.assertFalse(bool(flags & CompilerFlags.OPTIMIZED))
self.assertTrue(bool(flags & CompilerFlags.NOFREE))
code.append(ConcreteInstr("STORE_DEREF", 2))
code.update_flags()
self.assertFalse(bool(code.flags & CompilerFlags.OPTIMIZED))
self.assertFalse(bool(code.flags & CompilerFlags.NOFREE))
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_constructor(self):
code = ControlFlowGraph()
self.assertEqual(code.name, "<module>")
self.assertEqual(code.filename, "<string>")
self.assertEqual(code.flags, 0)
self.assertBlocksEqual(code, [])
def test_empty_code(self):
cfg = ControlFlowGraph()
del cfg[0]
self.assertEqual(cfg.compute_stacksize(), 0)
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_repr(self):
r = repr(ControlFlowGraph())
self.assertIn("ControlFlowGraph", r)
self.assertIn("1", r)
def test_to_code(self):
# test resolution of jump labels
bytecode = ControlFlowGraph()
bytecode.first_lineno = 3
bytecode.argcount = 3
if sys.version_info > (3, 8):
bytecode.posonlyargcount = 0
bytecode.kwonlyargcount = 2
bytecode.name = "func"
bytecode.filename = "hello.py"
bytecode.flags = 0x43
bytecode.argnames = ("arg", "arg2", "arg3", "kwonly", "kwonly2")
bytecode.docstring = None
block0 = bytecode[0]
block1 = bytecode.add_block()
block2 = bytecode.add_block()
block0.extend([
Instr("LOAD_FAST", "x", lineno=4),
Instr("POP_JUMP_IF_FALSE", block2, lineno=4),
])
block1.extend([
Instr("LOAD_FAST", "arg", lineno=5),
Instr("STORE_FAST", "x", lineno=5)
])
block2.extend([
Instr("LOAD_CONST", 3, lineno=6),
Instr("STORE_FAST", "x", lineno=6),
Instr("LOAD_FAST", "x", lineno=7),
Instr("RETURN_VALUE", lineno=7),
])
if WORDCODE:
expected = (b"|\x05"
b"r\x08"
b"|\x00"
b"}\x05"
b"d\x01"
b"}\x05"
b"|\x05"
b"S\x00")
else:
expected = (b"|\x05\x00"
b"r\x0c\x00"
b"|\x00\x00"
b"}\x05\x00"
b"d\x01\x00"
b"}\x05\x00"
b"|\x05\x00"
b"S")
code = bytecode.to_code()
self.assertEqual(code.co_consts, (None, 3))
self.assertEqual(code.co_argcount, 3)
if sys.version_info > (3, 8):
self.assertEqual(code.co_posonlyargcount, 0)
self.assertEqual(code.co_kwonlyargcount, 2)
self.assertEqual(code.co_nlocals, 6)
self.assertEqual(code.co_stacksize, 1)
# FIXME: don't use hardcoded constants
self.assertEqual(code.co_flags, 0x43)
self.assertEqual(code.co_code, expected)
self.assertEqual(code.co_names, ())
self.assertEqual(code.co_varnames,
("arg", "arg2", "arg3", "kwonly", "kwonly2", "x"))
self.assertEqual(code.co_filename, "hello.py")
self.assertEqual(code.co_name, "func")
self.assertEqual(code.co_firstlineno, 3)
# verify stacksize argument is honored
explicit_stacksize = code.co_stacksize + 42
code = bytecode.to_code(stacksize=explicit_stacksize)
self.assertEqual(code.co_stacksize, explicit_stacksize)
def test_delitem(self):
cfg = ControlFlowGraph()
b = cfg.add_block()
del cfg[b]
self.assertEqual(len(cfg.get_instructions()), 0)
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 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)
