feat[venom]: stack spilling to mem

tests/functional/codegen/features/test_constructor.py

@@ -3,7 +3,6 @@
 import pytest
 
 from tests.evm_backends.base_env import _compile
-from vyper.exceptions import StackTooDeep
 from vyper.utils import method_id
 
 
@@ -216,7 +215,6 @@ def get_foo() -> DynArray[DynArray[uint256, 3], 3]:
     assert c.get_foo() == [[37, 41, 73], [37041, 41073, 73037], [146, 123, 148]]
 
 
-@pytest.mark.venom_xfail(raises=StackTooDeep, reason="stack scheduler regression")
 def test_initialise_nested_dynamic_array_2(env, get_contract):
     code = """
 foo: DynArray[DynArray[DynArray[int128, 3], 3], 3]

tests/functional/codegen/features/test_immutable.py

@@ -1,7 +1,6 @@
 import pytest
 
 from vyper.compiler.settings import OptimizationLevel
-from vyper.exceptions import StackTooDeep
 
 
 @pytest.mark.parametrize(
@@ -199,7 +198,6 @@ def get_idx_two() -> uint256:
     assert c.get_idx_two() == expected_values[2][2]
 
 
-@pytest.mark.venom_xfail(raises=StackTooDeep, reason="stack scheduler regression")
 def test_nested_dynarray_immutable(get_contract):
     code = """
 my_list: immutable(DynArray[DynArray[DynArray[int128, 3], 3], 3])

tests/functional/codegen/features/test_transient.py

@@ -2,7 +2,7 @@
 
 from tests.utils import ZERO_ADDRESS
 from vyper.compiler import compile_code
-from vyper.exceptions import EvmVersionException, StackTooDeep, VyperException
+from vyper.exceptions import EvmVersionException, VyperException
 
 pytestmark = pytest.mark.requires_evm_version("cancun")
 
@@ -343,7 +343,6 @@ def get_idx_two(_a: uint256, _b: uint256, _c: uint256) -> uint256:
     assert c.get_idx_two(*values) == expected_values[2][2]
 
 
-@pytest.mark.venom_xfail(raises=StackTooDeep, reason="stack scheduler regression")
 def test_nested_dynarray_transient(get_contract, tx_failed, env):
     set_list = """
     self.my_list = [

tests/functional/codegen/types/test_lists.py

@@ -7,7 +7,7 @@
 from tests.utils import check_precompile_asserts, decimal_to_int
 from vyper.compiler.settings import OptimizationLevel
 from vyper.evm.opcodes import version_check
-from vyper.exceptions import ArrayIndexException, OverflowException, StackTooDeep, TypeMismatch
+from vyper.exceptions import ArrayIndexException, OverflowException, TypeMismatch
 
 
 def _map_nested(f, xs):
@@ -597,7 +597,6 @@ def bar(_baz: Foo[3]) -> String[96]:
     assert c.bar(c_input) == "Hello world!!!!"
 
 
-@pytest.mark.venom_xfail(raises=StackTooDeep, reason="stack scheduler regression")
 def test_list_of_nested_struct_arrays(get_contract):
     code = """
 struct Ded:

tests/unit/compiler/venom/test_stack_spill.py

