Control Flow Graph from Python Bytecode

Artifacts

Important: Pyodide takes time to initialize. Initialization completion is indicated by a red border around Run all button.

We previous discussed how to extract the control flow graph from the Python AST. However, the Python AST can change from version to version, with the introduction of new control flow structures. The byte code, in comparison, stays relatively stable. Hence, we can use the bytecode to recover the control flow graph too.

First, we need the following imports. The dis package gives us access to the Python disassembly, and networkx and matplotlib lets us draw.

System Imports

These are available from Pyodide, but you may wish to make sure that they are installed if you are attempting to run the program directly on the machine.

matplotlib
networkx

We also need pydot for drawing

Available Packages

These are packages that refer either to my previous posts or to pure python packages that I have compiled, and is available in the below locations. As before, install them if you need to run the program directly on the machine.

pydot-1.4.1-py2.py3-none-any.whl

Let us start by defining a few functions that we want to extract the control-flow graphs of.

Next, we define the node in the graph.

Now, we come to the main class

A small helper

We define an entry point for the graph

Next, the meat of generation. We look handle the jump instructions specifically. Everything else is a simple block.

class CFG(CFG):
    def gen(self):
        enter = self.entry()
        last = enter
        self.jump_to = {}
        self.opcodes = {}
        for i,ins in enumerate(dis.get_instructions(self.myfn)):
            byte = i * 2
            node = CFGNode(ins, byte)
            self.opcodes[byte] = node
            #print(i,ins)
            if ins.opname in ['LOAD_CONST',
                              'LOAD_FAST',
                              'STORE_FAST',
                              'COMPARE_OP',
                              'INPLACE_ADD',
                              'INPLACE_SUBTRACT',
                              'RETURN_VALUE',
                              'BINARY_MODULO',
                              'STORE_NAME',
                              'MAKE_FUNCTION',
                              'POP_BLOCK']:
                last.add_child(node)
                last = node
            elif ins.opname == 'POP_JUMP_IF_FALSE':
                #print("will jump to", ins.arg)
                self.lstadd(self.jump_to, ins.arg, node)
                node.props['jmp'] = True
                last.add_child(node)
                last = node
            elif ins.opname == 'JUMP_FORWARD':
                node.props['jmp'] = True
                self.lstadd(self.jump_to, (i+1)*2 + ins.arg, node)
                #print("will jump to", (i+1)*2 + ins.arg)
                last.add_child(node)
                last = node
            elif ins.opname == 'SETUP_LOOP':
                #print("setuploop: ", byte , ins.arg)
                last.add_child(node)
                last = node
            elif ins.opname == 'JUMP_ABSOLUTE':
                #print("will jump to", ins.arg)
                self.lstadd(self.jump_to, ins.arg, node)
                node.props['jmp'] = True
                last.add_child(node)
                last = node
            else:
                assert False

We need to fix the jumps now.

Making the graph to be displayed.

This finishes our implementation.

Show the image

Artifacts

The runnable Python source for this notebook is available here.

Contents

Artifacts