ASTI: introduce basic block fragment for predicated instructions

sipma · sipma · commit b2234eeada78 · 2025-07-08T10:18:26.000-07:00
diff --git a/chb/app/BasicBlock.py b/chb/app/BasicBlock.py
@@ -4,7 +4,7 @@
 # ------------------------------------------------------------------------------
 # The MIT License (MIT)
 #
-# Copyright (c) 2021-2024  Aarno Labs LLC
+# 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
@@ -52,12 +52,135 @@
     import chb.app.Function
 
 
+class BasicBlockFragment:
+    """Represents a basic block fragment without ast control flow.
+
+    In ARM instructions may be predicated, e.g.:
+
+    MOVEQ R0, R1  : if condition then R0 := R1
+
+    In the ocaml analyzer this additional control flow can be accomodated
+    directly in CHIF without the need to create a separate basic block
+    in the CFG (i.e. lightweight control flow). In decompilation, some of
+    these instructions may be accomodated without explicit control flow (e.g.,
+    by the C ternary operation), but this is not always possible.
+
+    However, even when predicated instructions cannot be converted into
+    expressions, it is not necessary to create top-level basic blocks
+    that are subject to the CFG-to-AST conversion. A more light-weight
+    solution is to embed the necessary control flow (notably limited
+    to branches and instruction sequences) within the block created for
+    the original basic block.
+
+    The basic block is partitioned into a linear sequence of BasicBlock
+    Fragments, where each fragment can be one of the following:
+    - a (linear) instruction sequence statement
+    - a branch statement containing a condition and a single (if) branch
+    - a branch statement containing a condition and a then and an else
+       branch.
+    In case of the branch statement either branch can have one or more
+    instructions that have the same condition setter and the same
+    condition.
+    """
+
+    def __init__(self, instr: Instruction) -> None:
+        self._linear: List[Instruction] = []
+        self._thenbranch: List[Instruction] = []
+        self._elsebranch: List[Instruction] = []
+        self._setter: Optional[str] = None
+        self._condition: Optional[str] = None
+        self.add_instr(instr)
+
+    @property
+    def condition(self) -> Optional[str]:
+        return self._condition
+
+    @property
+    def setter(self) -> Optional[str]:
+        return self._setter
+
+    @property
+    def is_predicated(self) -> bool:
+        return self.condition is not None
+
+    @property
+    def is_then_only(self) -> bool:
+        return self.is_predicated and len(self.elsebranch) == 0
+
+    @property
+    def linear(self) -> List[Instruction]:
+        return self._linear
+
+    @property
+    def thenbranch(self) -> List[Instruction]:
+        return self._thenbranch
+
+    @property
+    def elsebranch(self) -> List[Instruction]:
+        return self._elsebranch
+
+    @property
+    def is_empty(self) -> bool:
+        return (
+            len(self.linear) + len(self.thenbranch) + len(self.elsebranch) == 0)
+
+    def add_predicated_instr(self, instr: Instruction) -> None:
+        if self.is_empty:
+            self._condition = instr.get_instruction_cc()
+            self._setter = instr.get_instruction_condition_setter()
+            self.thenbranch.append(instr)
+        elif self.is_predicated:
+            if self.condition == instr.get_instruction_cc():
+                self.thenbranch.append(instr)
+            else:
+                self.elsebranch.append(instr)
+        else:
+            raise UF.CHBError("Cannot add predicated instruction to linear frag")
+
+    def add_linear_instr(self, instr: Instruction) -> None:
+        if self.is_empty or (not self.is_predicated):
+            self.linear.append(instr)
+        else:
+            raise UF.CHBError(
+                "Cannot add unpredicated instruction to predicate fragment")
+
+    def add_instr(self, instr: Instruction) -> None:
+        if instr.has_control_flow():
+            self.add_predicated_instr(instr)
+        else:
+            self.add_linear_instr(instr)
+
+    def __str__(self) -> str:
+        lines: List[str] = []
+        if self.condition:
+            setter = " (" + self.setter + ")" if self.setter else ""
+            lines.append("condition: " + self.condition + setter)
+        if self.linear:
+            lines.append("linear")
+            for i in self.linear:
+                lines.append("  " + str(i))
+        if self.thenbranch:
+            lines.append("then:")
+            for i in self.thenbranch:
+                lines.append("  " + str(i))
+        if self.elsebranch:
+            lines.append("else:")
+            for i in self.elsebranch:
+                lines.append("  " + str(i))
+        return "\n".join(lines)
+
+
 class BasicBlock(ABC):
 
     def __init__(
             self,
             xnode: ET.Element) -> None:
         self.xnode = xnode
+        self._partition: Dict[str, BasicBlockFragment] = {}
+
+    @property
+    def partition(self) -> Dict[str, BasicBlockFragment]:
+        return self._partition
 
     @property
     def baddr(self) -> str:
@@ -99,6 +222,55 @@ def has_return(self) -> bool:
     def has_conditional_return(self) -> bool:
         return self.last_instruction.is_conditional_return_instruction
 
