Beef up tests a little and ...

rocky · rocky · commit 199d6e5c86b8 · 2024-05-10T18:38:06.000-04:00
minor tweaks to control flow code to make things
clearer.
diff --git a/control_flow/bb.py b/control_flow/bb.py
@@ -1,5 +1,4 @@
-#!/usr/bin/env python
-# Copyright (c) 2021, 2023 by Rocky Bernstein <rb@dustyfeet.com>
+# Copyright (c) 2021, 2023-2024 by Rocky Bernstein <rb@dustyfeet.com>
 import sys
 from xdis import next_offset
 from xdis.version_info import PYTHON_VERSION_TRIPLE, IS_PYPY
@@ -23,6 +22,7 @@
     BB_LOOP,
     BB_NOFOLLOW,
     BB_RETURN,
+    FLAG2NAME,
 )
 
 # The byte code versions we support
@@ -141,7 +141,8 @@ def __repr__(self):
         else:
             exception_text = ""
         if len(self.flags) > 0:
-            flag_text = f", flags={sorted(self.flags)}"
+            flag_str = ",".join(FLAG2NAME[flag] for flag in sorted(self.flags))
+            flag_text = ", flags={%s}"% flag_str
         else:
             flag_text = ""
         return "BasicBlock(#%d range: %s%s, follow_offset=%s, edge_count=%d%s%s)" % (
diff --git a/control_flow/cfg.py b/control_flow/cfg.py
@@ -1,9 +1,11 @@
-# Copyright (c) 2021 by Rocky Bernstein <rb@dustyfeet.com>
+# Copyright (c) 2021, 2024 by Rocky Bernstein <rb@dustyfeet.com>
 #
 from operator import attrgetter
+from typing import Dict, Optional, Tuple
 from control_flow.graph import (
     DiGraph,
     Node,
+    TreeGraph,
     jump_flags,
     BB_JUMP_CONDITIONAL,
     BB_LOOP,
@@ -13,7 +15,7 @@
 )
 
 
-class ControlFlowGraph(object):
+class ControlFlowGraph:
     """
     Performs the control-flow analysis on set of basic blocks. It
     iterates over its bytecode and builds basic blocks with flag
@@ -25,13 +27,13 @@ def __init__(self, bb_mgr):
         self.block_offsets = {}
         self.seen_blocks = set()
         self.blocks = bb_mgr.bb_list
-        self.offset2block = {}
-        self.offset2block_sorted = {}
+        self.offset2block: Dict[int, Node] = {}
+        self.offset2block_sorted: Tuple[int, Node] = tuple()
         self.block_nodes = {}
         self.graph = None
         self.entry_node = None
         self.exit_node = bb_mgr.exit_block
-        self.dom_tree = None
+        self.dom_tree: Optional[TreeGraph] = None
         self.analyze(self.blocks, bb_mgr.exit_block)
 
     def analyze(self, blocks, exit_block):
@@ -42,7 +44,7 @@ def analyze(self, blocks, exit_block):
         assert (
             len(blocks) >= 2
         ), "Should have at least a start block and an exception exit block"
-        self.entry = blocks[1]
+        self.entry_node = blocks[1]
         self.build_flowgraph(blocks, exit_block)
 
     def build_flowgraph(self, blocks, exit_block):
@@ -191,7 +193,7 @@ def get_node(self, offset: int) -> Node:
         block = self.offset2block_sorted[0]
 
         # FIXME: use binary search
-        for next_offset, block in self.offset2block_sorted:
+        for _, block in self.offset2block_sorted:
             if block.bb.start_offset <= offset <= block.bb.end_offset:
                 break
             pass
diff --git a/control_flow/dominators.py b/control_flow/dominators.py
@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-# Copyright (c) 2021, 2023 by Rocky Bernstein <rb@dustyfeet.com>
+# Copyright (c) 2021, 2023, 2024 by Rocky Bernstein <rb@dustyfeet.com>
 """
   Dominator tree
 
@@ -18,7 +18,7 @@ def __str__(self) -> str:
         return f"DominatorSet<{sorted_set}>"
 
 
-class DominatorTree(object):
+class DominatorTree:
     """
     Handles the dominator trees (dominator/post-dominator), and the
     computation of the dominance/post-dominance frontier.
@@ -27,6 +27,7 @@ class DominatorTree(object):
     def __init__(self, cfg, debug=False):
         self.cfg = cfg
         self.debug = debug
+        self.root = cfg.entry_node
         self.build()
 
     def build(self):
@@ -243,7 +244,8 @@ def build_dom_set1(node, do_pdoms, debug=False):
             node.bb.dom_set |= child.bb.dom_set
 
 
-# Note: this has to be done after calling tree
+# Note: this has to be done after calling tree()
+# which builds the dominator tree.
 def dfs_forest(t, do_pdoms):
     """
     Builds data flow graph using Depth-First search.
diff --git a/control_flow/graph.py b/control_flow/graph.py
@@ -7,6 +7,8 @@
   :copyright: (c) 2014 by Romain Gaucher (@rgaucher)
 """
 
+from typing import Optional
+
 # First or Basic block that we entered on. Usually
 # at offset 0.
 # Does this need to be a set?
@@ -310,7 +312,7 @@ def postorder_traverse1(self, node):
 def write_dot(
     name: str,
     prefix: str,
-    graph: DiGraph,
+    graph: Optional[DiGraph],
     write_png: bool = False,
     debug=True,
     dominator_info_format=False,
@@ -321,6 +323,8 @@ def write_dot(
 
       dot is converted to PNG and dumped if `write_bool` is True.
     """
+    if graph is None:
+        return
     path_safe = name.translate(name.maketrans(" <>", "_[]"))
     dot_path = f"{prefix}{path_safe}.dot"
     open(dot_path, "w").write(graph.to_dot(False, dominator_info_format))
diff --git a/pytest/example_fns.py b/pytest/example_fns.py
@@ -1,6 +1,5 @@
-def two_basic_blocks():
-    """An example of code with a two basic blocks. The body of the code
-    and one for execption-branching exit.
+def one_basic_block():
+    """An example of code with a one basic block.
     """
     return 5
 
diff --git a/pytest/test_bb.py b/pytest/test_bb.py
@@ -1,9 +1,9 @@
 """Test control_flow.bb: basic blocks and basd-block management"""
 from control_flow.bb import basic_blocks, BB_ENTRY, BB_EXIT, BB_RETURN
-from example_fns import two_basic_blocks, if_else_blocks
+from example_fns import one_basic_block, if_else_blocks
 
-debug = False
-if debug:
+DEBUG = False
+if DEBUG:
     import dis
 
 
@@ -15,7 +15,7 @@ def check_blocks(bb_list: list):
     exit_count = 0
     return_count = 0
     for bb in bb_list:
-        if debug:
+        if DEBUG:
             print(bb)
         if BB_ENTRY in bb.flags:
             entry_count += 1
@@ -34,8 +34,8 @@ def check_blocks(bb_list: list):
 def test_basic():
 
     offset2inst_index = {}
-    for fn in (two_basic_blocks, if_else_blocks):
-        if debug:
+    for fn in (one_basic_block, if_else_blocks):
+        if DEBUG:
             print(fn.__name__)
             dis.dis(fn)
             print()
diff --git a/pytest/test_cfg.py b/pytest/test_cfg.py
@@ -1,20 +1,37 @@
-"""Test control_flow.bb: basic blocks and basd-block management"""
+"""Test control_flow.bb: basic blocks and basic-block management"""
+
 from typing import Callable
 from xdis.bytecode import get_instructions_bytes
 from xdis.std import opc
 from control_flow.bb import basic_blocks
 from control_flow.cfg import ControlFlowGraph
-from control_flow.graph import write_dot
-from example_fns import two_basic_blocks, if_else_blocks
+from control_flow.graph import BB_ENTRY, write_dot
+from example_fns import one_basic_block, if_else_blocks
 
-debug = True
-if debug:
+DEBUG = True
+if DEBUG:
     import dis
 
 
-def check_cfg(fn: Callable, cfg: ControlFlowGraph):
+def check_cfg(fn: Callable, cfg: ControlFlowGraph, check_dict: dict):
+    """
+    Check validity of congtrol-flow graph `cfg`. Values in `check_dict()`
+    are used to assist.
+    """
     bytecode = fn.__code__.co_code
 
+    # Prefix used in assert failures:
+    prefix = f"In {fn.__name__}:"
+
+    assert cfg.entry_node, f"{prefix} control-flow-graph should have non-null entry"
+    assert (
+        BB_ENTRY in cfg.entry_node.flags
+    ), f"{prefix}: the root should be marked as an entry node"
+
+    assert (
+        len(cfg.blocks) == check_dict["count"]
+    ), f"{prefix} graph should have {check_dict['count']} blocks"
+
     # Check that all get_node returns the correct node
     # for all instruction offsets in bytecode
     current_block = cfg.offset2block[0]
@@ -42,7 +59,7 @@ def check_cfg(fn: Callable, cfg: ControlFlowGraph):
     # asking for each offset above
     assert all(
         (inst.offset in offset2block for inst in get_instructions_bytes(bytecode, opc))
-    )
+    ), (f"{prefix}" "all offsets should be covered by cfg.offset2block")
 
     # Assert offset originally was in offset2block or was added in cache
     assert len(offset2block) == cached_offsets + cache_diff
@@ -51,16 +68,19 @@ def check_cfg(fn: Callable, cfg: ControlFlowGraph):
 
 def test_basic():
     offset2inst_index = {}
-    for fn in (two_basic_blocks, if_else_blocks):
-        if debug:
+    for fn, check_dict in (
+        (one_basic_block, {"count": 2}),
+        (if_else_blocks, {"count": 5}),
+    ):
+        if DEBUG:
             print(fn.__name__)
             dis.dis(fn)
             print()
         bb_mgr = basic_blocks(fn.__code__, offset2inst_index)
         cfg = ControlFlowGraph(bb_mgr)
-        if debug:
+        if DEBUG:
             write_dot(fn.__name__, "/tmp/test_cfg-", cfg.graph, write_png=True)
-        check_cfg(fn, cfg)
+        check_cfg(fn, cfg, check_dict)
 
 
 if __name__ == "__main__":
diff --git a/pytest/test_dom.py b/pytest/test_dom.py
@@ -1,41 +1,58 @@
+#!/usr/bin/env python
 """Test control_flow.bb: basic blocks and basd-block management"""
-from typing import Callable
 
 # from xdis.bytecode import get_instructions_bytes
 # from xdis.std import opc
 from control_flow.bb import basic_blocks
 from control_flow.cfg import ControlFlowGraph
 from control_flow.dominators import DominatorTree, dfs_forest, build_dom_set
-from control_flow.graph import write_dot
-from example_fns import two_basic_blocks, if_else_blocks
+from control_flow.graph import BB_ENTRY, write_dot
+from example_fns import if_else_blocks, one_basic_block
 
-debug = False
-if debug:
+DEBUG = True
+if DEBUG:
     import dis
 
 
-def check_dom(fn: Callable, dt, cfg: ControlFlowGraph):
+def check_dom(
+    dom_tree: DominatorTree, check_dict: dict, fn_name: str, dead_code_count: int = 0
+):
+
+    # Prefix used in assert failures:
+    prefix = f"In {fn_name}:"
+
+    assert dom_tree.root, f"{prefix} tree should have non-null root"
+    assert (
+        BB_ENTRY in dom_tree.root.flags
+    ), f"{prefix}: the root should be marked as an entry node"
+    assert check_dict["count"] == len(dom_tree.doms) - dead_code_count
+    assert (
+        dom_tree.cfg.graph.nodes == dom_tree.root.dom_set
+        ), f"{prefix} the root node should dominate all nodes"
     return
 
 
 def test_basic():
     offset2inst_index = {}
-    for fn in (two_basic_blocks, if_else_blocks):
+    for fn, check_dict in (
+        (one_basic_block, {"count": 2}),
+        (if_else_blocks, {"count": 5}),
+    ):
         name = fn.__name__
-        if debug:
+        if DEBUG:
             print(name)
             dis.dis(fn)
             print()
         bb_mgr = basic_blocks(fn.__code__, offset2inst_index)
         cfg = ControlFlowGraph(bb_mgr)
-        if debug:
+        if DEBUG:
             write_dot(name, "/tmp/test_dom-", cfg.graph, write_png=True)
-        dt = DominatorTree(cfg)
-        cfg.dom_tree = dt.tree(False)
+        dom_tree = DominatorTree(cfg)
+        cfg.dom_tree = dom_tree.tree(False)
         dfs_forest(cfg.dom_tree, False)
         build_dom_set(cfg.dom_tree, False)
         write_dot(name, "/tmp/test_dom-dom-", cfg.dom_tree, write_png=True)
-        check_dom(fn, dt, cfg)
+        check_dom(dom_tree, check_dict, fn.__name__)
 
 
 if __name__ == "__main__":
