AST: add viewstmt to pirinspector

Author: sipma

chb/app/CHVersion.py

@@ -1 +1 @@
-chbversion: str = "0.3.0-20250722"
+chbversion: str = "0.3.0-20250723"

chb/arm/ARMInstruction.py

@@ -441,8 +441,9 @@ def to_string(
             opcodewidth: int = 40,
             typingrules: bool = False,
             sp: bool = False) -> str:
+
+        lines: List[str] = []
         if typingrules:
-            lines: List[str] = []
             rulesapplied = self.app.type_constraints.rules_applied_to_instruction(
                 self.armfunction.faddr, self.iaddr)
             for r in sorted(str(r) for r in rulesapplied):

chb/ast/ASTDeserializer.py

@@ -71,6 +71,15 @@ def low_level_ast(self) -> AST.ASTStmt:
     def high_unreduced_ast(self) -> AST.ASTStmt:
         return self.asts[1]
 
+    def get_stmt(self, stmtid: int) -> AST.ASTStmt:
+        for node in self.nodes.values():
+            if node.is_ast_stmt:
+                node = cast(AST.ASTStmt, node)
+                if node.stmtid == stmtid:
+                    return node
+        else:
+            raise Exception("No stmt found with id: " + str(stmtid))
+
     def get_instruction(self, instrid: int) -> AST.ASTInstruction:
         for node in self.nodes.values():
             if node.is_ast_instruction:

chb/ast/ASTNOPVisitor.py

@@ -148,6 +148,9 @@ def visit_question_expression(self, expr: AST.ASTQuestion) -> None:
     def visit_address_of_expression(self, expr: AST.ASTAddressOf) -> None:
         pass
 
+    def visit_start_of_expression(self, expr: AST.ASTStartOf) -> None:
+        pass
+
     def visit_void_typ(self, typ: AST.ASTTypVoid) -> None:
         pass
 

chb/ast/ASTProvenanceCollector.py

@@ -0,0 +1,76 @@
+# ------------------------------------------------------------------------------
+# CodeHawk Binary Analyzer
+# Author: Henny Sipma
+# ------------------------------------------------------------------------------
+# The MIT License (MIT)
+#
+# Copyright (c) 2021-2025  Aarno Labs LLC
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+# ------------------------------------------------------------------------------
+"""Collects stmt and instruction provenance for pir inspector."""
+
+from typing import Dict, List, TYPE_CHECKING
+
+import chb.ast.ASTNode as AST
+from chb.ast.ASTNOPVisitor import ASTNOPVisitor
+
+if TYPE_CHECKING:
+    from chb.ast.ASTDeserializer import ASTFunctionDeserialization
+
+
+class ASTProvenanceCollector(ASTNOPVisitor):
+
+    def __init__(self, dfn: "ASTFunctionDeserialization") -> None:
+        self._dfn = dfn
+        self._hl_ll_instr_mapping: Dict[int, List[int]] = {}
+
+    @property
+    def dfn(self) -> "ASTFunctionDeserialization":
+        return self._dfn
+
+    @property
+    def provenance(self) -> Dict[int, List[AST.ASTInstruction]]:
+        result: Dict[int, List[AST.ASTInstruction]] = {}
+        for hlinstrid in self._hl_ll_instr_mapping:
+            llinstrids = self._hl_ll_instr_mapping[hlinstrid]
+            result[hlinstrid] = [self.dfn.get_instruction(i) for i in llinstrids]
+        return result
+
+    def instruction_provenance(
+            self, stmt: AST.ASTStmt) -> Dict[int, List[AST.ASTInstruction]]:
+        stmt.accept(self)
+        return self.provenance
+
+    def visit_loop_stmt(self, stmt: AST.ASTLoop) -> None:
+        stmt.body.accept(self)
+
+    def visit_block_stmt(self, stmt: AST.ASTBlock) -> None:
+        for s in stmt.stmts:
+            s.accept(self)
+
+    def visit_instruction_sequence_stmt(self, stmt: AST.ASTInstrSequence) -> None:
+        for i in stmt.instructions:
+            if i.instrid in self.dfn.astree.provenance.instruction_mapping:
+                self._hl_ll_instr_mapping[i.instrid] = (
+                    self.dfn.astree.provenance.instruction_mapping[i.instrid])
+
+    def visit_branch_stmt(self, stmt: AST.ASTBranch) -> None:
+        stmt.ifstmt.accept(self)
+        stmt.elsestmt.accept(self)

chb/ast/ASTViewer.py

@@ -95,6 +95,16 @@ def to_graph(self, stmt: AST.ASTStmt) -> DU.ASTDotGraph:
         stmt.accept(self)
         return self.dotgraph
 