+    def has_control_flow(self) -> bool:
+        """Returns true if this block contains predicated instructions that
+        are not otherwise covered in aggregates or by other conditions.
+
+        The case of a block with a conditional return is already handled in
+        the Cfg, so it is excluded here.
+
+        The case of a block with a conditional return and other conditional
+        instructions is not yet handled.
+        """
+
+        count = len([i for i in self.instructions.values() if i.has_control_flow()])
+        if count == 1 and self.has_conditional_return:
+            return False
+
+        return any(i.has_control_flow() for i in self.instructions.values())
+
+    def partition_control_flow(self) -> None:
+        curblock: Optional[BasicBlockFragment] = None
+        curaddr: Optional[str] = None
+
+        for (a, i) in sorted(self.instructions.items()):
+            if curaddr is None or curblock is None:
+                curaddr = a
+                curblock = BasicBlockFragment(i)
+            else:
+                if i.has_control_flow():
+                    if curblock.is_predicated:
+                        if i.get_instruction_condition_setter() == curblock.setter:
+                            curblock.add_instr(i)
+                        else:
+                            self._partition[curaddr] = curblock
+                            curblock = BasicBlockFragment(i)
+                            curaddr = a
+                    else:
+                        self._partition[curaddr] = curblock
+                        curblock = BasicBlockFragment(i)
+                        curaddr = a
+                else:
+                    if curblock.is_predicated:
+                        self._partition[curaddr] = curblock
+                        curblock = BasicBlockFragment(i)
+                        curaddr = a
+                    else:
+                        curblock.add_instr(i)
+
+        if curaddr is not None and curblock is not None:
+            self._partition[curaddr] = curblock
+
     @property
     @abstractmethod
     def instructions(self) -> Mapping[str, Instruction]:
