From 81a382612b7ddeed6b64ef94302a06c0c83744e8 Mon Sep 17 00:00:00 2001
From: Fish <fishw@asu.edu>
Date: Mon, 14 Oct 2024 19:10:49 -0700
Subject: [PATCH 1/3] AILBlockWalker: Implement Phi handler.

---
 ailment/block_walker.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/ailment/block_walker.py b/ailment/block_walker.py
index e0a6203..cfc0f7e 100644
--- a/ailment/block_walker.py
+++ b/ailment/block_walker.py
@@ -537,6 +537,22 @@ 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:
+        changed = False
+
+        src_and_vvars = None
+        for idx, (src, vvar) in enumerate(expr.src_and_vvars):
+            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
     ):

From e57943727e9035131d53474779fece26a77544f9 Mon Sep 17 00:00:00 2001
From: Fish <fishw@asu.edu>
Date: Mon, 14 Oct 2024 19:36:56 -0700
Subject: [PATCH 2/3] Handle cases where vvar is None.

---
 ailment/block_walker.py | 4 ++++
 ailment/expression.py   | 2 +-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/ailment/block_walker.py b/ailment/block_walker.py
index cfc0f7e..7b57926 100644
--- a/ailment/block_walker.py
+++ b/ailment/block_walker.py
@@ -542,6 +542,10 @@ def _handle_Phi(self, expr_id: int, expr: Phi, stmt_idx: int, stmt: Statement, b
 
         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
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)

From 5dfa7126e95d1cf0a362a5c77b8719c3b58a1580 Mon Sep 17 00:00:00 2001
From: Fish <fishw@asu.edu>
Date: Mon, 14 Oct 2024 21:26:16 -0700
Subject: [PATCH 3/3] Default to the read-only version of Phi handler.

---
 ailment/block_walker.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/ailment/block_walker.py b/ailment/block_walker.py
index 7b57926..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:
         """
@@ -538,6 +544,10 @@ def _handle_ITE(self, expr_idx: int, expr: ITE, stmt_idx: int, stmt: Statement,
         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