@@ -0,0 +1,261 @@
+import pytest
+
+from vyper.ir.compile_ir import Label
+from vyper.venom.basicblock import IRLiteral, IRVariable
+from vyper.venom.context import IRContext
+from vyper.venom.parser import parse_venom
+from vyper.venom.stack_model import StackModel
+from vyper.venom.stack_spiller import StackSpiller
+from vyper.venom.venom_to_assembly import VenomCompiler
+
+
+def _build_stack(count: int) -> tuple[StackModel, list[IRLiteral]]:
+    stack = StackModel()
+    ops = [IRLiteral(i) for i in range(count)]
+    for op in ops:
+        stack.push(op)
+    return stack, ops
+
+
+def _ops_only_strings(assembly) -> list[str]:
+    return [op for op in assembly if isinstance(op, str)]
+
+
+def _dummy_dfg():
+    class _DummyDFG:
+        def are_equivalent(self, a, b):
+            return False
+
+    return _DummyDFG()
+
+
+def test_swap_spills_deep_stack() -> None:
+    compiler = VenomCompiler(IRContext())
+    stack, ops = _build_stack(40)
+    assembly: list = []
+
+    target = ops[-18]
+    before = stack._stack.copy()
+
+    depth = stack.get_depth(target)
+    assert isinstance(depth, int) and depth < -16
+    swap_idx = -depth
+
+    compiler.spiller.swap(assembly, stack, depth)
+
+    expected = before.copy()
+    top_index = len(expected) - 1
+    target_index = expected.index(target)
+    expected[top_index], expected[target_index] = expected[target_index], expected[top_index]
+    assert stack._stack == expected
+
+    ops_str = _ops_only_strings(assembly)
+    assert ops_str.count("MSTORE") == swap_idx + 1
+    assert ops_str.count("MLOAD") == swap_idx + 1
+    assert all(int(op[4:]) <= 16 for op in ops_str if op.startswith("SWAP"))
+
+
+def test_dup_spills_deep_stack() -> None:
+    compiler = VenomCompiler(IRContext())
+    stack, ops = _build_stack(40)
+    assembly: list = []
+
+    target = ops[-18]
+    before = stack._stack.copy()
+
+    depth = stack.get_depth(target)
+    assert isinstance(depth, int) and depth < -16
+    dup_idx = 1 - depth
+
+    compiler.spiller.dup(assembly, stack, depth)
+
+    expected = before.copy()
+    expected.append(target)
+    assert stack._stack == expected
+
+    ops_str = _ops_only_strings(assembly)
+    assert ops_str.count("MSTORE") == dup_idx
+    assert ops_str.count("MLOAD") == dup_idx + 1
+    assert all(int(op[3:]) <= 16 for op in ops_str if op.startswith("DUP"))
+
+
+def test_stack_reorder_spills_before_swap() -> None:
+    ctx = IRContext()
+    compiler = VenomCompiler(ctx)
+    compiler.dfg = _dummy_dfg()
+
+    compiler.spiller = StackSpiller(ctx, initial_offset=0x10000)
+
+    stack = StackModel()
+    vars_on_stack = [IRVariable(f"%v{i}") for i in range(40)]
+    for var in vars_on_stack:
+        stack.push(var)
+
+    spilled: dict = {}
+    assembly: list = []
+
+    target = vars_on_stack[21]  # depth 18 from top for 40 items
+
+    compiler._stack_reorder(assembly, stack, [target], spilled, dry_run=False)
+
+    assert stack.get_depth(target) == 0
+    assert len(spilled) == 2  # spilled top two values to reduce depth to <= 16
+
+    ops_str = _ops_only_strings(assembly)
+    assert ops_str.count("MSTORE") == 2
+    assert all(int(op[4:]) <= 16 for op in ops_str if op.startswith("SWAP"))
+
+    # restoring a spilled variable should reload it via MLOAD
+    restore_assembly: list = []
+    spilled_var = next(iter(spilled))
+    compiler.spiller.restore_spilled_operand(restore_assembly, stack, spilled, spilled_var)
+    restore_ops = _ops_only_strings(restore_assembly)
+    assert restore_ops.count("MLOAD") == 1
+    assert spilled_var not in spilled
+    assert stack.get_depth(spilled_var) == 0
+
+
+def test_branch_spill_integration() -> None:
+    venom_src = """
+    function spill_demo {
+        main:
+            %v0 = mload 0
+            %v1 = mload 32
+            %v2 = mload 64
+            %v3 = mload 96
+            %v4 = mload 128
+            %v5 = mload 160
+            %v6 = mload 192
+            %v7 = mload 224
+            %v8 = mload 256
+            %v9 = mload 288
+            %v10 = mload 320
+            %v11 = mload 352
+            %v12 = mload 384
+            %v13 = mload 416
+            %v14 = mload 448
+            %v15 = mload 480
+            %v16 = mload 512
+            %v17 = mload 544
+            %v18 = mload 576
+            %v19 = mload 608
+            %cond = mload 640
+            jnz %cond, @then, @else
+        then:
+            %then_sum = add %v0, %v19
+            %res_then = add %then_sum, %cond
+            jmp @join
+        else:
+            %else_sum = add %v1, %v19
+            %res_else = add %else_sum, %cond
+            jmp @join
+        join:
+            %phi = phi @then, %res_then, @else, %res_else
+            %acc1 = add %phi, %v1
+            %acc2 = add %acc1, %v2
+            %acc3 = add %acc2, %v3
+            %acc4 = add %acc3, %v4
+            %acc5 = add %acc4, %v5
+            %acc6 = add %acc5, %v6
+            %acc7 = add %acc6, %v7
+            %acc8 = add %acc7, %v8
+            %acc9 = add %acc8, %v9
+            %acc10 = add %acc9, %v10
+            %acc11 = add %acc10, %v11
+            %acc12 = add %acc11, %v12
+            %acc13 = add %acc12, %v13
+            %acc14 = add %acc13, %v14
+            %acc15 = add %acc14, %v15
+            %acc16 = add %acc15, %v16
+            %acc17 = add %acc16, %v17
+            %acc18 = add %acc17, %v18
+            return %acc18
+    }
+    """
+
+    ctx = parse_venom(venom_src)
+    compiler = VenomCompiler(ctx)
+    compiler.generate_evm_assembly()
+
+    fn = next(iter(ctx.functions.values()))
+    assert any(inst.opcode == "alloca" for inst in fn.entry.instructions)
+
+    asm = compiler.generate_evm_assembly()
+    opcodes = [op for op in asm if isinstance(op, str)]
+
+    for op in opcodes:
+        if op.startswith("SWAP"):
+            assert int(op[4:]) <= 16
+        if op.startswith("DUP"):
+            assert int(op[3:]) <= 16
+
+    def _find_spill_ops(kind: str) -> list[int]:
+        matches: list[int] = []
+        idx = 0
+        while idx < len(asm):
+            item = asm[idx]
+            if isinstance(item, str) and item.startswith("PUSH"):
+                try:
+                    push_bytes = int(item[4:])
+                except ValueError:
+                    push_bytes = 0
+                target_idx = idx + 1 + push_bytes
+                if target_idx < len(asm) and asm[target_idx] == kind:
+                    matches.append(idx)
+                idx = target_idx + 1
+            else:
+                idx += 1
+        return matches
+
+    store_indices = _find_spill_ops("MSTORE")
+    load_indices = _find_spill_ops("MLOAD")
+    assert store_indices
+    assert load_indices
+
+    join_idx = next(
+        idx for idx, op in enumerate(asm) if isinstance(op, Label) and str(op) == "LABEL join"
+    )
+
+    assert any(idx < join_idx for idx in store_indices)
+    assert any(idx > join_idx for idx in store_indices)
+    assert any(idx < join_idx for idx in load_indices)
+    assert any(idx > join_idx for idx in load_indices)
+
+
+def test_dup_op_operand_not_in_stack() -> None:
+    compiler = VenomCompiler(IRContext())
+    stack = StackModel()
+    assembly: list = []
+
+    ops = [IRVariable(f"%{i}") for i in range(5)]
+    for op in ops:
+        stack.push(op)
+
+    not_in_stack = IRVariable("%99")
+
+    with pytest.raises(AssertionError):
+        compiler.dup_op(assembly, stack, not_in_stack)
+
+
+def test_stack_reorder_operand_not_in_stack_but_spilled() -> None:
+    ctx = IRContext()
+    compiler = VenomCompiler(ctx)
+    compiler.dfg = _dummy_dfg()
+
+    stack = StackModel()
+    for i in range(5):
+        stack.push(IRVariable(f"%{i}"))
+
+    spilled_var = IRVariable("%spilled")
+    spilled: dict = {spilled_var: 0x10000}
+
+    assembly: list = []
+
+    # Try to reorder with spilled_var as target (should restore it from memory)
+    compiler._stack_reorder(assembly, stack, [spilled_var], spilled, dry_run=False)
+
+    # Should have restored the spilled variable
+    assert stack.get_depth(spilled_var) == 0  # Should be on top of stack
+    assert spilled_var not in spilled  # Should have been removed from spilled dict
+    # Assembly should contain PUSH and MLOAD to restore
+    assert "MLOAD" in assembly

vyper/utils.py

@@ -428,6 +428,7 @@ class MemoryPositions:
     FREE_VAR_SPACE = 0
     FREE_VAR_SPACE2 = 32
     RESERVED_MEMORY = 64
+    STACK_SPILL_BASE = 0x10000  # scratch space used for spilling deep stacks
 
 
 # Sizes of different data types. Used to clamp types.