From 6e68179c4fe5616af49f8219e7bc967f87cbed91 Mon Sep 17 00:00:00 2001
From: Fish <ltfish@users.noreply.github.com>
Date: Mon, 14 Oct 2024 22:43:38 -0600
Subject: [PATCH] AILBlockWalker: Implement Phi handler. (#247)

* AILBlockWalker: Implement Phi handler.

* Handle cases where vvar is None.

* Default to the read-only version of Phi handler.
---
 ailment/block_walker.py | 32 +++++++++++++++++++++++++++++++-
 ailment/expression.py   |  2 +-
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/ailment/block_walker.py b/ailment/block_walker.py
index e0a6203..9898084 100644
--- a/ailment/block_walker.py
+++ b/ailment/block_walker.py
@@ -198,11 +198,17 @@ class AILBlockWalker(AILBlockWalkerBase):
 
     :ivar update_block: True if the block should be updated in place, False if a new block should be created and
                         returned as the result of walk().
+    :ivar replace_phi_stmt: True if you want _handle_Phi be called and vvars potentially replaced; False otherwise.
+                            Default to False because in the most majority cases you do not want vvars in a Phi
+                            variable be replaced.
     """
 
-    def __init__(self, stmt_handlers=None, expr_handlers=None, update_block: bool = True):
+    def __init__(
+        self, stmt_handlers=None, expr_handlers=None, update_block: bool = True, replace_phi_stmt: bool = False
+    ):
         super().__init__(stmt_handlers=stmt_handlers, expr_handlers=expr_handlers)
         self._update_block = update_block
+        self._replace_phi_stmt = replace_phi_stmt
 
     def walk(self, block: Block) -> Block | None:
         """
@@ -537,6 +543,30 @@ def _handle_ITE(self, expr_idx: int, expr: ITE, stmt_idx: int, stmt: Statement,
             return new_expr
         return None
 
+    def _handle_Phi(self, expr_id: int, expr: Phi, stmt_idx: int, stmt: Statement, block: Block | None) -> Phi | None:
+        if not self._replace_phi_stmt:
+            # fallback to the read-only version
+            return super()._handle_Phi(expr_id, expr, stmt_idx, stmt, block)
+
+        changed = False
+
+        src_and_vvars = None
+        for idx, (src, vvar) in enumerate(expr.src_and_vvars):
+            if vvar is None:
+                if src_and_vvars is not None:
+                    src_and_vvars.append((src, None))
+                continue
+            new_vvar = self._handle_expr(idx, vvar, stmt_idx, stmt, block)
+            if new_vvar is not None and new_vvar is not vvar:
+                changed = True
+                if src_and_vvars is None:
+                    src_and_vvars = expr.src_and_vvars[:idx]
+                src_and_vvars.append((src, new_vvar))
+            elif src_and_vvars is not None:
+                src_and_vvars.append((src, vvar))
+
+        return Phi(expr.idx, expr.bits, src_and_vvars, **expr.tags) if changed else None
+
     def _handle_DirtyExpression(
         self, expr_idx: int, expr: DirtyExpression, stmt_idx: int, stmt: Statement, block: Block | None
     ):
diff --git a/ailment/expression.py b/ailment/expression.py
index 36b8eb7..979e4a4 100644
--- a/ailment/expression.py
+++ b/ailment/expression.py
@@ -324,7 +324,7 @@ def __init__(
         self,
         idx,
         bits,
-        src_and_vvars: list[tuple[tuple[int, int], VirtualVariable]],
+        src_and_vvars: list[tuple[tuple[int, int], VirtualVariable | None]],
         **kwargs,
     ):
         super().__init__(idx, **kwargs)