fix[venom]: fix various well formedness issues

tests/unit/compiler/venom/test_calling_convention.py

@@ -391,7 +391,7 @@ def test_dret_rejects_mixing_ret_and_dret():
 main:
     %p1 = param
     %p2 = param
-    %retpc = param
+    %retpc = retpc_param
     %x = %p1
     %x = %p2
     mstore 0, %x

tests/unit/compiler/venom/test_dalloca.py

@@ -834,7 +834,7 @@ def test_dret_full_pipeline_with_and_without_inlining(env):
 
     function callee {
         callee:
-            %retpc = param
+            %retpc = retpc_param
             %p = dalloca 32
             mstore %p, 123
             dret 1, %p, 32, %retpc
@@ -928,7 +928,7 @@ def test_unreachable_raw_dret_does_not_block_inlining():
 
 function dret_callee {
     dret_callee:
-        %retpc = param
+        %retpc = retpc_param
         %p = dalloca 32
         mstore %p, 7
         dret 1, %p, 32, %retpc
@@ -1133,7 +1133,7 @@ def test_inlined_publishing_callee_host_does_not_publish(env):
 
     function callee {
         callee:
-            %retpc = param
+            %retpc = retpc_param
             %p = dalloca 32
             mstore %p, 123
             dret 1, %p, 32, %retpc
@@ -1758,7 +1758,7 @@ def test_escaped_getfmp_capture_pins_reclaim(env):
 
 function callee {
     callee:
-        %retpc = param
+        %retpc = retpc_param
         %s1 = dalloca 32
         %s2 = dalloca 32
         mstore %s1, 11

tests/unit/compiler/venom/test_inliner.py

@@ -1,10 +1,13 @@
 from tests.venom_utils import parse_venom
 from vyper.compiler.settings import OptimizationLevel, VenomOptimizationFlags
+from vyper.evm.assembler.core import assembly_to_evm
+from vyper.evm.assembler.instructions import PUSHLABEL, Label
 from vyper.venom.analysis.analysis import IRAnalysesCache
 from vyper.venom.analysis.fcg import FCGGlobalAnalysis
 from vyper.venom.basicblock import IRLabel
 from vyper.venom.check_venom import check_venom_ctx
 from vyper.venom.passes import FunctionInlinerPass, SimplifyCFGPass
+from vyper.venom.venom_to_assembly import VenomCompiler
 
 
 def test_inliner_phi_invalidation():
@@ -35,9 +38,10 @@ def test_inliner_phi_invalidation():
 
     function f {
     main:
+        %retpc = retpc_param
         %p = source
         %1 = add %p, 1
-        ret %1
+        ret %retpc, %1
     }
     """
 
@@ -91,9 +95,10 @@ def test_inliner_phi_invalidation_inner():
 
     function f {
     main:
+        %retpc = retpc_param
         %p = source
         %1 = add %p, 1
-        ret %1
+        ret %retpc, %1
     }
     """
 
@@ -152,7 +157,7 @@ def test_inliner_preserves_memory_read_max_size():
         %dst = param
         %src = param
         %size = param
-        %retpc = param
+        %retpc = retpc_param
         mcopy %dst, %src, %size
         ret %retpc
     }
@@ -184,9 +189,9 @@ def test_noinline_annotation():
 
     function f [noinline] {
     f:
-        %retpc = param
+        %retpc = retpc_param
         %1 = 42
-        ret %1, %retpc
+        ret %retpc, %1
     }
     """
 
@@ -203,3 +208,37 @@ def run_inliner(source):
 
     # sanity check: without the annotation, the function gets inlined
     assert IRLabel("f") not in run_inliner(src.replace(" [noinline]", "")).functions
+
+
+def test_inliner_retpc_retain():
+    code = """
+    function main {
+    main:
+        %1 = invoke @f
+        %ptr = alloca 32
+        mstore %ptr, %1
+        return %ptr, 32
+    }
+
+    function f {
+    f:
+        %retpc = retpc_param
+        %1 = 42
+        ret %retpc, %1
+    }
+    """
+
+    flags = VenomOptimizationFlags(level=OptimizationLevel.CODESIZE)
+
+    def run_inliner(source):
+        ctx = parse_venom(source)
+        analyses = {fn: IRAnalysesCache(fn) for fn in ctx.functions.values()}
+        FunctionInlinerPass(analyses, ctx, flags).run_pass()
+        for fn in ctx.functions.values():
+            SimplifyCFGPass(analyses[fn], fn).run_pass()
+        return ctx
+
+    ctx = run_inliner(code)
+    asm = VenomCompiler(ctx).generate_evm_assembly()
+    assert PUSHLABEL(Label("f")) not in asm
+    _ = assembly_to_evm(asm)

tests/unit/compiler/venom/test_removeunused.py

@@ -1,3 +1,5 @@
+import pytest
+
 from tests.venom_utils import assert_ctx_eq, parse_from_basic_block
 from vyper.evm.address_space import MEMORY
 from vyper.venom.analysis import MemSSA
@@ -112,3 +114,19 @@ def test_removeunused_invalidates_memory_ssa():
     # with fresh MemorySSA the store has no reader left, so DSE removes it
     DeadStoreElimination(ac, fn).run_pass(MEMORY)
     assert not any(inst.opcode == "mstore" for inst in fn.entry.instructions)
+
+
+@pytest.mark.parametrize(
+    "opcode", ["balance", "selfbalance", "extcodesize", "extcodehash", "returndatasize"]
+)
+def test_removeunused_non_alias_supported_opcodes(opcode):
+    pre = f"""
+    main:
+        %1 = {opcode}
+        stop
+    """
+    post = """
+    main:
+        stop
+    """
+    _check_pre_post(pre, post)

tests/unit/compiler/venom/test_venom_to_assembly.py

@@ -1,3 +1,4 @@
+from tests.venom_utils import parse_from_basic_block
 from vyper.compiler import compile_code
 from vyper.compiler.settings import Settings
 from vyper.venom.basicblock import IRVariable
@@ -77,6 +78,29 @@ def test_invoke_middle_output_unused():
     assert pop_idx > 0 and asm[pop_idx - 1] == "SWAP1", asm
 
 
+def test_unused_clean_stack_from_cfg_in():
+    code = """
+    main:
+        %unused = offset @else, 32
+        %cond = calldataload 0
+        jnz %cond, @then, @else
+    then:
+        %a = calldataload 0
+        jmp @after
+    else:
+        jmp @after
+    after:
+        %x = phi @else, %cond, @then, %a
+        %ptr = alloca 32
+        mstore %ptr, %x
+        return %ptr, 32
+    """
+
+    ctx = parse_from_basic_block(code)
+    asm = VenomCompiler(ctx).generate_evm_assembly()
+    assert "POP" not in asm
+
+
 def test_popmany_bulk_removal_of_suffix():
     compiler = VenomCompiler(IRContext())
     stack = StackModel()

vyper/venom/analysis/mem_alias.py

@@ -244,3 +244,7 @@ def mem_alias_type_factory(addr_space: AddrSpace) -> type[MemoryAliasAnalysisAbs
         return TransientAliasAnalysis
     else:  # pragma: nocover
         raise ValueError(f"Invalid address space: {addr_space}")
+
+
+def can_create_mem_alias(addr_space: AddrSpace) -> bool:
+    return addr_space in (MEMORY, STORAGE, TRANSIENT)

vyper/venom/passes/branch_optimization.py

@@ -1,6 +1,12 @@
 from vyper.utils import OrderedSet
 from vyper.venom.analysis import CFGAnalysis, DFGAnalysis, LivenessAnalysis
-from vyper.venom.basicblock import COMPARATOR_INSTRUCTIONS, IRBasicBlock, IRInstruction, IRLiteral
+from vyper.venom.basicblock import (
+    COMPARATOR_INSTRUCTIONS,
+    IRBasicBlock,
+    IRInstruction,
+    IRLiteral,
+    IRVariable,
+)
 from vyper.venom.passes.base_pass import InstUpdater, IRPass
 
 
@@ -45,6 +51,8 @@ def _optimize_branches(self) -> None:
             cost_a, cost_b = len(fst_liveness), len(snd_liveness)
 
             cond = term_inst.operands[0]
+            if not isinstance(cond, IRVariable):
+                continue
             prev_inst = self.dfg.get_producing_instruction(cond)
             assert prev_inst is not None
 

vyper/venom/passes/function_inliner.py

@@ -7,7 +7,7 @@
 from vyper.venom.analysis.fcg import FCGGlobalAnalysis
 from vyper.venom.analysis.readonly_memory_args import ReadonlyMemoryArgsGlobalAnalysis
 from vyper.venom.basicblock import IRBasicBlock, IRInstruction, IRLabel, IROperand, IRVariable
-from vyper.venom.call_layout import InvokeLayout, has_dret
+from vyper.venom.call_layout import FunctionCallLayout, InvokeLayout, has_dret
 from vyper.venom.context import IRContext
 from vyper.venom.function import IRFunction
 from vyper.venom.passes.base_pass import IRGlobalPass
@@ -132,13 +132,18 @@ def _inline_call_site(self, func: IRFunction, call_site: IRInstruction) -> None:
         # operands[1:] + [operands[0]] reorder for raw IR.
         binding_ops = InvokeLayout(self.ctx, call_site).bound_params
 
+        layout = FunctionCallLayout(func_copy)
+        retpc_inst = layout.retpc_param_opcode_inst
+        assert retpc_inst is not None
+        retpc_op = retpc_inst.output
+        retpc_inst.make_nop()
         for bb in func_copy.get_basic_blocks():
             bb.parent = call_site_func
             call_site_func.append_basic_block(bb)
             param_idx = 0
             for inst in bb.instructions:
                 if inst.is_param:
-                    # NOTE: one of these params is the return pc.
+                    assert inst.opcode != "retpc_param"
                     inst.opcode = "assign"
                     val = binding_ops[param_idx]
                     inst.operands = [val]
@@ -156,6 +161,7 @@ def _inline_call_site(self, func: IRFunction, call_site: IRInstruction) -> None:
                     # host's own reclaim governs the inlined data. Whether the
                     # host publishes is determined solely by the host's own
                     # terminators (plain `ret` is callee-save).
+                    assert retpc_op in inst.operands, retpc_op
                     ret_values = [op for op in inst.operands[:-1] if not isinstance(op, IRLabel)]
 
                     # Map each returned value to corresponding callsite outputs
@@ -173,6 +179,7 @@ def _inline_call_site(self, func: IRFunction, call_site: IRInstruction) -> None:
 
             for inst in bb.instructions:
                 if not inst.annotation:
+                    assert retpc_op not in inst.operands, (inst, retpc_op)
                     inst.annotation = f"from {func.name}"
 
         call_site_bb.instructions = call_site_bb.instructions[:call_idx]

vyper/venom/passes/remove_unused_variables.py

@@ -1,13 +1,17 @@
+from vyper.evm import address_space
 from vyper.utils import OrderedSet, uniq
+from vyper.venom import effects
 from vyper.venom.analysis import BasePtrAnalysis, DFGAnalysis, LivenessAnalysis
 from vyper.venom.analysis.load_analysis import LoadAnalysis
 from vyper.venom.analysis.mem_alias import (
-    MemoryAliasAnalysis,
     StorageAliasAnalysis,
     TransientAliasAnalysis,
+    can_create_mem_alias,
+    mem_alias_type_factory,
 )
 from vyper.venom.analysis.mem_ssa import MemSSA, StorageSSA, TransientSSA
 from vyper.venom.basicblock import IRInstruction
+from vyper.venom.effects import EMPTY, FMP
 from vyper.venom.passes.base_pass import IRPass
 
 
@@ -19,8 +23,11 @@ class RemoveUnusedVariablesPass(IRPass):
     dfg: DFGAnalysis
     work_list: OrderedSet[IRInstruction]
 
+    invalidate_alias: set[address_space.AddrSpace]
+
     def run_pass(self):
         self.dfg = self.analyses_cache.request_analysis(DFGAnalysis)
+        self.invalidate_alias = set()
 
         work_list = OrderedSet()
         self.work_list = work_list
@@ -45,14 +52,16 @@ def run_pass(self):
         # invalidations below only cascade to them when the parent is
         # actually cached, but the alias analyses can also be requested
         # (and cached) on their own
-        self.analyses_cache.invalidate_analysis(MemoryAliasAnalysis)
         self.analyses_cache.invalidate_analysis(StorageAliasAnalysis)
         self.analyses_cache.invalidate_analysis(TransientAliasAnalysis)
         self.analyses_cache.invalidate_analysis(LoadAnalysis)
         self.analyses_cache.invalidate_analysis(MemSSA)
         self.analyses_cache.invalidate_analysis(StorageSSA)
         self.analyses_cache.invalidate_analysis(TransientSSA)
         self.analyses_cache.invalidate_analysis(LivenessAnalysis)
+        for space in self.invalidate_alias:
+            alias_analysis = mem_alias_type_factory(space)
+            self.analyses_cache.invalidate_analysis(alias_analysis)
 
     def _process_instruction(self, inst) -> bool:
         outputs = inst.get_outputs()
@@ -72,5 +81,22 @@ def _process_instruction(self, inst) -> bool:
             new_uses = self.dfg.get_uses(operand)
             self.work_list.addmany(new_uses)
 
+        # instructions that handle FMP can be removed if there is no use for the
+        # since they either bump or set only at the end of the function so the
+        # removal of the instruction will not effect other instructions.
+        # Other write effect should not be removed by this analysis
+        assert inst.get_write_effects() == EMPTY or inst.get_write_effects() == FMP
+        effs = inst.get_read_effects()
+        if effs != EMPTY:
+            for eff in effs:
+                space = effects.to_addr_space(eff)
+                if space is None:
+                    # sanity check
+                    continue
+
+                # Mem alias does not use all address spaces
+                if can_create_mem_alias(space):
+                    self.invalidate_alias.add(space)
+
         inst.make_nop()
         return True

vyper/venom/venom_to_assembly.py

@@ -645,7 +645,13 @@ def clean_stack_from_cfg_in(
         # for making sure its output stack layout works no matter which
         # bb it jumps into).
         layout = self.liveness.out_vars(in_bb)
-        to_pop = list(layout.difference(inputs))
+        to_pop = []
+        for item in stack._stack:
+            if not isinstance(item, IRVariable):
+                continue
+            if item not in inputs and item not in layout:
+                to_pop.append(item)
+        to_pop.extend(list(layout.difference(inputs)))
         self._assert_dead_stack_prefix(stack)
         physical_to_pop = [
             var for var in to_pop if stack.get_depth(var) is not StackModel.NOT_IN_STACK