bug: checks memory expansion costs utils (#1322)

<!--- Please provide a general summary of your changes in the title
above -->

<!-- Give an estimate of the time you spent on this PR in terms of work
days.
Did you spend 0.5 days on this PR or rather 2 days?  -->

Time spent on this PR:

## Pull request type

<!-- Please try to limit your pull request to one type,
submit multiple pull requests if needed. -->

Please check the type of change your PR introduces:

- [x] Bugfix
- [ ] Feature
- [ ] Code style update (formatting, renaming)
- [ ] Refactoring (no functional changes, no api changes)
- [ ] Build related changes
- [ ] Documentation content changes
- [ ] Other (please describe):

## What is the current behavior?

<!-- Please describe the current behavior that you are modifying,
or link to a relevant issue. -->

Resolves #1266

- `memory_expansion_cost_saturated`: offset.high + size.high != 0 can be
broken when selecting `(P - 1)/2` and `(P + 1)/2`
- `calculate_gas_extend_memory`: max_offset is a felt possibly > 2**128
so one needs to use `is_le_felt`

## What is the new behavior?

<!-- Please describe the behavior or changes that are being added by
this PR. -->

- Condition is updated for `memory_expansion_cost_saturated`
- Usage of `is_le_felt` in `calculate_gas_extend_memory`

<!-- Reviewable:start -->
- - -
This change is [<img src="https://reviewable.io/review_button.svg"
height="34" align="absmiddle"
alt="Reviewable"/>](https://reviewable.io/reviews/kkrt-labs/kakarot/1322)
<!-- Reviewable:end -->

src/kakarot/gas.cairo

@@ -73,15 +73,18 @@ namespace Gas {
         return memory_cost;
     }
 
-    // @notice Compute the expansion cost of max_offset for the memory
+    // @notice Compute the expansion cost of max_offset for the memory.
+    // @dev Assumption max_offset < 2**133 necessary for unsigned_div_rem usage.
     // @param words_len The current length of the memory.
     // @param max_offset The target max_offset to be applied to the given memory.
     // @return cost The expansion gas cost: 0 if no expansion is triggered, and the new size of the memory
     func calculate_gas_extend_memory{range_check_ptr}(
         words_len: felt, max_offset: felt
     ) -> model.MemoryExpansion {
         alloc_locals;
-        let memory_expansion = is_nn(max_offset - (words_len * 32 - 1));
+        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 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);
             return expansion;
@@ -111,14 +114,14 @@ namespace Gas {
             return expansion;
         }
 
-        if (offset.high + size.high != 0) {
-            // Hardcoded value of cost(2**128) and size of 2**128 bytes = 2**123 words of 32 bytes
-            // This offset would produce an OOG error in any case
-            let expansion = model.MemoryExpansion(cost=MEMORY_COST_U128, new_words_len=2 ** 123);
-            return expansion;
+        let is_low_part_overflowing = is_le_felt(2 ** 128, offset.low + size.low);
+        if (offset.high == 0 and size.high == 0 and is_low_part_overflowing == 0) {
+            return calculate_gas_extend_memory(words_len, offset.low + size.low);
         }
-
-        return calculate_gas_extend_memory(words_len, offset.low + size.low);
+        // Hardcoded value of cost(2**128) and size of 2**128 bytes = 2**123 words of 32 bytes
+        // This offset would produce an OOG error in any case
+        let expansion = model.MemoryExpansion(cost=MEMORY_COST_U128, new_words_len=2 ** 123);
+        return expansion;
     }
 
     // @notice Given two memory chunks, compute the maximum expansion cost

tests/src/kakarot/test_gas.cairo

@@ -1,4 +1,5 @@
 %builtins range_check
+from starkware.cairo.common.alloc import alloc
 
 from kakarot.gas import Gas
 from starkware.cairo.common.uint256 import Uint256
@@ -51,6 +52,24 @@ func test__max_memory_expansion_cost{range_check_ptr}() -> felt {
     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 tests.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;

tests/src/kakarot/test_gas.py

@@ -15,23 +15,21 @@ def test_should_return_same_as_execution_specs(self, cairo_run, max_offset):
             assert calculate_memory_gas_cost(max_offset) == output
 
         @given(
-            bytes_len=integers(min_value=0, max_value=0xFFFFFF),
-            added_offset=integers(min_value=0, max_value=0xFFFFFF),
+            bytes_len=integers(min_value=0, max_value=2**128 - 1),
+            added_offset=integers(min_value=0, max_value=2**128 - 1),
         )
         def test_should_return_correct_expansion_cost(
             self, cairo_run, bytes_len, added_offset
         ):
-            cost_before = calculate_memory_gas_cost(bytes_len)
-            cost_after = calculate_memory_gas_cost(bytes_len + added_offset)
-            diff = cost_after - cost_before
-
-            words_len = (bytes_len + 31) // 32
             max_offset = bytes_len + added_offset
             output = cairo_run(
                 "test__memory_expansion_cost",
-                words_len=words_len,
+                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(
@@ -62,6 +60,26 @@ def test_should_return_max_expansion_cost(
                 ).cost
             )
 
+        @given(
+            offset=integers(min_value=0, max_value=2**256 - 1),
+            size=integers(min_value=0, max_value=2**256 - 1),
+        )
+        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,
+            )
+            if size == 0:
+                cost = 0
+            elif offset + size > 2**128 - 1:
+                cost = 0x200000000000000000000000000018000000000000000000000000000000
+            else:
+                cost = calculate_gas_extend_memory(b"", [(offset, size)]).cost
+
+            assert cost == output
+
     class TestMessageGas:
         @pytest.mark.parametrize(
             "gas_param, gas_left, expected",