refactor[venom]: refactor equivalent vars analysis

tests/unit/compiler/venom/test_equivalent_variables.py

@@ -0,0 +1,88 @@
+from collections import defaultdict
+
+from tests.venom_utils import parse_from_basic_block
+from vyper.venom.analysis import IRAnalysesCache, VarEquivalenceAnalysis
+from vyper.venom.basicblock import IRLiteral
+
+
+def _check_expected(code, expected):
+    ctx = parse_from_basic_block(code)
+    fn = next(iter(ctx.functions.values()))
+    ac = IRAnalysesCache(fn)
+    eq = ac.request_analysis(VarEquivalenceAnalysis)
+
+    tmp = defaultdict(list)
+    for var, bag in eq._bags.items():
+        if not isinstance(var, IRLiteral):
+            tmp[bag].append(var)
+
+    ret = []
+    for varset in tmp.values():
+        ret.append(tuple(var.value for var in varset))
+
+    assert tuple(ret) == expected
+
+
+def test_simple_equivalent_vars():
+    code = """
+    main:
+        %1 = 5
+        %2 = %1
+    """
+    expected = (("%1", "%2"),)
+    _check_expected(code, expected)
+
+
+def test_equivalent_vars2():
+    code = """
+    main:
+        # graph with multiple edges from root: %1 => %2 and %1 => %3
+        %1 = 5
+        %2 = %1
+        %3 = %1
+    """
+    expected = (("%1", "%2", "%3"),)
+    _check_expected(code, expected)
+
+
+def test_equivalent_vars3():
+    code = """
+    main:
+        # even weirder graph
+        %1 = 5
+        %2 = %1
+        %3 = %2
+        %4 = %2
+        %5 = %1
+        %6 = 7  ; disjoint
+    """
+    expected = (("%1", "%2", "%3", "%4", "%5"), ("%6",))
+    _check_expected(code, expected)
+
+
+def test_equivalent_vars4():
+    code = """
+    main:
+        # even weirder graph
+        %1 = 5
+        %2 = %1
+        %3 = 5  ; not disjoint, equality on 5
+        %4 = %3
+    """
+    expected = (("%1", "%2", "%3", "%4"),)
+    _check_expected(code, expected)
+
+
+def test_equivalent_vars5():
+    """
+    Test with non-literal roots
+    """
+    code = """
+    main:
+        %1 = param
+        %2 = %1
+        %3 = param  ; disjoint
+        %4 = %3
+    """
+    expected = (("%1", "%2"), ("%3", "%4"))
+    _check_expected(code, expected)

vyper/venom/analysis/equivalent_vars.py

@@ -1,40 +1,79 @@
-from vyper.venom.analysis import DFGAnalysis, IRAnalysis
-from vyper.venom.basicblock import IRVariable
+from typing import Optional
+
+from vyper.venom.analysis import IRAnalysis
+from vyper.venom.basicblock import IRInstruction, IRLiteral, IROperand, IRVariable
 
 
 class VarEquivalenceAnalysis(IRAnalysis):
     """
-    Generate equivalence sets of variables. This is used to avoid swapping
-    variables which are the same during venom_to_assembly. Theoretically,
-    the DFTPass should order variable declarations optimally, but, it is
-    not aware of the "pickaxe" heuristic in venom_to_assembly, so they can
-    interfere.
+    Generate equivalence sets of variables. This is used in passes so that
+    they can "peer through" store chains
     """
 
     def analyze(self):
-        dfg = self.analyses_cache.request_analysis(DFGAnalysis)
+        # map from variables to "equivalence set" of variables, denoted
+        # by "bag" (an int).
+        self._bags: dict[IROperand, int] = {}
+
+        # dict from bags to literal values
+        self._literals: dict[int, IRLiteral] = {}
+
+        # the root of the store chain
+        self._root_instructions: dict[int, IRInstruction] = {}
 
-        equivalence_set: dict[IRVariable, int] = {}
+        bag = 0
+        for bb in self.function.get_basic_blocks():
+            for inst in bb.instructions:
+                if inst.output is None:
+                    continue
+                if inst.opcode != "store":
+                    self._handle_nonstore(inst, bag)
+                else:
+                    self._handle_store(inst, bag)
+                bag += 1
 
-        for bag, (var, inst) in enumerate(dfg._dfg_outputs.items()):
-            if inst.opcode != "store":
-                continue
+    def _handle_nonstore(self, inst: IRInstruction, bag: int):
+        assert inst.output is not None  # help mypy
+        if bag in self._bags:
+            bag = self._bags[inst.output]
+        else:
+            self._bags[inst.output] = bag
+        self._root_instructions[bag] = inst
 
-            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 isinstance(source, (IRVariable, IRLiteral))
+        assert var not in self._bags  # invariant
 
-        self._equivalence_set = equivalence_set
+        if source in self._bags:
+            bag = self._bags[source]
+            self._bags[var] = bag
+        else:
+            self._bags[source] = bag
+            self._bags[var] = bag
 
-    def equivalent(self, var1, var2):
-        if var1 not in self._equivalence_set:
+        if isinstance(source, IRLiteral):
+            self._literals[bag] = source
+
+    def equivalent(self, var1: IROperand, var2: IROperand):
+        if var1 == var2:
+            return True
+        if var1 not in self._bags:
             return False
-        if var2 not in self._equivalence_set:
+        if var2 not in self._bags:
             return False
-        return self._equivalence_set[var1] == self._equivalence_set[var2]
+        return self._bags[var1] == self._bags[var2]
+
+    def get_literal(self, var: IROperand) -> Optional[IRLiteral]:
+        if isinstance(var, IRLiteral):
+            return var
+        if (bag := self._bags.get(var)) is None:
+            return None
+        return self._literals.get(bag)
+
+    def get_root_instruction(self, var: IROperand):
+        bag = self._bags[var]
+        return self._root_instructions[bag]

vyper/venom/passes/load_elimination.py

@@ -22,7 +22,7 @@ def run_pass(self):
         self.analyses_cache.invalidate_analysis(DFGAnalysis)
 
     def equivalent(self, op1, op2):
-        return op1 == op2 or self.equivalence.equivalent(op1, op2)
+        return self.equivalence.equivalent(op1, op2)
 
     def _process_bb(self, bb, eff, load_opcode, store_opcode):
         # not really a lattice even though it is not really inter-basic block;