fix: memory expansion early return condition

cairo/src/gas.cairo

@@ -92,7 +92,7 @@ namespace Gas {
     ) -> model.MemoryExpansion {
         alloc_locals;
         let is_memory_length_not_zero = is_not_zero(words_len);
-        let current_memory_length = (words_len * 32 - 1) * is_memory_length_not_zero;
+        let current_memory_length = (words_len * 32) * is_memory_length_not_zero;
         let memory_expansion = is_le_felt(current_memory_length, max_offset);
         if (memory_expansion == FALSE) {
             let expansion = model.MemoryExpansion(cost=0, new_words_len=words_len);

cairo/tests/src/test_gas.cairo

@@ -0,0 +1,84 @@
+%builtins range_check
+from starkware.cairo.common.alloc import alloc
+
+from src.gas import Gas
+from starkware.cairo.common.uint256 import Uint256
+from starkware.cairo.lang.compiler.lib.registers import get_fp_and_pc
+
+func test__memory_cost{range_check_ptr}() -> felt {
+    tempvar words_len: felt;
+    %{ ids.words_len = program_input["words_len"]; %}
+    let cost = Gas.memory_cost(words_len);
+
+    return cost;
+}
+
+func test__memory_expansion_cost{range_check_ptr}() -> felt {
+    tempvar words_len: felt;
+    tempvar max_offset: felt;
+    %{
+        ids.words_len = program_input["words_len"];
+        ids.max_offset = program_input["max_offset"];
+    %}
+    let memory_expansion = Gas.calculate_gas_extend_memory(words_len, max_offset);
+
+    return memory_expansion.cost;
+}
+
+func test__max_memory_expansion_cost{range_check_ptr}() -> felt {
+    alloc_locals;
+    let fp_and_pc = get_fp_and_pc();
+    local __fp__: felt* = fp_and_pc.fp_val;
+    local words_len: felt;
+    local offset_1: Uint256;
+    local size_1: Uint256;
+    local offset_2: Uint256;
+    local size_2: Uint256;
+    %{
+        ids.words_len = program_input["words_len"];
+        ids.offset_1.low = program_input["offset_1"][0]
+        ids.offset_1.high = program_input["offset_1"][1]
+        ids.size_1.low = program_input["size_1"][0]
+        ids.size_1.high = program_input["size_1"][1]
+        ids.offset_2.low = program_input["offset_2"][0]
+        ids.offset_2.high = program_input["offset_2"][1]
+        ids.size_2.low = program_input["size_2"][0]
+        ids.size_2.high = program_input["size_2"][1]
+    %}
+    let memory_expansion = Gas.max_memory_expansion_cost(
+        words_len, &offset_1, &size_1, &offset_2, &size_2
+    );
+
+    return memory_expansion.cost;
+}
+
+func test__memory_expansion_cost_saturated{range_check_ptr}() -> felt {
+    alloc_locals;
+    local words_len: felt;
+    let (offset) = alloc();
+    let (size) = alloc();
+    %{
+        from src.utils.uint256 import int_to_uint256
+        ids.words_len = program_input["words_len"]
+        segments.write_arg(ids.offset, int_to_uint256(program_input["offset"]))
+        segments.write_arg(ids.size, int_to_uint256(program_input["size"]))
+    %}
+
+    let memory_expansion = Gas.memory_expansion_cost_saturated(
+        words_len, [cast(offset, Uint256*)], [cast(size, Uint256*)]
+    );
+    return memory_expansion.cost;
+}
+
+func test__compute_message_call_gas{range_check_ptr}() -> felt {
+    tempvar gas_param: Uint256;
+    tempvar gas_left: felt;
+    %{
+        ids.gas_param.low = program_input["gas_param"][0]
+        ids.gas_param.high = program_input["gas_param"][1]
+        ids.gas_left = program_input["gas_left"]
+    %}
+    let gas = Gas.compute_message_call_gas(gas_param, gas_left);
+
+    return gas;
+}

cairo/tests/src/test_gas.py

@@ -0,0 +1,121 @@
+import pytest
+from hypothesis import given
+from hypothesis.strategies import integers
+
+from ethereum.cancun.vm.gas import (
+    calculate_gas_extend_memory,
+    calculate_memory_gas_cost,
+)
+from src.utils.uint256 import int_to_uint256
+from tests.utils.strategies import felt, uint128, uint256
+
+
+class TestGas:
+    class TestCost:
+        @given(max_offset=integers(min_value=0, max_value=0xFFFFFF))
+        def test_should_return_same_as_execution_specs(self, cairo_run, max_offset):
+            output = cairo_run("test__memory_cost", words_len=(max_offset + 31) // 32)
+            assert calculate_memory_gas_cost(max_offset) == output
+
+        @given(
+            bytes_len=uint128,
+            added_offset=uint128,
+        )
+        def test_should_return_correct_expansion_cost(
+            self, cairo_run, bytes_len, added_offset
+        ):
+            max_offset = bytes_len + added_offset
+            output = cairo_run(
+                "test__memory_expansion_cost",
+                words_len=(bytes_len + 31) // 32,
+                max_offset=max_offset,
+            )
+            cost_before = calculate_memory_gas_cost(bytes_len)
+            cost_after = calculate_memory_gas_cost(max_offset)
+            diff = cost_after - cost_before
+            assert diff == output
+
+        @given(
+            offset_1=felt,
+            size_1=felt,
+            offset_2=felt,
+            size_2=felt,
+            words_len=integers(
+                min_value=0, max_value=0x3C0000
+            ),  # upper bound reaching 30M gas limit on expansion
+        )
+        def test_should_return_max_expansion_cost(
+            self, cairo_run, offset_1, size_1, offset_2, size_2, words_len
+        ):
+            memory_cost_u32 = calculate_memory_gas_cost(2**32 - 1)
+            output = cairo_run(
+                "test__max_memory_expansion_cost",
+                words_len=words_len,
+                offset_1=int_to_uint256(offset_1),
+                size_1=int_to_uint256(size_1),
+                offset_2=int_to_uint256(offset_2),
+                size_2=int_to_uint256(size_2),
+            )
+            expansion = calculate_gas_extend_memory(
+                b"\x00" * 32 * words_len,
+                [
+                    (offset_1, size_1),
+                    (offset_2, size_2),
+                ],
+            )
+
+            # If the memory expansion is greater than 2**27 words of 32 bytes
+            # We saturate it to the hardcoded value corresponding the the gas cost of a 2**32 memory size
+            expected_saturated = (
+                memory_cost_u32
+                if (words_len * 32 + expansion.expand_by) >= 2**32
+                else expansion.cost
+            )
+            assert output == expected_saturated
+
+        @given(
+            offset=uint256,
+            size=uint256,
+        )
+        def test_memory_expansion_cost_saturated(self, cairo_run, offset, size):
+            output = cairo_run(
+                "test__memory_expansion_cost_saturated",
+                words_len=0,
+                offset=offset,
+                size=size,
+            )
+
+            total_expansion = offset.wrapping_add(size)
+
+            if size == 0:
+                cost = 0
+            elif (
+                total_expansion > 2**32
+                or total_expansion < offset
+                or total_expansion < size
+            ):
+                cost = calculate_memory_gas_cost(2**32)
+            else:
+                cost = calculate_gas_extend_memory(b"", [(offset, size)]).cost
+
+            assert cost == output
+
+    class TestMessageGas:
+        @pytest.mark.parametrize(
+            "gas_param, gas_left, expected",
+            [
+                (0, 0, 0),
+                (10, 100, 10),
+                (100, 100, 99),
+                (100, 10, 10),
+            ],
+        )
+        def test_should_return_message_base_gas(
+            self, cairo_run, gas_param, gas_left, expected
+        ):
+            output = cairo_run(
+                "test__compute_message_call_gas",
+                gas_param=int_to_uint256(gas_param),
+                gas_left=gas_left,
+            )
+            assert output == expected