+    def stmt_to_graph(
+            self,
+            stmt: AST.ASTStmt,
+            provinstrs: Dict[int, List[AST.ASTInstruction]] = {}) -> DU.ASTDotGraph:
+        stmt.accept(self)
+        for pl in provinstrs.values():
+            for p in pl:
+                p.accept(self)
+        return self.dotgraph
+
     def instr_to_graph(
             self,
             instr: AST.ASTInstruction,
@@ -124,6 +134,9 @@ def get_expr_connections(self, expr: AST.ASTExpr) -> str:
         if self.astree.provenance.has_reaching_definitions(expr.exprid):
             ids = self.astree.provenance.reaching_definitions[expr.exprid]
             result += "\\nrdefs:[" + ",".join(str(id) for id in ids) + "]"
+        if expr.exprid in self.astree.expr_spanmap():
+            span = self.astree.expr_spanmap()[expr.exprid]
+            result += "\\ncc:" + span
         return result
 
     def stmt_name(self, stmt: AST.ASTStmt) -> str:

chb/ast/AbstractSyntaxTree.py

@@ -287,6 +287,25 @@ def get_exprid(self, exprid: Optional[int]) -> int:
     def add_span(self, span: ASTSpanRecord) -> None:
         self._spans.append(span)
 
+    def add_stmt_span(
+            self, locationid: int, spans: List[Tuple[str, str]]) -> None:
+        """Add a span for the ast instructions contained in a stmt.
+
+        Note: this is currently done only for if statements originating from
+        predicated instructions.
+        """
+        spaninstances: List[Dict[str, Union[str, int]]] = []
+        for (iaddr, bytestring) in spans:
+            span: Dict[str, Union[str, int]] = {}
+            span["base_va"] = iaddr
+            span["size"] = len(bytestring) // 2
+            spaninstances.append(span)
+        spanrec: Dict[str, Any] = {}
+        spanrec["locationid"] = locationid
+        spanrec["spans"] = spaninstances
+        self.add_span(cast(ASTSpanRecord, spanrec))
+
+
     def add_instruction_span(
             self, locationid: int, base: str, bytestring: str) -> None:
         """Add a span for an ast instruction."""
@@ -1224,7 +1243,7 @@ def _cast_if_needed(self, e: AST.ASTExpr, opunsigned: bool) -> AST.ASTExpr:
         # CIL encodes signedness in the shift operator but C infers
         # the operator flavor from the type of the left operand, so
         # we may need to insert a cast to ensure that serialization
-        # through C code preserves the AST. 
+        # through C code preserves the AST.
 
         mb_t = e.ctype(ASTBasicCTyper(self.globalsymboltable))
         force_cast = mb_t is None

chb/ast/astutil.py

@@ -39,6 +39,7 @@
 from chb.ast.ASTCPrettyPrinter import ASTCPrettyPrinter
 from chb.ast.ASTDeserializer import ASTDeserializer
 import chb.ast.ASTNode as AST
+from chb.ast.ASTProvenanceCollector import ASTProvenanceCollector
 from chb.ast.ASTViewer import ASTViewer
 import chb.ast.astdotutil as DU
 
@@ -191,6 +192,36 @@ def viewastcmd(args: argparse.Namespace) -> NoReturn:
     exit(0)
 
 
+def viewstmtcmd(args: argparse.Namespace) -> NoReturn:
+
+    # arguments
+    pirfile: str = args.pirfile
+    function: Optional[str] = args.function
+    stmtid: int = args.stmtid
+    provenance: bool = args.provenance
+    outputfilename: str = args.output
+
+    with open(pirfile, "r") as fp:
+        pirjson = json.load(fp)
+
+    faddr = get_function_addr(pirjson, function)
+    deserializer = ASTDeserializer(pirjson)
+    (globaltable, dfns) = deserializer.deserialize()
+    for dfn in dfns:
+        if dfn.astree.faddr == faddr:
+            stmt = dfn.get_stmt(stmtid)
+            viewer = ASTViewer(faddr, dfn.astree)
+            if provenance:
+                provcollector = ASTProvenanceCollector(dfn)
+                provinstrs = provcollector.instruction_provenance(stmt)
+                g = viewer.stmt_to_graph(stmt, provinstrs)
+            else:
+                g = viewer.to_graph(stmt)
+
+        DU.print_dot(outputfilename, g)
+    exit(0)
+
+
 def viewinstrcmd(args: argparse.Namespace) -> NoReturn:
 
     # arguments

chb/ast/pirinspector

@@ -93,6 +93,24 @@ def parse() -> argparse.Namespace:
     )
     viewastcmd.set_defaults(func=AU.viewastcmd)
 
+    # --- view selected stmt ast
+    viewstmtcmd = viewparsers.add_parser("stmt")
+    viewstmtcmd.add_argument("pirfile", help="name of json file with ast information")
+    viewstmtcmd.add_argument("--function", help="name or address of function to view")
+    viewstmtcmd.add_argument(
+        "--stmtid", help="stmt id of statement", type=int)
+    viewstmtcmd.add_argument(
+        "--provenance",
+        help="show associated low-level instructions",
+        action="store_true",
+    )
+    viewstmtcmd.add_argument(
+        "-o", "--output",
+        help="name of output graph file (without extension)",
+        required=True,
+    )
+    viewstmtcmd.set_defaults(func=AU.viewstmtcmd)
+
     # --- view single instruction ast
     viewinstrcmd = viewparsers.add_parser("instruction")
     viewinstrcmd.add_argument("pirfile", help="name of json file with ast information")

chb/astinterface/ASTInterface.py

@@ -585,7 +585,10 @@ def add_formal(
         return nextindex
 
     def add_span(self, span: ASTSpanRecord) -> None:
-        self.astree.add_span
+        self.astree.add_span(span)
+
+    def add_stmt_span(self, id: int, spans: List[Tuple[str, str]]) -> None:
+        self.astree.add_stmt_span(id, spans)
 
     def add_instruction_span(self, id: int, base: str, bytestring: str) -> None:
         self.astree.add_instruction_span(id, base, bytestring)