From c4e702665e82e1bfe3ce9272b83e6fa639e78aff Mon Sep 17 00:00:00 2001
From: Charles Cooper <cooper.charles.m@gmail.com>
Date: Tue, 17 Dec 2024 22:42:04 -0500
Subject: [PATCH] allow peering through to original value

---
 vyper/venom/analysis/equivalent_vars.py | 47 ++++++++++++++++---------
 1 file changed, 31 insertions(+), 16 deletions(-)

diff --git a/vyper/venom/analysis/equivalent_vars.py b/vyper/venom/analysis/equivalent_vars.py
index 895895651a..4a16d1c483 100644
--- a/vyper/venom/analysis/equivalent_vars.py
+++ b/vyper/venom/analysis/equivalent_vars.py
@@ -1,5 +1,5 @@
-from vyper.venom.analysis import DFGAnalysis, IRAnalysis
-from vyper.venom.basicblock import IRVariable
+from vyper.venom.analysis import IRAnalysis
+from vyper.venom.basicblock import IRInstruction, IRLiteral, IROperand
 
 
 class VarEquivalenceAnalysis(IRAnalysis):
@@ -12,25 +12,35 @@ class VarEquivalenceAnalysis(IRAnalysis):
     """
 
     def analyze(self):
-        dfg = self.analyses_cache.request_analysis(DFGAnalysis)
+        self._equivalence_set: dict[IROperand, int] = {}
 
-        equivalence_set: dict[IRVariable, int] = {}
+        # dict from bags to literal values
+        self._literals: dict[int, IRLiteral] = {}
 
-        for bag, (var, inst) in enumerate(dfg._dfg_outputs.items()):
-            if inst.opcode != "store":
-                continue
+        bag = 0
+        for bb in self.function.get_basic_blocks():
+            for inst in bb.instructions:
+                if inst.opcode != "store":
+                    continue
+                self._handle_store(inst, bag)
+                bag += 1
 
-            source = inst.operands[0]
+    def _handle_store(self, inst: IRInstruction, bag: int):
+        var = inst.output
+        source = inst.operands[0]
 
-            assert var not in equivalence_set  # invariant
-            if source in equivalence_set:
-                equivalence_set[var] = equivalence_set[source]
-                continue
-            else:
-                equivalence_set[var] = bag
-                equivalence_set[source] = bag
+        assert var is not None  # help mypy
+        assert var not in self._equivalence_set  # invariant
 
-        self._equivalence_set = equivalence_set
+        if source in self._equivalence_set:
+            bag = self._equivalence_set[source]
+            self._equivalence_set[var] = bag
+        else:
+            self._equivalence_set[source] = bag
+            self._equivalence_set[var] = bag
+
+        if isinstance(source, IRLiteral):
+            self._literals[bag] = source
 
     def equivalent(self, var1, var2):
         if var1 not in self._equivalence_set:
@@ -38,3 +48,8 @@ def equivalent(self, var1, var2):
         if var2 not in self._equivalence_set:
             return False
         return self._equivalence_set[var1] == self._equivalence_set[var2]
+
+    def get_literal(self, var):
+        if (bag := self._equivalence_set.get(var)) is None:
+            return None
+        return self._literals.get(bag)