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)