diff --git a/chb/app/Cfg.py b/chb/app/Cfg.py
@@ -687,13 +687,14 @@ def ast(self,
                 succ = self.successors(n)[0]
                 instr = astblock.last_instruction
                 rv = instr.return_value()
+                rvcondition = instr.ast_condition(astree)
                 astexpr: Optional[AST.ASTExpr] = None
                 if rv is not None and not astree.returns_void():
                     astexpr = XU.xxpr_to_ast_def_expr(
                         rv, instr.xdata, instr.iaddr, astree)
                     rtnstmt = astree.mk_return_stmt(
                         astexpr, instr.iaddr, instr.bytestring)
-                    rvcondition = instr.ast_condition(astree)
+
                     if rvcondition is not None:
                         elsebr = astree.mk_instr_sequence([])
                         brstmt = cast(AST.ASTBranch, astree.mk_branch(
@@ -702,7 +703,15 @@ def ast(self,
                     else:
                         blockstmts[n] = [blocknode, rtnstmt]
                 else:
-                    blockstmts[n] = [blocknode]
+                    rtnstmt = astree.mk_return_stmt(
+                        None, instr.iaddr, instr.bytestring)
+                    if rvcondition is not None:
+                        elsebr = astree.mk_instr_sequence([])
+                        brstmt = cast(AST.ASTBranch, astree.mk_branch(
+                            rvcondition, rtnstmt, elsebr, succ))
+                        blockstmts[n] = [blocknode, brstmt]
+                    else:
+                        blockstmts[n] = [blocknode, rtnstmt]
 
             elif astblock.has_return:
                 instr = astblock.last_instruction
diff --git a/chb/app/InstrXData.py b/chb/app/InstrXData.py
@@ -642,6 +642,22 @@ def get_return_cxpr(self) -> XXpr:
             raise UF.CHBError("Unexpected error in C return value")
         return self.xprdictionary.xpr(rval)
 
+    @property
+    def is_predicate_assignment(self) -> bool:
+        return "agg:predassign" in self.tags
+
+    @property
+    def is_nondet_predicate_assignment(self) -> bool:
+        return "agg:predassign:nd" in self.tags
+
+    @property
+    def is_ternary_assignment(self) -> bool:
+        return "agg:ternassign" in self.tags
+
+    @property
+    def is_nondet_ternary_assignment(self) -> bool:
+        return "agg:ternassign:nd" in self.tags
+
     @property
     def is_aggregate_jumptable(self) -> bool:
         return "agg-jt" in self.tags
diff --git a/chb/app/Instruction.py b/chb/app/Instruction.py
@@ -4,7 +4,7 @@
 # ------------------------------------------------------------------------------
 # The MIT License (MIT)
 #
-# Copyright (c) 2021-2024  Aarno Labs LLC
+# 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
@@ -145,6 +145,18 @@ def rev_bytestring(self) -> str:
         revb = "".join(i+j for i, j in zip(b[:-1][::-2], b[::-2]))
         return revb
 
+    def has_control_flow(self) -> bool:
+        """Returns true if this instruction is predicated and not covered
+        by an enclosing aggregate or other condition."""
+
+        return self.xnode.get("brcc") is not None
+
+    def get_instruction_cc(self) -> Optional[str]:
+        return self.xnode.get("brcc")
+
+    def get_instruction_condition_setter(self) -> Optional[str]:
+        return self.xnode.get("brsetter")
+
     def md5(self) -> str:
         m = hashlib.md5()
         m.update(self.bytestring.encode("utf-8"))
@@ -292,10 +304,16 @@ def ast_prov(self, astree: ASTInterface) -> Tuple[
     def is_condition_true(self) -> bool:
         return False
 
-    def ast_condition_prov(self, astree: ASTInterface, reverse: bool = False) -> Tuple[
-            Optional[AST.ASTExpr], Optional[AST.ASTExpr]]:
+    def ast_condition_prov(
+            self, astree: ASTInterface, reverse: bool = False
+    ) -> Tuple[Optional[AST.ASTExpr], Optional[AST.ASTExpr]]:
         raise UF.CHBError("ast-condition-prov not defined")
 
+    def ast_cc_condition_prov(
+            self, astree: ASTInterface
+    ) -> Tuple[Optional[AST.ASTExpr], Optional[AST.ASTExpr]]:
+        raise UF.CHBError("ast-cc-codntiion-prov not defined")
+
     def assembly_ast_condition(
             self,
             astree: ASTInterface,
diff --git a/chb/astinterface/ASTInterfaceBasicBlock.py b/chb/astinterface/ASTInterfaceBasicBlock.py
@@ -4,7 +4,7 @@
 # ------------------------------------------------------------------------------
 # The MIT License (MIT)
 #
-# Copyright (c) 2022-2024  Aarno Labs LLC
+# Copyright (c) 2022-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
@@ -26,7 +26,7 @@
 # ------------------------------------------------------------------------------
 """Basic block in an abstract syntax tree."""
 
-from typing import cast, Dict, List, Optional, Set, TYPE_CHECKING
+from typing import cast, Dict, List, Optional, Set, Tuple, TYPE_CHECKING
 
 import chb.ast.ASTNode as AST
 
@@ -40,7 +40,7 @@
 
 if TYPE_CHECKING:
     from chb.arm.ARMCfgBlock import ARMCfgBlock
-    from chb.app.BasicBlock import BasicBlock
+    from chb.app.BasicBlock import BasicBlock, BasicBlockFragment
     from chb.arm.ARMInstruction import ARMInstruction
     from chb.arm.opcodes.ARMLogicalShiftLeft import ARMLogicalShiftLeft
     from chb.arm.opcodes.ARMReverseSubtract import ARMReverseSubtract
@@ -168,12 +168,57 @@ def assembly_ast(self, astree: "ASTInterface") -> AST.ASTStmt:
             instrs.extend(i.assembly_ast(astree))
         return astree.mk_instr_sequence(instrs)
 
+    def ast_fragment(
+            self, astree: "ASTInterface", frag: "BasicBlockFragment") -> AST.ASTStmt:
+        if frag.is_predicated:
+            theninstrs = [ASTInterfaceInstruction(i) for i in frag.thenbranch]
+            elseinstrs = [ASTInterfaceInstruction(i) for i in frag.elsebranch]
+            thenstmt = self.linear_ast(astree, theninstrs)
+            elsestmt = self.linear_ast(astree, elseinstrs)
+            cinstr = theninstrs[0]
+            brcond = cinstr.ast_cc_condition(astree)
+            if brcond is None:
+                chklogger.logger.warning(
+                    "No instruction predicate expression found at address %s",
+                    cinstr.iaddr)
+                brcond = astree.mk_temp_lval_expression()
+            return astree.mk_branch(brcond, thenstmt, elsestmt, "0x0")
+        else:
+            instrs = [ASTInterfaceInstruction(i) for i in frag.linear]
+            return self.linear_ast(astree, instrs)
+
+    def fragmented_ast(self, astree: "ASTInterface") -> AST.ASTStmt:
+
+        if len(self.basicblock.partition) == 0:
+            raise UF.CHBError("Error in fragmented ast")
+
+        stmts: List[AST.ASTStmt] = []
+
+        for (a, bf) in sorted(self.basicblock.partition.items()):
+            stmt = self.ast_fragment(astree, bf)
+            stmts.append(stmt)
+
+        return astree.mk_block(stmts)
+
     def ast(self, astree: "ASTInterface") -> AST.ASTStmt:
         if self.is_trampoline:
             return self.trampoline_ast(astree)
 
+        if self.basicblock.has_control_flow():
+            self.basicblock.partition_control_flow()
+            return self.fragmented_ast(astree)
+
+        else:
+
+            return self.linear_ast(
+                astree, sorted(self.instructions.values(), key = lambda p:p.iaddr))
+
+    def linear_ast(
+            self,
+            astree: "ASTInterface",
+            instritems: List[ASTInterfaceInstruction]) -> AST.ASTStmt:
         instrs: List[AST.ASTInstruction] = []
-        for (a, i) in sorted(self.instructions.items(), key=lambda p: p[0]):
+        for i in instritems:
             instrs.extend(i.ast(astree))
         return astree.mk_instr_sequence(instrs)
 
diff --git a/chb/astinterface/ASTInterfaceInstruction.py b/chb/astinterface/ASTInterfaceInstruction.py
