feat[venom]: offset instruction (#4180)

this commit introduces an `offset` instruction that is emitted in the
algebraic pass when the add instruction calculates an offset from a
code label, which is used for immutables. this allows compilation
directly to the magic `OFST` assembly instruction, which does
additional constant folding after symbol resolution.

---------

Co-authored-by: Charles Cooper <[email protected]>

tests/unit/compiler/venom/test_algebraic_optimizer.py

@@ -127,3 +127,54 @@ def test_interleaved_case(interleave_point):
         assert bb.instructions[-1].operands[0] == op3_inv
     else:
         assert bb.instructions[-1].operands[0] == op3
+
+
+def test_offsets():
+    ctx = IRContext()
+    fn = ctx.create_function("_global")
+
+    bb = fn.get_basic_block()
+
+    br1 = IRBasicBlock(IRLabel("then"), fn)
+    fn.append_basic_block(br1)
+    br2 = IRBasicBlock(IRLabel("else"), fn)
+    fn.append_basic_block(br2)
+
+    p1 = bb.append_instruction("param")
+    op1 = bb.append_instruction("store", 32)
+    op2 = bb.append_instruction("add", 0, IRLabel("mem"))
+    op3 = bb.append_instruction("store", 64)
+    bb.append_instruction("dloadbytes", op1, op2, op3)
+    op5 = bb.append_instruction("mload", op3)
+    op6 = bb.append_instruction("iszero", op5)
+    bb.append_instruction("jnz", op6, br1.label, br2.label)
+
+    op01 = br1.append_instruction("store", 32)
+    op02 = br1.append_instruction("add", 0, IRLabel("mem"))
+    op03 = br1.append_instruction("store", 64)
+    br1.append_instruction("dloadbytes", op01, op02, op03)
+    op05 = br1.append_instruction("mload", op03)
+    op06 = br1.append_instruction("iszero", op05)
+    br1.append_instruction("return", p1, op06)
+
+    op11 = br2.append_instruction("store", 32)
+    op12 = br2.append_instruction("add", 0, IRLabel("mem"))
+    op13 = br2.append_instruction("store", 64)
+    br2.append_instruction("dloadbytes", op11, op12, op13)
+    op15 = br2.append_instruction("mload", op13)
+    op16 = br2.append_instruction("iszero", op15)
+    br2.append_instruction("return", p1, op16)
+
+    ac = IRAnalysesCache(fn)
+    MakeSSA(ac, fn).run_pass()
+    AlgebraicOptimizationPass(ac, fn).run_pass()
+    RemoveUnusedVariablesPass(ac, fn).run_pass()
+
+    offset_count = 0
+    for bb in fn.get_basic_blocks():
+        for instruction in bb.instructions:
+            assert instruction.opcode != "add"
+            if instruction.opcode == "offset":
+                offset_count += 1
+
+    assert offset_count == 3

vyper/venom/passes/algebraic_optimization.py

@@ -1,6 +1,6 @@
 from vyper.venom.analysis.dfg import DFGAnalysis
 from vyper.venom.analysis.liveness import LivenessAnalysis
-from vyper.venom.basicblock import IRInstruction, IROperand
+from vyper.venom.basicblock import IRInstruction, IRLabel, IRLiteral, IROperand
 from vyper.venom.passes.base_pass import IRPass
 
 
@@ -58,10 +58,24 @@ def _get_iszero_chain(self, op: IROperand) -> list[IRInstruction]:
         chain.reverse()
         return chain
 
+    def _handle_offsets(self):
+        for bb in self.function.get_basic_blocks():
+            for inst in bb.instructions:
+                # check if the instruction is of the form
+                # `add <ptr> <label>`
+                # this works only if store chains have been eliminated.
+                if (
+                    inst.opcode == "add"
+                    and isinstance(inst.operands[0], IRLiteral)
+                    and isinstance(inst.operands[1], IRLabel)
+                ):
+                    inst.opcode = "offset"
+
     def run_pass(self):
         self.dfg = self.analyses_cache.request_analysis(DFGAnalysis)
 
         self._optimize_iszero_chains()
+        self._handle_offsets()
 
         self.analyses_cache.invalidate_analysis(DFGAnalysis)
         self.analyses_cache.invalidate_analysis(LivenessAnalysis)

vyper/venom/passes/extract_literals.py

@@ -20,7 +20,7 @@ def _process_bb(self, bb):
         i = 0
         while i < len(bb.instructions):
             inst = bb.instructions[i]
-            if inst.opcode == "store":
+            if inst.opcode in ("store", "offset"):
                 i += 1
                 continue
 

vyper/venom/venom_to_assembly.py

@@ -401,6 +401,12 @@ def _generate_evm_for_instruction(
                 stack.poke(depth, ret)
             return apply_line_numbers(inst, assembly)
 
+        if opcode == "offset":
+            assembly.extend(["_OFST", f"_sym_{inst.operands[1].value}", inst.operands[0].value])
+            assert isinstance(inst.output, IROperand), "Offset must have output"
+            stack.push(inst.output)
+            return apply_line_numbers(inst, assembly)
+
         # Step 2: Emit instruction's input operands
         self._emit_input_operands(assembly, inst, operands, stack)