pcode-graph 1.0.0__tar.gz → 1.1.0__tar.gz

{pcode_graph-1.0.0 → pcode_graph-1.1.0}/.github/workflows/release.yml

@@ -29,4 +29,4 @@ jobs:
         with:
           name: dist
           path: dist/
-      - uses: pypa/gh-action-pypi-publish@release/v1
+      - uses: pypa/gh-action-pypi-publish@release/v1

pcode_graph-1.1.0/.github/workflows/test.yml

@@ -0,0 +1,17 @@
+name: Tests
+on:
+  pull_request:
+  push:
+    branches: [main]
+  workflow_call:
+
+jobs:
+  pytest:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - run: pip install .[test,torch]
+      - run: python -m pytest tests

{pcode_graph-1.0.0 → pcode_graph-1.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: pcode_graph
-Version: 1.0.0
+Version: 1.1.0
 Summary: Build, display and vectorize semantic graphs from binary code
 Project-URL: Homepage, https://github.com/quarkslab/pcode_graph
 Project-URL: Issues, https://github.com/quarkslab/pcode_graph/issues

{pcode_graph-1.0.0 → pcode_graph-1.1.0}/pcode_graph/analysis.py

@@ -3,7 +3,7 @@ from enum import Enum
 from typing import Iterator, override
 from pypcode import OpCode, PcodeOp, Varnode
 from pcode_graph.arch import Arch
-from pcode_graph.pcode import OpCodes, dump_operation, from_native_opcode
+from pcode_graph.pcode import dump_operation, from_native_opcode
 from pcode_graph.visitor import PcodeVisitor
 
 OpIndex = int
@@ -56,7 +56,8 @@ class Var:
 def wrap(var: Varnode) -> Var:
     match var.space.name:
         case "register":
-            return Var(VarKinds.Register, var.size, var.offset, var.getRegisterName())
+            name = var.getRegisterName() or "UNKNOWN"
+            return Var(VarKinds.Register, var.size, var.offset, name)
         case "unique":
             return Var(VarKinds.Unique, var.size, var.offset, f"${var.offset}")
         case "const":

{pcode_graph-1.0.0 → pcode_graph-1.1.0}/pcode_graph/charac.py

@@ -11,13 +11,27 @@ from pcode_graph.pcode import OpCodes
 from collections import deque
 from itertools import chain
 
-class Diameters(NamedTuple):
+
+class Stats(NamedTuple):
     data_only: int
     control_only: int
     both: int
 
 
-def get_graph_diameters(cdg: CDG) -> Diameters:
+def count_edges(cdg: CDG) -> Stats:
+
+    data_only: int = 0
+
+    for edge in cdg.edges:
+        if edge.kind == EdgeKinds.Data:
+            data_only += 1
+
+    return Stats(data_only, len(cdg.edges) - data_only, len(cdg.edges))
+
+
+def get_graph_diameters(cdg: CDG) -> Stats:
+    """Computes the diameter of the given graph."""
+
     n = len(cdg.nodes)
 
     def build_adjacency(edges):
@@ -47,17 +61,19 @@ def get_graph_diameters(cdg: CDG) -> Diameters:
                     if dist[neigbor] == -1:
                         dist[neigbor] = dist[current] + 1
                         front.append(neigbor)
-            d = max(d, max(dist))            
+            d = max(d, max(dist))
         return d
 
-    return Diameters(
+    return Stats(
         diameter(build_adjacency(data_edges)),
         diameter(build_adjacency(ctrl_edges)),
         diameter(build_adjacency(chain(data_edges, ctrl_edges))),
     )
 
+
 class FunctionCharacteristics(IntFlag):
-    """Helps evaluating the level of difficulty of this function."""
+    """Helps evaluating the level of difficulty of a function,
+    regarding miscalleneous ML tasks."""
 
     data = auto()
     condition = auto()
@@ -68,6 +84,7 @@ class FunctionCharacteristics(IntFlag):
 
 
 def get_function_characteristics(graph: CDG) -> FunctionCharacteristics:
+    """Characterizes a function regarding control flow features."""
 
     # Every function contains data-flow easiest level
     function_characteristics = FunctionCharacteristics.data

{pcode_graph-1.0.0 → pcode_graph-1.1.0}/pcode_graph/checker.py

@@ -34,8 +34,8 @@ def check_graph(graph: CDG):
         )
 
         proxy = proxy_dict[index]
-        control_succs = proxy.get_control_successors()
-        data_preds = proxy.get_data_predecessors()
+        control_succs = list(proxy.get_control_successors())
+        data_preds = list(proxy.get_data_predecessors())
 
         if node.kind != NodeKinds.Operation:
             assert node.opcode is None
@@ -46,11 +46,11 @@ def check_graph(graph: CDG):
 
             data_numbers: list[int] = []
             control_numbers: list[int] = []
-            for edge in proxy.in_edges:
+            for edge in proxy.get_in_edges():
                 if edge.kind == EdgeKinds.Data:
                     assert edge.operand_number is not None
                     data_numbers.append(edge.operand_number)
-            for edge in proxy.out_edges:
+            for edge in proxy.get_out_edges():
                 if edge.kind == EdgeKinds.Control:
                     assert edge.operand_number is not None
                     control_numbers.append(edge.operand_number)

{pcode_graph-1.0.0 → pcode_graph-1.1.0}/pcode_graph/cli.py

@@ -3,7 +3,7 @@
 Dumps Control and Data Graph (CDG) and various data from a binary or assembly file.
 
 Usage:
-    cdg ( asm | functions | html | md | pcode | table ) <input_path> [options]
+    cdg ( asm | functions | html | md | pcode | table | stats ) <input_path> [options]
 
 Dump commands:
     asm        Assembly dump
@@ -12,6 +12,7 @@ Dump commands:
     md         CDG in markdown/mermaid format
     pcode      P-Code operations
     table      Markdown table of the P-Code enriched by static analysis results
+    stats      Prints graph statistics
 
 Options:
     -f, --function <name>       Take only the function with given symbol name.
@@ -21,10 +22,11 @@ Options:
     -D, --dataflow-only         Only extract dataflow graph (and no control flow).
     -o, --output <path>         Where to write the output.
     -h, --help                  This help.
-    -d, --debug                 Debug mode: do no catch exceptions.    
+    -d, --debug                 Debug mode: do no catch exceptions.
 """
 
 from pathlib import Path
+import re
 import sys
 from docopt import docopt
 from lief import Binary
@@ -33,6 +35,7 @@ from loguru import logger
 import pypcode
 from pcode_graph.arch import Arch
 from pcode_graph.asm import Assembler
+from pcode_graph.charac import count_edges, get_graph_diameters
 from pcode_graph.lief_importer import (
     get_function,
     iter_code_sections,
@@ -104,6 +107,10 @@ def main():
                     code = function.content
                 else:
                     # Lift piece of code
+                    if not re.search("[a-fA-FxX]", chunk_begin_end):
+                        logger.warning(
+                            f"This does not look like hexadecimal addresses: {chunk_begin_end}"
+                        )
                     start_address, end_address = [
                         int(a, 16) for a in chunk_begin_end.split(":")
                     ]
@@ -164,6 +171,17 @@ def main():
             print(str(cdg), file=output)
             sys.exit(0)
 
+        if args["stats"]:
+            stats = dict(
+                diameter=get_graph_diameters(cdg)._asdict(),
+                edges=count_edges(cdg)._asdict(),
+                nodes=len(cdg.nodes),
+            )
+
+            print(repr(stats), file=output)
+            sys.exit(0)
+
+        assert args["html"], "Unknown command"
         assert output_path, "HTML output file is mandatory"
         output.close()
         draw_graph(cdg, Path(output_path))

{pcode_graph-1.0.0 → pcode_graph-1.1.0}/pcode_graph/emulator.py

@@ -69,7 +69,7 @@ class Emulator(PcodeVisitor):
                             OperandNumber | None, tuple[NodeIndex, Value]
                         ] = {}
 
-                        for in_edge in p.in_edges:
+                        for in_edge in p.get_in_edges():
                             if in_edge.kind != EdgeKinds.Data:
                                 continue
 
@@ -88,7 +88,8 @@ class Emulator(PcodeVisitor):
                             propagated = True
 
                     case NodeKinds.OutputRegister:
-                        (in_edge,) = p.get_inputs_edges()
+                        (in_edge_index,) = p.get_input_edges()
+                        in_edge = self.graph.edges[in_edge_index]
                         value = self.values.get(in_edge.source_node)
                         if value is not None:
                             self.values[index] = value

{pcode_graph-1.0.0 → pcode_graph-1.1.0}/pcode_graph/graph.py

@@ -1,9 +1,11 @@
 from dataclasses import dataclass, replace
 from enum import Enum, auto
+from typing import Iterator
 from pcode_graph.pcode import OpCodes
 from pcode_graph.utils import unsigned_to_signed
 
 NodeIndex = int
+EdgeIndex = int
 OperandNumber = int
 RegisterId = int
 Value = int
@@ -31,7 +33,7 @@ class EdgeKinds(Enum):
     Control = 1
 
 
-@dataclass
+@dataclass(slots=True)
 class Edge:
     source_node: NodeIndex
 
@@ -44,7 +46,7 @@ class Edge:
     operand_number: OperandNumber | None
 
 
-@dataclass
+@dataclass(slots=True)
 class Node:
     # Kind of node
     kind: NodeKinds
@@ -99,43 +101,54 @@ class Node:
                 raise ValueError(self.kind)
 
 
-@dataclass
+@dataclass(slots=True)
 class EdgeProxy:
     """Easy access to edges for one node."""
 
-    in_edges: list[Edge]
-    out_edges: list[Edge]
-
-    def get_data_predecessors(self) -> list[NodeIndex]:
-        return [
-            edge.source_node for edge in self.in_edges if edge.kind == EdgeKinds.Data
-        ]
-
-    def get_control_predecessors(self) -> list[NodeIndex]:
-        return [
-            edge.source_node for edge in self.in_edges if edge.kind == EdgeKinds.Control
-        ]
-
-    def get_data_successors(self) -> list[NodeIndex]:
-        return [
-            edge.destination_node
-            for edge in self.out_edges
-            if edge.kind == EdgeKinds.Data
-        ]
-
-    def get_control_successors(self) -> list[NodeIndex]:
-        return [
-            edge.destination_node
-            for edge in self.out_edges
-            if edge.kind == EdgeKinds.Control
-        ]
-
-    def get_inputs_edges(self) -> list[Edge]:
+    graph: CDG
+    in_edge_indexes: list[EdgeIndex]
+    out_edge_indexes: list[EdgeIndex]
+
+    def get_in_edges(self) -> Iterator[Edge]:
+        for index in self.in_edge_indexes:
+            yield self.graph.edges[index]
+
+    def get_out_edges(self) -> Iterator[Edge]:
+        for index in self.out_edge_indexes:
+            yield self.graph.edges[index]
+
+    def get_data_predecessors(self) -> Iterator[NodeIndex]:
+        for e in self.in_edge_indexes:
+            edge = self.graph.edges[e]
+            if edge.kind == EdgeKinds.Data:
+                yield edge.source_node
+
+    def get_control_predecessors(self) -> Iterator[NodeIndex]:
+        for e in self.in_edge_indexes:
+            edge = self.graph.edges[e]
+            if edge.kind == EdgeKinds.Control:
+                yield edge.source_node
+
+    def get_data_successors(self) -> Iterator[NodeIndex]:
+        for e in self.out_edge_indexes:
+            edge = self.graph.edges[e]
+            if edge.kind == EdgeKinds.Data:
+                yield edge.destination_node
+
+    def get_control_successors(self) -> Iterator[NodeIndex]:
+        for e in self.out_edge_indexes:
+            edge = self.graph.edges[e]
+            if edge.kind == EdgeKinds.Control:
+                yield edge.destination_node
+
+    def get_input_edges(self) -> list[EdgeIndex]:
         """Returns data input operands, by number."""
-        operands: list[tuple[OperandNumber | None, Edge]] = []
+
+        operands: list[tuple[OperandNumber | None, EdgeIndex]] = []
         hasNone = False
         hasNumber = False
-        for in_edge in self.in_edges:
+        for in_edge_index in self.in_edge_indexes:
+            in_edge = self.graph.edges[in_edge_index]
             if in_edge.kind == EdgeKinds.Control:
                 continue
             n = in_edge.operand_number
@@ -143,14 +156,14 @@ class EdgeProxy:
                 hasNone = True
             else:
                 hasNumber = True
-            operands.append((n, in_edge))
+            operands.append((n, in_edge_index))
         assert hasNone ^ hasNumber
         if hasNumber:
             operands.sort()
         return [o[1] for o in operands]
 
 
-@dataclass
+@dataclass(slots=True)
 class CDG:
     """
     Graph of P-Code operations with Control and Data flow edges.
@@ -167,10 +180,10 @@ class CDG:
         return CDG([replace(n) for n in self.nodes], [replace(e) for e in self.edges])
 
     def compute_edge_proxy(self) -> dict[NodeIndex, EdgeProxy]:
-        cache = {index: EdgeProxy([], []) for index in range(len(self.nodes))}
-        for edge in self.edges:
-            cache[edge.source_node].out_edges.append(edge)
-            cache[edge.destination_node].in_edges.append(edge)
+        cache = {index: EdgeProxy(self, [], []) for index in range(len(self.nodes))}
+        for e, edge in enumerate(self.edges):
+            cache[edge.source_node].out_edge_indexes.append(e)
+            cache[edge.destination_node].in_edge_indexes.append(e)
         return cache
 
     def add_node(self, node: Node) -> NodeIndex:

{pcode_graph-1.0.0 → pcode_graph-1.1.0}/pcode_graph/lief_importer.py

@@ -29,7 +29,7 @@ class FunctionOps(ChunkOps):
 
 
 def parse_binary(binary_path: Path) -> lief.Binary:
-    assert binary_path.exists(), f"{binary_path} not found"
+    assert binary_path.is_file(), f"{binary_path} is not a file"    
     logger.debug(f"Parse {binary_path}")
     b = lief.parse(binary_path.as_posix())
     assert isinstance(b, lief.Binary)

{pcode_graph-1.0.0 → pcode_graph-1.1.0}/pcode_graph/maker.py

@@ -26,10 +26,18 @@ class MakerFlags:
     build_cfg: bool = True
     """Add control flow links in addition to the data flow. Defaults to True."""
 
-    ignore_reg_outputs: set[str] | None = None
+    reg_outputs_black_list: set[str] | None = None
     """Do not consider the given registers as graph outputs to simplify the graphs.
     Arch-dependent register names.
-    If None, the result of pcode.get_special_register_names() is used.
+    If white-list is not None, the value of black-list is ignored.
+    If both black-list and white-list are None, the content of SPECIAL_REGISTERS is used
+    for black-list.
+    """
+
+    reg_outputs_white_list: set[str] | None = None
+    """Consider only the given registers as graph outputs to simplify the graphs.
+    Arch-dependent register names.
+    If white-list is not None, the value of black-list is ignored.
     """
 
 
@@ -45,10 +53,11 @@ def make_graph(prg: RichPcodeList, flags: MakerFlags | None = None) -> CDG:
     """
 
     if flags is None:
-        flags = MakerFlags()
-    ignore_reg_outputs = flags.ignore_reg_outputs
-    if ignore_reg_outputs is None:
-        ignore_reg_outputs = SPECIAL_REGISTERS[prg.arch]
+        flags = MakerFlags()    
+
+    black_list, white_list = flags.reg_outputs_black_list, flags.reg_outputs_white_list
+    if black_list is None and white_list is None:
+        black_list = SPECIAL_REGISTERS[prg.arch]    
 
     # Already created input nodes, by related input variable
     input_node_indexes: dict[Var, NodeIndex] = {}
@@ -189,13 +198,22 @@ def make_graph(prg: RichPcodeList, flags: MakerFlags | None = None) -> CDG:
     for var, def_op_indexes in prg.get_exit_defs().items():
         assert def_op_indexes
         match var.kind:
-            case VarKinds.Register:
-                if var.pretty_name.lower().startswith("tmp"):
-                    # Ignore writes to temporaries
-                    continue
-                if var.pretty_name in ignore_reg_outputs:
-                    # Ignore writes to special registers
-                    continue
+            case VarKinds.Register:                
+                if white_list is not None:
+                    if var.pretty_name not in white_list:
+                        continue
+                else:
+                    assert black_list is not None
+                    if var.pretty_name.lower().startswith("tmp"):
+                        # Ignore writes to temporaries
+                        continue
+                    if var.pretty_name in black_list:
+                        # Ignore writes to special registers
+                        continue
+                    if var.pretty_name == "UNKNOWN":
+                        # Register name not provided in P-Code
+                        continue
+                assert var.pretty_name
                 output_node = Node(
                     NodeKinds.OutputRegister,
                     var.size,
@@ -309,8 +327,7 @@ def make_graph(prg: RichPcodeList, flags: MakerFlags | None = None) -> CDG:
 
     graph = CDG(nodes, edges)
 
-    # Uncomment to locate a graph structure bug
-    # Comment to speed up
+    # Uncomment to locate a bug: graph structure vs simplifier
     # check_graph(graph)
 
     GraphSimplifier(graph).process()

{pcode_graph-1.0.0 → pcode_graph-1.1.0}/pcode_graph/simplifier.py

@@ -1,5 +1,13 @@
+
 from pcode_graph.checker import check_graph
-from pcode_graph.graph import CDG, Edge, EdgeKinds, NodeIndex, OpCodes, NodeKinds
+from pcode_graph.graph import (
+    CDG,
+    EdgeIndex,
+    EdgeKinds,
+    NodeIndex,
+    OpCodes,
+    NodeKinds,
+)
 
 from loguru import logger
 from pathlib import Path
@@ -13,81 +21,86 @@ class GraphSimplifier:
     ) -> None:
         self.graph = graph
 
-    def _update_proxy(self):
         self._proxy = self.graph.compute_edge_proxy()
 
+    def _finalize_removals(self):
+
+        remap: dict[int, int] = {}
+        shift = 0
+        kept_nodes = []
+        for index, node in enumerate(self.graph.nodes):
+            if node is None:
+                shift += 1
+            else:
+                remap[index] = index - shift
+                kept_nodes.append(node)
+
+        self.graph.nodes = kept_nodes
+
+        kept_edges = []
+        for edge in self.graph.edges:
+            if edge is None:
+                continue
+            edge.destination_node = remap[edge.destination_node]
+            edge.source_node = remap[edge.source_node]
+            kept_edges.append(edge)
+
+        self.graph.edges = kept_edges
+
     def _remove_node(self, index: NodeIndex):
-        def remap_index(n: NodeIndex):
-            assert n != index
-            return n - int(n > index)
 
         # logger.debug(f"Remove {self.graph.nodes[index]} at {index}")
 
-        # Remap edges
         proxy = self._proxy[index]
 
-        control_succs = proxy.get_control_successors()
+        control_succs = list(proxy.get_control_successors())
         assert (
             len(control_succs) <= 1
         ), "More than one control successor to the node to remove"
 
-        data_preds = proxy.get_data_predecessors()
+        data_preds = list(proxy.get_data_predecessors())
 
         # Convert to set to handle particular case of Phi node
         assert len(data_preds) <= 1 or (
             self.graph.nodes[index].kind == NodeKinds.Phi and len(set(data_preds)) == 1
         ), "More than one data predecessor to the node to remove"
 
-        kept_edges: list[Edge] = []
-
-        def remap_edge(
-            kind: EdgeKinds, source: NodeIndex, dest: NodeIndex, operand: int | None
-        ):
-            # logger.debug(
-            #     f"Remap {self.graph.dump_edge(Edge(kind=kind, source_node=source, destination_node=dest, operand_number=operand))}"
-            # )
-            edge = Edge(
-                source_node=remap_index(source),
-                destination_node=remap_index(dest),
-                kind=kind,
-                operand_number=operand,
-            )
-            kept_edges.append(edge)
-
-        for e in self.graph.edges:
+        for i, e in enumerate(self.graph.edges):
+            if e is None:
+                continue
+            remove = False
             if e.source_node == index:
                 if e.kind == EdgeKinds.Data:
-                    remap_edge(
-                        e.kind,
-                        data_preds[0],
-                        e.destination_node,
-                        e.operand_number,
-                    )
-            elif e.destination_node == index:
+                    self._proxy[e.source_node].out_edge_indexes.remove(i)
+                    e.source_node = data_preds[0]
+                    self._proxy[e.source_node].out_edge_indexes.append(i)
+                else:
+                    remove = True
+
+            if e.destination_node == index:
                 if e.kind == EdgeKinds.Control:
-                    remap_edge(
-                        e.kind,
-                        e.source_node,
-                        control_succs[0],
-                        e.operand_number,
-                    )
-            else:
-                remap_edge(e.kind, e.source_node, e.destination_node, e.operand_number)
+                    self._proxy[e.destination_node].in_edge_indexes.remove(i)
+                    e.destination_node = control_succs[0]
+                    self._proxy[e.destination_node].in_edge_indexes.append(i)
+                else:
+                    remove = True
 
-        # Patch the graph
-        self.graph.edges = kept_edges
-        self.graph.nodes.pop(index)
+            if remove:
+                self._remove_edge(i)
 
+        # Patch the graph
+        self.graph.nodes[index] = None  # type: ignore
         self._graph_changed = True
-        self._update_proxy()
 
         # self.debug_helper()
 
-    def _remove_edge(self, edge: Edge):
+    def _remove_edge(self, edge_index: int):
 
-        self.graph.remove_edge(edge)
+        edge = self.graph.edges[edge_index]
+        self._proxy[edge.source_node].out_edge_indexes.remove(edge_index)
+        self._proxy[edge.destination_node].in_edge_indexes.remove(edge_index)
+        self.graph.edges[edge_index] = None  # type: ignore
         self._graph_changed = True
-        self._update_proxy()
 
     def debug_helper(self):
         try:
@@ -102,7 +115,6 @@ class GraphSimplifier:
     def process(self):
 
         self._graph_changed = True
-        self._update_proxy()
 
         while self._graph_changed:
             self._graph_changed = False
@@ -110,6 +122,9 @@ class GraphSimplifier:
             # Remove useless phi nodes
             for index, node in enumerate(self.graph.nodes):
 
+                if node is None:
+                    continue
+
                 if node.kind == NodeKinds.Phi:
 
                     phi_data_preds = set(self._proxy[index].get_data_predecessors())
@@ -117,13 +132,12 @@ class GraphSimplifier:
                         # Keep this Phi node
                         continue
 
-                    self._remove_node(index)
-                    break
+                    self._remove_node(index)                    
 
                 elif node.opcode == OpCodes.COPY:
 
                     ctrl_preds = set(self._proxy[index].get_control_predecessors())
-                    ctrl_succs = self._proxy[index].get_control_successors()
+                    ctrl_succs = set(self._proxy[index].get_control_successors())
 
                     remove = True
                     if ctrl_preds and ctrl_succs:
@@ -131,55 +145,52 @@ class GraphSimplifier:
 
                         remove = False
                         # Remove copies that dominates their successor...
-                        if self._proxy[ctrl_succ].get_control_predecessors() == [index]:
+                        if list(self._proxy[ctrl_succ].get_control_predecessors()) == [
+                            index
+                        ]:
                             remove = True
                         elif len(ctrl_preds) == 1:
                             # ...or are dominated by their unique predecessor
                             (ctrl_pred,) = ctrl_preds
-                            if self._proxy[ctrl_pred].get_control_successors() == [
-                                index
-                            ]:
+                            if list(
+                                self._proxy[ctrl_pred].get_control_successors()
+                            ) == [index]:
                                 remove = True
 
                     if remove:
                         self._remove_node(index)
-                        break
 
-                elif node.kind == NodeKinds.Constant:                    
-                    if not self._proxy[index].get_data_successors():
+                elif node.kind == NodeKinds.Constant:
+                    if not list(self._proxy[index].get_data_successors()):
                         # Remove constants let orphan by other simplifications
                         self._remove_node(index)
-                        break
 
                 elif node.opcode in {
                     OpCodes.INT_ZEXT,
                     OpCodes.INT_SEXT,
                 }:
                     self._remove_node(index)
-                    break
 
                 elif node.opcode == OpCodes.SUBPIECE:
 
-                    operands: list[Edge] = self._proxy[index].get_inputs_edges()
-
+                    operands: list[EdgeIndex] = self._proxy[index].get_input_edges()
                     assert len(operands) == 2
-                    num_trunc_bytes = self.graph.nodes[operands[1].source_node]
+                    source_node_index = self.graph.edges[operands[1]].source_node
+                    num_trunc_bytes = self.graph.nodes[source_node_index]
                     assert num_trunc_bytes.kind == NodeKinds.Constant
                     if num_trunc_bytes.value == 0:
                         self._remove_edge(operands[1])
-                        self._remove_node(index)                        
-                        break
+                        self._remove_node(index)
 
                 elif node.opcode == OpCodes.INT_MULT:
                     # Remove mulitiplication by one
-                    operands: list[Edge] = self._proxy[index].get_inputs_edges()
-                    for op_edge in operands:
+                    for op_edge_index in self._proxy[index].get_input_edges():
+                        op_edge = self.graph.edges[op_edge_index]
                         o = op_edge.source_node
                         op = self.graph.nodes[o]
                         if op.kind == NodeKinds.Constant and op.value == 1:
-                            self._remove_edge(op_edge)
-                            self._remove_node(index)                            
+                            self._remove_edge(op_edge_index)
+                            self._remove_node(index)
                             break
-                    else:
-                        continue
-                    break
+
+        self._finalize_removals()

{pcode_graph-1.0.0 → pcode_graph-1.1.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pcode_graph"
-version = "1.0.0"
+version = "1.1.0"
 authors = [
   { name="Samuel Hangouët", email="samuel.hangouet@quarkslab.com" },
 ]

{pcode_graph-1.0.0 → pcode_graph-1.1.0}/tests/maker_test.py

@@ -98,7 +98,7 @@ ret
 """,
         x86_64,
         outdir,
-        flags=MakerFlags(ignore_reg_outputs=special_regs),
+        flags=MakerFlags(reg_outputs_black_list=special_regs),
     )
 
     check_inputs(g, "#0x8", "RSP", "MEMORY")
