4 Controlflow_Analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "014f2f87",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Control flow analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "211fdc54",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "So far, we have considered raw text, produced token streams with lexers, and syntax trees using parsers. The source code statements that we have captured in the abstract syntax tree are executed in sequence; this sequence is known as the _control flow_. In this paper we will consider how to represent and analyse the control flow of a program."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae23a8d3",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "We will consider the control flow only within individual functions in this lecture, such as the following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "287eb3ce",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "  public int foo(int x) {\n",
    "    int y = 0;\n",
    "\n",
    "    while(x >= 0) {\n",
    "        if(x % 2 == 0)\n",
    "            y += x;\n",
    "        x--;\n",
    "    }\n",
    "\n",
    "    return y;\n",
    "  }\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a8dcca1",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Similar to the last chapter, we are going to use [javalang](https://github.com/c2nes/javalang/) to extract syntax trees for individual functions. To do this, we re-use the helper function that defines a dummy class wrapped around the method, as the javalang parser expects a valid Java class as input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cecac22",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "import javalang"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f181ac3",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def parse_method(code):\n",
    "    class_code = \"class Dummy {\\n\" + code + \"\\n}\";\n",
    "    tokens = javalang.tokenizer.tokenize(class_code)\n",
    "    parser = javalang.parser.Parser(tokens)\n",
    "    ast = parser.parse()\n",
    "    _, node = list(ast.filter(javalang.tree.MethodDeclaration))[0]\n",
    "    return node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98169c4d",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11ff5d97",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "To recall the last chapter, let's look at the AST for this method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed4825dc",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "from graphviz import Digraph\n",
    "def print_tree(tree):\n",
    "    unique_id = 1\n",
    "    dot = Digraph()\n",
    "    for path, node in tree:\n",
    "        dot.node(str(id(node)), str(type(node)))\n",
    "        \n",
    "        for child in node.children:\n",
    "            if isinstance(child, javalang.ast.Node):\n",
    "                dot.edge(str(id(node)), str(id(child)))\n",
    "            elif type(child) == str:\n",
    "                strid = str(unique_id)\n",
    "                unique_id = unique_id + 1\n",
    "                dot.node(strid, child)\n",
    "                dot.edge(str(id(node)), strid)\n",
    "            elif type(child) == list:\n",
    "                for lc in child:\n",
    "                    dot.edge(str(id(node)), str(id(lc)))\n",
    "                 \n",
    "    return dot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20320025",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "print_tree(tree)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78ddb534",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Control flow graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92c3aa8d",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "The Control Flow Graph (CFG) models the flow of control in a program. A CFG is a directed graph $G = (N, E)$ where nodes $N$ correspond to _statements_ in the program, and edges $E$ represent where the control can flow. AST nodes that are children of statements (e.g. parameters) are not included in the CFG. (You are familiar with the relevant types of nodes since these are the roots of the statement trees you're extracting as part of the first assignment). By definition, our CFG will always contains two dedicated nodes representing an entry and exit node."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca0c3a8b",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "For reasons of convenience, we'll be using an existing graph library that provides some of the infrastructure we need when applying algorithms on the graph. The library we'll be using is [NetworkX](https://networkx.org/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12921cac",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "import networkx as nx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e98070d4",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "One of the reasons we're using NetworkX is that it can render graphs directly in our Jupyter notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f82024d",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "The first thing we need to do is define a class that represents a node in the CFG. We will store the AST node representing the statement as an attribute of a CFG node (could also use the statement tree), and add some code to produce readable labels for the CFG nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8c15a75",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "class CFGNode:\n",
    "    def __init__(self, ast_node):\n",
    "        self.ast_node = ast_node\n",
    "    \n",
    "    def __str__(self):\n",
    "        if type(self.ast_node) == str:\n",
    "            return self.ast_node\n",
    "        return str(self.ast_node.position.line - 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b1546b4",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "A control flow graph consists of a graph and dedicated start and end nodes. Since we are also going to create other graphs besides the CFG, which will only differ in the edges they contain, we define a general helper class `ProgramGraph` together with some wrapper functions to access the underlying NetworkX graph datastructure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f41f764",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "class ProgramGraph:\n",
    "    def __init__(self, graph, start, end):\n",
    "        self.graph = graph\n",
    "        self.start = start\n",
    "        self.end   = end\n",
    "        \n",
    "    def nodes(self):\n",
    "        return self.graph.nodes()\n",
    "    \n",
    "    def edges(self):\n",
    "        return self.graph.edges()\n",
    "    \n",
    "    def successors(self, node):\n",
    "        return self.graph.successors(node)\n",
    "    \n",
    "    def predecessors(self, node):\n",
    "        return self.graph.predecessors(node)\n",
    "    \n",
    "    def is_branch(self, node):\n",
    "        return self.graph.out_degree(node) > 1\n",
    "    \n",
    "    def is_merge(self, node):\n",
    "        return self.graph.in_degree(node) > 1\n",
    "    \n",
    "    def plot(self):\n",
    "        pos = nx.nx_agraph.graphviz_layout(self.graph, prog=\"dot\")\n",
    "        nx.draw_networkx_nodes(self.graph, pos=pos, node_size=800, node_color='#FFFFFF', edgecolors='#000000')\n",
    "        nx.draw_networkx_edges(self.graph, pos=pos, connectionstyle=\"arc\", arrowsize=20)\n",
    "        nx.draw_networkx_labels(self.graph, pos=pos)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d425e663",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "The process of building a CFG consists of traversing the nodes of the AST, and creating CFG nodes whenever we encounter statements, as well as the relevant edges where we need them. We'll create a class `CFGBuilder` that helps us construct a CFG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1098fbf2",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "class CFGBuilder:\n",
    "    def __init__(self, method_declaration):\n",
    "        # create graph\n",
    "        self.graph = nx.DiGraph()\n",
    "        \n",
    "        # create entry/end node\n",
    "        self.start = CFGNode(\"Start\")\n",
    "        self.end   = CFGNode(\"End\")\n",
    "        self.graph.add_node(self.start)\n",
    "        self.graph.add_node(self.end)\n",
    "        \n",
    "        # set entry as current target\n",
    "        self.frontier = [ self.start ]\n",
    "        \n",
    "        for node in method_declaration.body:\n",
    "            self.add_node(node)\n",
    "            \n",
    "        # Add edges from last nodes to end node\n",
    "        for parent in self.frontier:\n",
    "            self.graph.add_edge(parent, self.end)\n",
    "\n",
    "    def create_graph(self):\n",
    "        return ProgramGraph(self.graph, self.start, self.end)\n",
    "    \n",
    "    def add_node(self, node):\n",
    "        # TODO: implement\n",
    "        pass "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6968d35",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "We now just need to fill the method `add_node` with content. For sequential control flow, edges are added simply between successor statements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad73094f",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "  public void foo(int x) {\n",
    "    int y = x + 1;\n",
    "    System.out.println(\"y = \" + y);\n",
    "  }\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfb5b959",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Let's extend the `CFGBuilder` to handle this case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b27f4628",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "class CFGBuilder(CFGBuilder):\n",
    "\n",
    "    def add_node(self, node):\n",
    "        cfg_node = CFGNode(node)\n",
    "        self.graph.add_node(cfg_node)\n",
    "        for parent in self.frontier:\n",
    "            self.graph.add_edge(parent, cfg_node)\n",
    "        self.frontier = [ cfg_node ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96fd9e30",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d45be09",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Code isn't always sequential, so we need to consider different types of programming constructs that we might encounter; for example, we might encounter `BlockStatements`: A block statement is a single node in the AST, but it contains children as subnodes. If we don't treat blocks special, we will only get a single node in the CFG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc1828e0",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "  public void foo(int x) {\n",
    "    {\n",
    "      int y = x + 1;\n",
    "      System.out.println(\"y = \" + y);\n",
    "    }\n",
    "  }\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98e1bf29",
   "metadata": {
    "scrolled": false,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66ac6a08",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "To fix this issue, we need to override our `add_node` method to handle block statements differently. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a9cf1d7",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "class CFGBuilder(CFGBuilder):\n",
    "    def add_block_node(self, block_node):\n",
    "        for node in block_node.statements:\n",
    "            self.add_node(node)            \n",
    "\n",
    "    def add_statement_node(self, node):\n",
    "        cfg_node = CFGNode(node)\n",
    "        self.graph.add_node(cfg_node)\n",
    "        for parent in self.frontier:\n",
    "            self.graph.add_edge(parent, cfg_node)\n",
    "        self.frontier = [ cfg_node ]           \n",
    "\n",
    "    def add_node(self, node):\n",
    "        if type(node) == javalang.tree.StatementExpression or type(node) == javalang.tree.LocalVariableDeclaration:\n",
    "            self.add_statement_node(node)            \n",
    "        elif type(node) == javalang.tree.BlockStatement:\n",
    "            self.add_block_node(node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0853718d",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8292424f",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "As you can imagine, we will need to extend the `add_node` function to delegate for all possible types of AST nodes. To avoid creating a large if-else construct, and because javalang does not provide a visitor interface, we'll use Python's single dispatch to create overloaded methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b32a1a8c",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "from functools import singledispatchmethod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3896d2c",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "class CFGBuilder(CFGBuilder):\n",
    "    @singledispatchmethod\n",
    "    def add_node(self, node):\n",
    "        pass # Ignore, but could also throw NotImplementedError\n",
    "    \n",
    "    @add_node.register        \n",
    "    def add_block_node(self, block_node: javalang.tree.BlockStatement):\n",
    "        for node in block_node.statements:\n",
    "            self.add_node(node)            \n",
    "    \n",
    "    @add_node.register\n",
    "    def add_statement_node(self, node: javalang.tree.StatementExpression):\n",
    "        cfg_node = CFGNode(node)\n",
    "        self.graph.add_node(cfg_node)\n",
    "        for parent in self.frontier:\n",
    "            self.graph.add_edge(parent, cfg_node)\n",
    "        self.frontier = [ cfg_node ]\n",
    "        \n",
    "    @add_node.register\n",
    "    def add_declaration_node(self, node: javalang.tree.LocalVariableDeclaration):\n",
    "        self.add_statement_node(node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48e1a759",
   "metadata": {
    "scrolled": false,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6c0ee03",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Next let's consider branching statements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "992a736d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "  public void foo(int x) {\n",
    "    {\n",
    "      int y = x + 1;\n",
    "      if (y == 42) {\n",
    "          System.out.println(\"y = the answer\");\n",
    "      }\n",
    "    }\n",
    "  }\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bacc0139",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Our CFG builder doesn't know about branching yet, so it will treat an if statement like an individual statement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c2d442e",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "871b56ee",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Thus, we need to add a method that handles if nodes. Whenever we add such a new method, the Python single dispatch requires us to also add the `add_node` method again, so expect to see this again and again (but still better adding new branches to a large if-construct, I think)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f872fb5",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "class CFGBuilder(CFGBuilder):\n",
    "    \n",
    "    @singledispatchmethod\n",
    "    def add_node(cls, arg):\n",
    "        super().add_node(arg)\n",
    "    \n",
    "    @add_node.register\n",
    "    def add_if_node(self, if_node: javalang.tree.IfStatement):\n",
    "        cfg_node = CFGNode(if_node)\n",
    "        self.graph.add_node(cfg_node)\n",
    "\n",
    "        for parent in self.frontier:\n",
    "            self.graph.add_edge(parent, cfg_node)\n",
    "        self.frontier = [cfg_node]\n",
    "        self.add_node(if_node.then_statement)\n",
    "        \n",
    "        if if_node.else_statement:\n",
    "            current_frontier = self.frontier[:]\n",
    "            self.frontier = [cfg_node]\n",
    "            self.add_node(if_node.else_statement)\n",
    "            self.frontier.extend(current_frontier)\n",
    "        else:\n",
    "            self.frontier.append(cfg_node)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85459640",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ef208b2",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Let's also consider the case where there is an else-branch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f45898d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "  public void foo(int x) {\n",
    "    {\n",
    "      if (y == 42) {\n",
    "          System.out.println(\"y = the answer\");\n",
    "      } else {\n",
    "          y = 3;\n",
    "      }\n",
    "    }\n",
    "  }\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb59e473",
   "metadata": {
    "scrolled": false,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "508c1890",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Next up, while loops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef52aff2",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "  public void foo(int x) {\n",
    "     int y = 0;\n",
    "     while (y < x) {\n",
    "       System.out.println(\"Hallo\");\n",
    "       y += 1;\n",
    "     }\n",
    "  }\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42cc2c4f",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "class CFGBuilder(CFGBuilder):\n",
    "\n",
    "    @singledispatchmethod\n",
    "    def add_node(cls, arg):\n",
    "        super().add_node(arg)\n",
    "    \n",
    "    @add_node.register\n",
    "    def add_while_node(self, while_node: javalang.tree.WhileStatement):\n",
    "        cfg_node = CFGNode(while_node)\n",
    "        self.graph.add_node(cfg_node)\n",
    "        for parent in self.frontier:\n",
    "            self.graph.add_edge(parent, cfg_node)\n",
    "        self.frontier = [cfg_node]\n",
    "        self.add_node(while_node.body)\n",
    "        for parent in self.frontier:\n",
    "            self.graph.add_edge(parent, cfg_node)\n",
    "        self.frontier = [cfg_node]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f9168cd",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86ed7850",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Do-while loops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d9f0d42",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "  public void foo(int x) {\n",
    "     int y = 0;\n",
    "     do {\n",
    "       System.out.println(\"Hallo\");\n",
    "       y += 1;\n",
    "     } while (y < x);\n",
    "  }\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "760ebca2",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "class CFGBuilder(CFGBuilder):\n",
    "\n",
    "    @singledispatchmethod\n",
    "    def add_node(cls, arg):\n",
    "        super().add_node(arg)\n",
    "    \n",
    "    @add_node.register\n",
    "    def add_do_node(self, do_node:javalang.tree.DoStatement):\n",
    "\n",
    "        top = self.frontier[0]\n",
    "        self.add_node(do_node.body)\n",
    "        \n",
    "        cfg_node = CFGNode(do_node)\n",
    "        self.graph.add_node(cfg_node)\n",
    "        for parent in self.frontier:\n",
    "            self.graph.add_edge(parent, cfg_node)\n",
    "        \n",
    "        head = list(self.graph.successors(top))[0]\n",
    "        self.graph.add_edge(cfg_node, head)\n",
    "        self.frontier = [cfg_node]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "838d3836",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43617bbe",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "For `for`-loops we will not split up the loop header into three nodes (initialisation, condition, increment) to keep CFGs simple, but adding more nodes here would be straightforward."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4eebeb99",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "  public void foo(int x) {\n",
    "     for(int i = 0; i < 10; i++) {\n",
    "       System.out.println(\"Hallo\");\n",
    "       System.out.println(\"Hallo\");\n",
    "     }\n",
    "  }\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68015fbf",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "class CFGBuilder(CFGBuilder):\n",
    "    @singledispatchmethod\n",
    "    def add_node(cls, arg):\n",
    "        super().add_node(arg)\n",
    "    \n",
    "    @add_node.register\n",
    "    def add_for_node(self, for_node: javalang.tree.ForStatement):\n",
    "        cfg_node = CFGNode(for_node)\n",
    "        self.graph.add_node(cfg_node)\n",
    "        for parent in self.frontier:\n",
    "            self.graph.add_edge(parent, cfg_node)\n",
    "        self.frontier = [ cfg_node ]\n",
    "        self.add_node(for_node.body)\n",
    "        for loop_end in self.frontier:\n",
    "            self.graph.add_edge(loop_end, cfg_node)\n",
    "        self.frontier = [ cfg_node ]   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e85332c",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d016cbac",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Return statements need to be linked directly to end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4af597ae",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "  public boolean foo(int x) {\n",
    "     if (x == 42) {\n",
    "       return true;\n",
    "     } else if (x == 0) {\n",
    "       return false;\n",
    "     }\n",
    "     System.out.println(\"Hallo\");\n",
    "     return false;\n",
    "  }\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04b43c2e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "class CFGBuilder(CFGBuilder):\n",
    "    \n",
    "    @singledispatchmethod\n",
    "    def add_node(cls, arg):\n",
    "        super().add_node(arg)\n",
    "    \n",
    "    @add_node.register\n",
    "    def add_return_node(self, return_node: javalang.tree.ReturnStatement):\n",
    "        cfg_node = CFGNode(return_node)\n",
    "        self.graph.add_node(cfg_node)\n",
    "        for parent in self.frontier:\n",
    "            self.graph.add_edge(parent, cfg_node)\n",
    "        self.graph.add_edge(cfg_node, self.end)\n",
    "        self.frontier = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbbc1c97",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fca340c0",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Switch-statements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68ab8c53",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "  public int foo(int x) {\n",
    "     switch (x) {\n",
    "       case 1:\n",
    "         return true;\n",
    "       case 2:\n",
    "         return false;\n",
    "       default:\n",
    "         return false;\n",
    "     }\n",
    "  }\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40392af3",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "class CFGBuilder(CFGBuilder):\n",
    "    @singledispatchmethod\n",
    "    def add_node(cls, arg):\n",
    "        super().add_node(arg)\n",
    "    \n",
    "    @add_node.register\n",
    "    def add_switch_node(self, switch_node: javalang.tree.SwitchStatement):\n",
    "        cfg_node = CFGNode(switch_node)\n",
    "        self.graph.add_node(cfg_node)\n",
    "        for parent in self.frontier:\n",
    "            self.graph.add_edge(parent, cfg_node)\n",
    "        self.frontier = []\n",
    "        break_cases = []\n",
    "        for case in switch_node.cases:\n",
    "            self.frontier.append(cfg_node)\n",
    "            for statement in case.statements:\n",
    "                self.add_node(statement)\n",
    "            if isinstance(case.statements[-1], javalang.tree.BreakStatement):\n",
    "                # If last statement is a break, add edge to the end\n",
    "                break_cases.extend(self.frontier)\n",
    "                self.frontier = []\n",
    "        self.frontier.append(cfg_node)\n",
    "        self.frontier.extend(break_cases)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8ebb023",
   "metadata": {
    "scrolled": false,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80e29dfb",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "We can observe how the CFG differs depending on whether we forget (or purposely omit) a `break` statement in a switch-case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6031f6d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "  public int foo(int x) {\n",
    "     switch (x) {\n",
    "       case 1:\n",
    "         System.out.println(\"Hallo\");\n",
    "         break;\n",
    "       case 2:\n",
    "         System.out.println(\"Hallo\");\n",
    "         break;\n",
    "       case 3:\n",
    "         System.out.println(\"Hallo\");\n",
    "         break;\n",
    "     }\n",
    "     return false;\n",
    "  }\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02c9be0e",
   "metadata": {
    "scrolled": false,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c9fc353",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Here's the version with a missing `break`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4932603",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "  public int foo(int x) {\n",
    "     switch (x) {\n",
    "       case 1:\n",
    "         System.out.println(\"Hallo\");\n",
    "         break;\n",
    "       case 2:\n",
    "         System.out.println(\"Hallo\");\n",
    "         // Fall through\n",
    "       case 3:\n",
    "         System.out.println(\"Hallo\");\n",
    "         break;\n",
    "     }\n",
    "     return false;\n",
    "  }\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f43fd7cb",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e0d3552",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "There are still several aspects of Java programs we have not covered yet, such as:\n",
    "- `break` statements in loops\n",
    "- `continue` statements\n",
    "- `try`/`catch` statements\n",
    "- Exceptional control flow\n",
    "- `goto` statements\n",
    "\n",
    "However, the functionality we have now is sufficient to cover the examples we need."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cc43b34",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Basic Blocks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "468a02fd",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "A _basic block_ is a sequence of statements that must execute consecutively. Basic blocks can be used to reduce the number of nodes in a control flow graph. The control flow enters the first statement in the block and leaves the last statement, and no intermediate statements can exit the program or jump to another statement, thus we can generally also apply control flow analysis on basic blocks rather than individual statements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "695c2a5c",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "\n",
    "class BasicBlock:\n",
    "    def __init__(self):        \n",
    "        self.cfg_nodes = []\n",
    "        \n",
    "    def __str__(self):\n",
    "        return \",\".join([str(node) for node in self.cfg_nodes])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "073b48b9",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "The first step to create basic blocks is to identify _leaders_. The following nodes in the CFG are leaders:\n",
    "\n",
    "- The entry node\n",
    "- Every node that is the target of a branch\n",
    "- Every node that is the target of an implicit goto (e.g., header of a while or for loop). We can identify these nodes because they will have more than one incoming edge."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4232de3b",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def get_leaders(cfg):\n",
    "    leaders = set()\n",
    "    leaders.add(cfg.start)\n",
    "    leaders.add(cfg.end) #todo: check\n",
    "    \n",
    "    for node in cfg.nodes():\n",
    "        if cfg.is_branch(node):\n",
    "            for successor in cfg.successors(node):\n",
    "                leaders.add(successor)\n",
    "        if cfg.is_merge(node):\n",
    "            leaders.add(node)\n",
    "\n",
    "    return leaders"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f3a8389",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Each leader starts a basic block that covers the sequence of nodes until the next leader. Thus, to calculate basic blocks we do a search starting from leader blocks up to the next leader blocks, and add all these blocks to our BasicBlock datastructure. We also maintain a mapping of CFG nodes to basic blocks, such that we can then connect the basic blocks with control flow edges."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca4f96e1",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def convert_bb(cfg):\n",
    "    bbcfg = nx.DiGraph()\n",
    "    \n",
    "    leaders = get_leaders(cfg)\n",
    "    bb_mapping = {}\n",
    "    \n",
    "    for leader in leaders:\n",
    "        bb = BasicBlock()\n",
    "        stack = [leader]\n",
    "        while stack:\n",
    "            element = stack.pop()\n",
    "            if element not in bb.cfg_nodes and element not in stack:\n",
    "                bb.cfg_nodes.append(element)\n",
    "                bb_mapping[element] = bb\n",
    "                \n",
    "                if not cfg.is_branch(element):\n",
    "                    for successor in cfg.successors(element):\n",
    "                        if successor not in leaders:\n",
    "                            stack.append(successor)\n",
    "        bbcfg.add_node(bb)\n",
    "    \n",
    "    for s, t in cfg.edges():\n",
    "        bbfrom = bb_mapping[s]\n",
    "        bbto   = bb_mapping[t]\n",
    "        if bbfrom != bbto:\n",
    "            bbcfg.add_edge(bbfrom, bbto)\n",
    "\n",
    "    return ProgramGraph(bbcfg, bb_mapping[cfg.start], bb_mapping[cfg.end])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa942234",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Let's consider an example that contains some sequences of statements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7923d356",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "  public void foo(int x) {\n",
    "     int y = 0;\n",
    "     int z = x + 1;\n",
    "     while (y < x) {\n",
    "       System.out.println(\"Hallo\");\n",
    "       y += 1;\n",
    "     }\n",
    "  }\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35c1855f",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "First, here's the original CFG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cee3b65",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b18082f4",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Here's the version with basic blocks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "515da161",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "bbcfg = convert_bb(cfg)\n",
    "bbcfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80d4e005",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Note that the two edges between node 5 and the merged node 6,7 overlap."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5be58508",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Dominance"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ed8b127",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Node $a$ dominates node $b$ if every path from entry to $b$ must pass through $a$.\n",
    "\n",
    "Node $a$ dominates $b$ if and only if:\n",
    "- $a = b$ or\n",
    "- $a$ is the unique predecessor to $b$ or\n",
    "- $a$ is the unique predecessor to any of $b$’s predecessors.\n",
    "\n",
    "In other words, the entry node dominates all other nodes; every node dominates itself; and the entry node to a loop dominates all other nodes in the loop."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ad0fd06",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "The calculation of dominators is based on a fixpoint algorithm which updates a map $D$ of dominator nodes for each node in the CFG until no more changes are found:\n",
    "1. We already know that the entry node will only be dominated by itself, so the map is initialised for entry accordingly.\n",
    "2. For each other node the map is initialised with all other nodes to start with, and the sets are iteratively reduced.\n",
    "3. In every iteration of the algorithm, the dominator set is updated for each node except the entry node. For node $n$ the set is updated to $n \\cup (\\bigcap D(p))$ for all immediate predecessors $p$ of n, where $D$ defines the currently stored for $p$.\n",
    "4. If the dominator set for any of the nodes was changed, the algorithm repeats step 3. If no set was changed, a fixpoint has been reached."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5c3e316",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "class ProgramGraph(ProgramGraph):\n",
    "    def dominators(self):\n",
    "        entry = self.start\n",
    "        remaining_nodes = [node for node in self.nodes() if node != entry]\n",
    "\n",
    "        # D(entry) <- {entry}\n",
    "        dominators = {}\n",
    "        dominators[entry] = [entry]\n",
    "\n",
    "        # D(node \\ {entry}) <- nodes\n",
    "        for node in remaining_nodes:\n",
    "            dominators[node] = set(self.nodes())\n",
    "\n",
    "        changed = True\n",
    "        while changed:\n",
    "            changed = False\n",
    "\n",
    "            for node in remaining_nodes:\n",
    "                current_dominators = dominators[node]\n",
    "\n",
    "                # new_dominators = node \\cup intersection(dominators for all predecessors)\n",
    "                new_dominators = set(self.nodes())\n",
    "\n",
    "                for p in self.predecessors(node):\n",
    "                    new_dominators = new_dominators.intersection(dominators[p])\n",
    "\n",
    "                new_dominators.add(node)\n",
    "\n",
    "                if new_dominators.symmetric_difference(current_dominators):\n",
    "                    dominators[node] = new_dominators\n",
    "                    changed = True\n",
    "\n",
    "        return dominators"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a17dc71",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Let's see this algorithm in action. First, here's the CFG that is the starting point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a558134",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37764411",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "for node, dominators in cfg.dominators().items():\n",
    "    dom = \",\".join([str(d) for d in dominators])\n",
    "    print(f\"{node}: {dom}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09421aa8",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "For example, the end node is not dominated by nodes 6 and 7, since there exists a path from start to end that does not include these two nodes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fb86bc4",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "The _strict_ dominators for a CFG node $n$ are all of its dominators except $n$ itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "855432ca",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "class ProgramGraph(ProgramGraph):\n",
    "    def strict_dominators(self):\n",
    "        dominators = self.dominators()\n",
    "        for node, dom in dominators.items():\n",
    "            dom.remove(node)\n",
    "        return dominators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9458a4b",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Recreate CFG object since the method was added\n",
    "cfg = CFGBuilder(tree).create_graph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73b38bc5",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "for node, dominators in cfg.strict_dominators().items():\n",
    "    dom = \",\".join([str(d) for d in dominators])\n",
    "    print(f\"{node}: {dom}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58444426",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "The immediate dominator for a node $n$ is the last dominator of $n$ on any path from the entry to $n$. To calculate the immediate dominator for node $n$, we iterate over the strict dominators $s$ of $n$, and remove the dominators $D(s)$ from $D(n)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6eb8cfee",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "dominator_map = cfg.strict_dominators()\n",
    "queue = [ cfg.start ]\n",
    "while queue:\n",
    "    current = queue.pop()\n",
    "    for node in cfg.nodes():            \n",
    "        dominators = dominator_map[node]\n",
    "            \n",
    "        if current in dominators:\n",
    "            dominators.remove(current)\n",
    "            if not dominators:\n",
    "                print(f\"Immediate dominator of {node}: {current}\")\n",
    "                queue.append(node)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64d2144c",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Given the dominance information, we can create the _dominator tree_, which is a graph (tree) that contains the same nodes as the CFG, but a node’s parent is its immediate dominator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bad12d0",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def dominator_tree(cfg):\n",
    "    dt = nx.DiGraph()\n",
    "    dt.add_nodes_from(cfg.nodes())\n",
    "    dominator_map = cfg.strict_dominators()\n",
    "    queue = [ cfg.start ]\n",
    "    while queue:\n",
    "        current = queue.pop()\n",
    "        for node in cfg.nodes():            \n",
    "            dominators = dominator_map[node]\n",
    "            \n",
    "            if current in dominators:\n",
    "                dominators.remove(current)\n",
    "                if not dominators:\n",
    "                    dt.add_edge(current, node)\n",
    "                    queue.append(node)\n",
    "\n",
    "    return ProgramGraph(dt, cfg.start, cfg.end)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ebde388",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "dt = dominator_tree(cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d1500a7",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "dt.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "889b6ed2",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Post-Dominance"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d69ef17",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Node $b$ post-dominates $a$ if every path from $a$ to the exit node passes through $b$. The calculation of the post-dominance information uses the identical algorithm as the dominance calculation, except that, since the relation is inverted, the entire CFG is reversed before applying the algorithm. A graph is reversed by (1) inverting the direction of all its edges, and (2) swapping start and end nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f2e6034",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "class ProgramGraph(ProgramGraph):\n",
    "    def reverse(self):\n",
    "        reverse_cfg = self.graph.reverse()\n",
    "        return ProgramGraph(reverse_cfg, self.end, self.start)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81205b9b",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Here's a simple example to invert."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0b72453",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "  public boolean foo(int x) {\n",
    "     if (x == 42) {\n",
    "       return true;\n",
    "     } else if (x == 0) {\n",
    "       return false;\n",
    "     }\n",
    "     System.out.println(\"Hallo\");\n",
    "     return false;\n",
    "  }\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f6e5913",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "First we consider the original CFG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f46f120d",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1660ede",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "...and now the reversed CFG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4035a171",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "rcfg = cfg.reverse()\n",
    "rcfg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b17f882c",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "The post-dominators can now be immediately calculated by reversing the CFG, and then calculating the dominators on the resulting reversed CFG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb9b1eaa",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "class ProgramGraph(ProgramGraph):\n",
    "    def post_dominators(self):\n",
    "        rcfg = self.reverse()\n",
    "        return rcfg.dominators()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb7341eb",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "cfg = CFGBuilder(tree).create_graph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "129fc8c3",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "for node, postdominators in cfg.post_dominators().items():\n",
    "    dom = \",\".join([str(d) for d in postdominators])\n",
    "    print(f\"{node}: {dom}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54b6994b",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Similarly to the dominator tree, we can also define a postdominator-tree, in which each node has its immediate post-dominator as a parent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "664a8aa6",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def postdominator_tree(cfg):\n",
    "    rcfg = cfg.reverse()\n",
    "    return dominator_tree(rcfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd6e96ef",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "pdt = postdominator_tree(cfg)\n",
    "pdt.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "301b6c2b",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Control Dependence"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0abd127",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Intuitively, a statement $b$ in the program is control dependent on statement $a$, if $a$ directly determines whether $b$ will be executed or not.\n",
    "\n",
    "We define $b$ to be control dependent on $a$ if:\n",
    "- There exists a directed path from $a$ to $b$ in the control flow graph where all nodes (excl. $a$) are post-dominated by $b$.\n",
    "- $a$ is not post-dominated by $b$.\n",
    "- For $b$ to be control dependent on $a$, $a$ must have two exits, one of which goes through $b$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0fe7ccd",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "The algorithm to calculate control dependencies works as follows:\n",
    "\n",
    "1. Create a new graph that contains the same node as the CFG (and DT and PDT).\n",
    "2. Find $S$, a set of edges $(a,b)$ in the CFG such that $b$ is not an ancestor of $a$ in the post-dominator tree.\n",
    "3. For each edge $(a,b) \\in S$, find $L$, the least common ancestor $L$ of $a$ and $b$ in the post-dominator tree.\n",
    "4. Consider each edge $(a,b) \\in S$ and its corresponding $L$. Traverse backwards in the post-dominator tree from $b$ to $L$, marking each node visited; mark $L$ only if $L = a$.\n",
    "5. Statements representing all marked nodes are control dependent on $a$; add edges to $a$ from all marked nodes.\n",
    "6. All remaining nodes that are not connected in the end are only dependent on entry\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a7f093d",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "We thus need a helper function that gives us the least common ancestor of two nodes in a graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5b45a3b",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "class ProgramGraph(ProgramGraph):\n",
    "    def transitive_successors(self, node):\n",
    "        return nx.nodes(nx.dfs_tree(self.graph, node))\n",
    "    \n",
    "    def least_common_ancestor(self, node1, node2):\n",
    "        current = node1\n",
    "        successors = self.transitive_successors(current)\n",
    "\n",
    "        while not (node1 in successors) or not (node2 in successors):\n",
    "            current = list(self.graph.predecessors(current))[0]\n",
    "            successors = self.transitive_successors(current)\n",
    "        \n",
    "        return current  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eeb3f705",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "The first step of the algorithm requires a function to help us determine if one node is an ancestor of another node. We determine this using reachability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "510f9ca5",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "class ProgramGraph(ProgramGraph):\n",
    "    def is_reachable(self, source, target):\n",
    "        successors = self.transitive_successors(source)\n",
    "        return target in successors"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53fd637d",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Whenever we have found control dependencies we add edges, and in the end some nodes will not be connected since they are executed unconditionally. To find these nodes, we can look for nodes without incoming edges."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a46dbf9",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "class ProgramGraph(ProgramGraph):\n",
    "    def in_degree(self, node):\n",
    "        return self.graph.in_degree(node)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88a12905",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "We will just connect all unconditionally executed nodes to the start node."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6936dbd8",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "The actual algorithm just puts together these steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "591c42e7",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def control_dependence_graph(cfg):\n",
    "    pdt = postdominator_tree(cfg)\n",
    "    cdg = nx.DiGraph()\n",
    "    cdg.add_nodes_from(cfg.nodes())\n",
    "    \n",
    "    # 1. Find S, a set of edges (a,b) in the CFG such that\n",
    "    # b is not an ancestor of a in the post-dominator tree.\n",
    "    edges = set()\n",
    "    for a, b in cfg.edges():\n",
    "        if not pdt.is_reachable(b, a):\n",
    "            edges.add((a, b))\n",
    "\n",
    "    for a, b in edges:\n",
    "        # 2. For each edge (a,b) in S, find L, the least common ancestor\n",
    "        # of a and b in the post-dominator tree.\n",
    "        lca = pdt.least_common_ancestor(a, b)\n",
    "\n",
    "        # Traverse backwards in the post-dominator tree from b to L,\n",
    "        # marking each node visited; mark L only if L = a.\n",
    "        # Statements representing all marked nodes are control dependent on a.\n",
    "        current = b\n",
    "        while current != lca:\n",
    "            cdg.add_edge(a, current)\n",
    "            current = list(pdt.predecessors(current))[0]\n",
    "\n",
    "        if a == lca:\n",
    "            cdg.add_edge(a, lca)\n",
    "\n",
    "    # Any nodes not connected at this point are only dependent on entry\n",
    "    entry = cfg.start\n",
    "    for node in cdg.nodes():\n",
    "        if node != entry and cdg.in_degree(node) == 0:\n",
    "            cdg.add_edge(entry, node)\n",
    "\n",
    "    return ProgramGraph(cdg, cfg.start, cfg.end)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db994d70",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(code)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe3ef007",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "We will now look at all the relevant intermediate graphs and the resulting control dependence graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a55edd3",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7620c50",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "dt = dominator_tree(cfg)\n",
    "dt.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff9c0eea",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "pdt = postdominator_tree(cfg)\n",
    "pdt.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ccb78d3",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "cdg = control_dependence_graph(cfg)\n",
    "cdg.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d17de3e7",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "source": [
    "Finally, we consider a slightly trickier example which is also included on the lecture slides, to see that the control dependence graph is actually not a tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d845cb8",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "public boolean testMe(int x, int y) {\n",
    "    if (x <= y) {\n",
    "        if (x == y) {\n",
    "            System.out.println(\"Some output\");\n",
    "        }\n",
    "        if (x > 0) {\n",
    "            if (y == 17) {\n",
    "               return true;\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "    return false;\n",
    "}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c721509",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "tree = parse_method(code)\n",
    "cfg = CFGBuilder(tree).create_graph()\n",
    "cfg.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf2be26d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "dt = dominator_tree(cfg)\n",
    "dt.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "866acd82",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "pdt = postdominator_tree(cfg)\n",
    "pdt.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57f1bf55",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "cdg = control_dependence_graph(cfg)\n",
    "cdg.plot()"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}