Implement the first version of constant propagation support. (#426)

* Implement the first version of constant propagation support.

pyvex/block.py

@@ -5,6 +5,7 @@
 
 from . import expr, stmt
 from .const import get_type_size
+from .const_val import ConstVal
 from .data_ref import DataRef
 from .enums import VEXObject
 from .errors import SkipStatementsError
@@ -63,11 +64,13 @@ class IRSB(VEXObject):
         "default_exit_target",
         "_instruction_addresses",
         "data_refs",
+        "const_vals",
     )
 
     # The following constants shall match the defs in pyvex.h
     MAX_EXITS = 400
     MAX_DATA_REFS = 2000
+    MAX_CONST_VALS = 1000
 
     def __init__(
         self,
@@ -135,6 +138,7 @@ def __init__(
         self._exit_statements: tuple[tuple[int, int, IRStmt], ...] | None = None
         self.default_exit_target = None
         self.data_refs = ()
+        self.const_vals = ()
         self._instruction_addresses: tuple[int, ...] = ()
 
         if data is not None:
@@ -578,9 +582,16 @@ def _from_c(self, lift_r, skip_stmts=False):
         self.data_refs = None
         if lift_r.data_ref_count > 0:
             if lift_r.data_ref_count > self.MAX_DATA_REFS:
-                raise SkipStatementsError("data_ref_count exceeded MAX_DATA_REFS (%d)" % self.MAX_DATA_REFS)
+                raise SkipStatementsError(f"data_ref_count exceeded MAX_DATA_REFS ({self.MAX_DATA_REFS})")
             self.data_refs = [DataRef.from_c(lift_r.data_refs[i]) for i in range(lift_r.data_ref_count)]
 
+        # Const values
+        self.const_vals = None
+        if lift_r.const_val_count > 0:
+            if lift_r.const_val_count > self.MAX_CONST_VALS:
+                raise SkipStatementsError(f"const_val_count exceeded MAX_CONST_VALS ({self.MAX_CONST_VALS})")
+            self.const_vals = [ConstVal.from_c(lift_r.const_vals[i]) for i in range(lift_r.const_val_count)]
+
     def _set_attributes(
         self,
         statements=None,

pyvex/const_val.py

@@ -0,0 +1,26 @@
+class ConstVal:
+    """
+    A constant value object. Indicates a constant value assignment to a VEX tmp variable.
+
+    :ivar tmp:          The tmp variable being assigned to.
+    :ivar value:        The value of the tmp variable.
+    :ivar stmt_idx:     The IRSB statement index containing the data access
+    """
+
+    __slots__ = (
+        "tmp",
+        "value",
+        "stmt_idx",
+    )
+
+    def __init__(self, tmp: int, value: int, stmt_idx: int):
+        self.tmp = tmp
+        self.value = value
+        self.stmt_idx = stmt_idx
+
+    def __repr__(self):
+        return f"<ConstVal {self.tmp} = {self.value:#x} @ {self.stmt_idx}>"
+
+    @classmethod
+    def from_c(cls, r):
+        return cls(r.tmp, r.value, r.stmt_idx)

pyvex/lifting/libvex.py

@@ -78,10 +78,6 @@ def _lift(self):
             if strict_block_end is None:
                 strict_block_end = True
 
-            collect_data_refs = 1 if self.collect_data_refs else 0
-            if collect_data_refs != 0 and self.load_from_ro_regions:
-                collect_data_refs |= 2  # the second bit stores load_from_ro_regions
-
             if self.cross_insn_opt:
                 px_control = VexRegisterUpdates.VexRegUpdUnwindregsAtMemAccess
             else:
@@ -99,7 +95,9 @@ def _lift(self):
                 self.traceflags,
                 self.allow_arch_optimizations,
                 strict_block_end,
-                collect_data_refs,
+                1 if self.collect_data_refs else 0,
+                1 if self.load_from_ro_regions else 0,
+                1 if self.const_prop else 0,
                 px_control,
                 self.bytes_offset,
             )

pyvex/lifting/lift_function.py

@@ -34,6 +34,7 @@ def lift(
     collect_data_refs=False,
     cross_insn_opt=True,
     load_from_ro_regions=False,
+    const_prop=False,
 ):
     """
     Recursively lifts blocks using the registered lifters and postprocessors. Tries each lifter in the order in
@@ -144,6 +145,7 @@ def lift(
                     collect_data_refs=collect_data_refs,
                     cross_insn_opt=cross_insn_opt,
                     load_from_ro_regions=load_from_ro_regions,
+                    const_prop=const_prop,
                 )
             except SkipStatementsError:
                 assert skip_stmts is True
@@ -160,6 +162,7 @@ def lift(
                     collect_data_refs=collect_data_refs,
                     cross_insn_opt=cross_insn_opt,
                     load_from_ro_regions=load_from_ro_regions,
+                    const_prop=const_prop,
                 )
             break
         except LiftingException as ex:
@@ -210,6 +213,8 @@ def lift(
                 strict_block_end=strict_block_end,
                 skip_stmts=False,
                 collect_data_refs=collect_data_refs,
+                load_from_ro_regions=load_from_ro_regions,
+                const_prop=const_prop,
             )
 
         next_addr = addr + final_irsb.size
@@ -235,6 +240,8 @@ def lift(
                 inner=True,
                 skip_stmts=False,
                 collect_data_refs=collect_data_refs,
+                load_from_ro_regions=load_from_ro_regions,
+                const_prop=const_prop,
             )
             if more_irsb.size:
                 # Successfully decoded more bytes
@@ -273,6 +280,8 @@ def lift(
                     inner=inner,
                     skip_stmts=False,
                     collect_data_refs=collect_data_refs,
+                    load_from_ro_regions=load_from_ro_regions,
+                    const_prop=const_prop,
                 )
             except LiftingException:
                 continue

pyvex/lifting/lifter.py

@@ -21,6 +21,7 @@ class Lifter:
         "addr",
         "cross_insn_opt",
         "load_from_ro_regions",
+        "const_prop",
         "disasm",
         "dump_irsb",
     )
@@ -63,6 +64,7 @@ def lift(
         collect_data_refs: bool = False,
         cross_insn_opt: bool = True,
         load_from_ro_regions: bool = False,
+        const_prop: bool = False,
         disasm: bool = False,
         dump_irsb: bool = False,
     ):
@@ -102,6 +104,7 @@ def lift(
         self.irsb = irsb
         self.cross_insn_opt = cross_insn_opt
         self.load_from_ro_regions = load_from_ro_regions
+        self.const_prop = const_prop
         self.disasm = disasm
         self.dump_irsb = dump_irsb
         self._lift()