@@ -131,7 +131,7 @@ ldp x29, x30, [sp], #0x10
 """,
         arm64,
         outdir,
-        flags=MakerFlags(ignore_reg_outputs=special_regs),
+        flags=MakerFlags(reg_outputs_black_list=special_regs),
     )
 
     load1, load2 = [n for n, node in enumerate(g.nodes) if node.opcode == OpCodes.LOAD]
@@ -166,10 +166,10 @@ ldr x2, [sp, 0x10]
 
     (load_index,) = [n for n, node in enumerate(g.nodes) if node.opcode == OpCodes.LOAD]
     proxy = g.compute_edge_proxy()
-    phi_index, add_index = proxy[load_index].get_data_predecessors()
+    phi_index, add_index = list(proxy[load_index].get_data_predecessors())
     assert g.nodes[phi_index].kind == NodeKinds.Phi
     assert g.nodes[add_index].opcode == OpCodes.INT_ADD
-    assert len(proxy[phi_index].get_data_predecessors()) == 3
+    assert len(list(proxy[phi_index].get_data_predecessors())) == 3
 
 
 def test_call_ind_arm(outdir: Path, arm64: ArchContext, log):
@@ -199,7 +199,7 @@ mul x1, x0, x0
     else:
         assert False, "Call not found"
 
-    assert len(g.compute_edge_proxy()[call].get_control_successors()) == 2
+    assert len(list(g.compute_edge_proxy()[call].get_control_successors())) == 2
 
 
 def test_div(outdir: Path, x86_64: ArchContext, log):
@@ -240,7 +240,7 @@ mov x0, x6
 
     proxy = g.compute_edge_proxy()
 
-    assert proxy[x0_output].get_data_predecessors() == [x1_input]
+    assert list(proxy[x0_output].get_data_predecessors()) == [x1_input]
 
 
 def test_various_opcodes(outdir: Path, arm64: ArchContext, log):
@@ -409,14 +409,14 @@ ret
 
     proxy = g.compute_edge_proxy()
     (load,) = [i for i, n in enumerate(g.nodes) if n.opcode == OpCodes.LOAD]
-    assert len(proxy[load].get_control_predecessors()) == 2
+    assert len(list(proxy[load].get_control_predecessors())) == 2
 
     (rax,) = [
         i
         for i, n in enumerate(g.nodes)
         if n.kind == NodeKinds.OutputRegister and n.register_name == "RAX"
     ]
-    (rax_pred,) = proxy[rax].get_data_predecessors()
+    (rax_pred,) = list(proxy[rax].get_data_predecessors())
     assert g.nodes[rax_pred].kind == NodeKinds.Phi
 
 
@@ -450,7 +450,7 @@ jmp rax
 
     (end,) = [i for i, n in enumerate(g.nodes) if n.kind == NodeKinds.External]
     (branchind,) = [i for i, n in enumerate(g.nodes) if n.opcode == OpCodes.BRANCHIND]
-    assert g.compute_edge_proxy()[end].get_control_predecessors() == [branchind]
+    assert list(g.compute_edge_proxy()[end].get_control_predecessors()) == [branchind]
 
 
 def test_imul(outdir: Path, x86_64: ArchContext):
@@ -569,7 +569,7 @@ call 0x546fb
         asm,
         x86_64,
         outdir / "full",
-        flags=MakerFlags(ignore_reg_outputs=ignored_regs),
+        flags=MakerFlags(reg_outputs_black_list=ignored_regs),
     )
     check_inputs(g2, "RSP", "#0x8", "#0x5")  # 0x5 is return address
     check_outputs(g2, "RSP", "MEMORY")
@@ -1131,7 +1131,7 @@ RET
 
     for n, node in enumerate(g.nodes):
         if node.kind == NodeKinds.Constant:
-            assert proxy[n].get_data_successors(), f"Orphan constant {node}"
+            assert list(proxy[n].get_data_successors()), f"Orphan constant {node}"
 
 
 def test_lea_simplification(outdir: Path, x86_64: ArchContext, log):

pcode_graph-1.1.0/todo.md

@@ -0,0 +1,3 @@
+- Add Basic Blocks in Analysis stage to speed up Reaching Defs computation
+- Specify CDGv2 with BB nodes and check it helps to reduce graph diameter while keeping CFG invariant by instruction reordering.
+