perf(interpreter): rewrite gas accounting for memory expansion (bluealloy#1361)

* feat(interpreter): add helpers for spending all gas

* perf(interpreter): rewrite gas accounting for memory expansion

* chore: do not inline

* chore: restore Cargo.toml

Author: DaniPopes

Diff for: bins/revm-test/src/bin/snailtracer.rs

@@ -19,7 +19,7 @@ pub fn simple_example() {
         })
         .build();
 
-    let _ = evm.transact();
+    let _ = evm.transact().unwrap();
 }
 
 fn main() {

Diff for: crates/interpreter/src/gas.rs

@@ -14,10 +14,6 @@ pub struct Gas {
     limit: u64,
     /// The remaining gas.
     remaining: u64,
-    /// The remaining gas, without memory expansion.
-    remaining_nomem: u64,
-    /// The **last** memory expansion cost.
-    memory: u64,
     /// Refunded gas. This is used only at the end of execution.
     refunded: i64,
 }
@@ -29,8 +25,6 @@ impl Gas {
         Self {
             limit,
             remaining: limit,
-            remaining_nomem: limit,
-            memory: 0,
             refunded: 0,
         }
     }
@@ -41,8 +35,6 @@ impl Gas {
         Self {
             limit,
             remaining: 0,
-            remaining_nomem: 0,
-            memory: 0,
             refunded: 0,
         }
     }
@@ -55,8 +47,11 @@ impl Gas {
 
     /// Returns the **last** memory expansion cost.
     #[inline]
+    #[deprecated = "memory expansion cost is not tracked anymore; \
+                    calculate it using `SharedMemory::current_expansion_cost` instead"]
+    #[doc(hidden)]
     pub const fn memory(&self) -> u64 {
-        self.memory
+        0
     }
 
     /// Returns the total amount of gas that was refunded.
@@ -87,15 +82,13 @@ impl Gas {
     /// Erases a gas cost from the totals.
     #[inline]
     pub fn erase_cost(&mut self, returned: u64) {
-        self.remaining_nomem += returned;
         self.remaining += returned;
     }
 
     /// Spends all remaining gas.
     #[inline]
     pub fn spend_all(&mut self) {
         self.remaining = 0;
-        self.remaining_nomem = 0;
     }
 
     /// Records a refund value.
@@ -128,30 +121,13 @@ impl Gas {
     ///
     /// Returns `false` if the gas limit is exceeded.
     #[inline]
+    #[must_use]
     pub fn record_cost(&mut self, cost: u64) -> bool {
         let (remaining, overflow) = self.remaining.overflowing_sub(cost);
-        if overflow {
-            return false;
-        }
-
-        self.remaining_nomem -= cost;
-        self.remaining = remaining;
-        true
-    }
-
-    /// Records memory expansion gas.
-    ///
-    /// Used in [`resize_memory!`](crate::resize_memory).
-    #[inline]
-    pub fn record_memory(&mut self, gas_memory: u64) -> bool {
-        if gas_memory > self.memory {
-            let (remaining, overflow) = self.remaining_nomem.overflowing_sub(gas_memory);
-            if overflow {
-                return false;
-            }
-            self.memory = gas_memory;
+        let success = !overflow;
+        if success {
             self.remaining = remaining;
         }
-        true
+        success
     }
 }

Diff for: crates/interpreter/src/gas/calc.rs

@@ -342,13 +342,18 @@ pub const fn warm_cold_cost(is_cold: bool) -> u64 {
     }
 }
 
-/// Memory expansion cost calculation.
+/// Memory expansion cost calculation for a given memory length.
 #[inline]
-pub const fn memory_gas(a: usize) -> u64 {
-    let a = a as u64;
+pub const fn memory_gas_for_len(len: usize) -> u64 {
+    memory_gas(crate::interpreter::num_words(len as u64))
+}
+
+/// Memory expansion cost calculation for a given number of words.
+#[inline]
+pub const fn memory_gas(num_words: u64) -> u64 {
     MEMORY
-        .saturating_mul(a)
-        .saturating_add(a.saturating_mul(a) / 512)
+        .saturating_mul(num_words)
+        .saturating_add(num_words.saturating_mul(num_words) / 512)
 }
 
 /// Initial gas that is deducted for transaction to be included.

Diff for: crates/interpreter/src/instructions/macros.rs

@@ -89,27 +89,23 @@ macro_rules! resize_memory {
         $crate::resize_memory!($interp, $offset, $len, ())
     };
     ($interp:expr, $offset:expr, $len:expr, $ret:expr) => {
-        let size = $offset.saturating_add($len);
-        if size > $interp.shared_memory.len() {
-            // We are fine with saturating to usize if size is close to MAX value.
-            let rounded_size = $crate::interpreter::next_multiple_of_32(size);
-
+        let new_size = $offset.saturating_add($len);
+        if new_size > $interp.shared_memory.len() {
             #[cfg(feature = "memory_limit")]
-            if $interp.shared_memory.limit_reached(size) {
+            if $interp.shared_memory.limit_reached(new_size) {
                 $interp.instruction_result = $crate::InstructionResult::MemoryLimitOOG;
                 return $ret;
             }
 
-            // Gas is calculated in evm words (256 bits).
-            let words_num = rounded_size / 32;
-            if !$interp
-                .gas
-                .record_memory($crate::gas::memory_gas(words_num))
-            {
+            // Note: we can't use `Interpreter` directly here because of potential double-borrows.
+            if !$crate::interpreter::resize_memory(
+                &mut $interp.shared_memory,
+                &mut $interp.gas,
+                new_size,
+            ) {
                 $interp.instruction_result = $crate::InstructionResult::MemoryOOG;
                 return $ret;
             }
-            $interp.shared_memory.resize(rounded_size);
         }
     };
 }

Diff for: crates/interpreter/src/interpreter.rs

@@ -6,12 +6,12 @@ mod shared_memory;
 mod stack;
 
 pub use contract::Contract;
-pub use shared_memory::{next_multiple_of_32, SharedMemory, EMPTY_SHARED_MEMORY};
+pub use shared_memory::{num_words, SharedMemory, EMPTY_SHARED_MEMORY};
 pub use stack::{Stack, STACK_LIMIT};
 
 use crate::EOFCreateOutcome;
 use crate::{
-    primitives::Bytes, push, push_b256, return_ok, return_revert, CallOutcome, CreateOutcome,
+    gas, primitives::Bytes, push, push_b256, return_ok, return_revert, CallOutcome, CreateOutcome,
     FunctionStack, Gas, Host, InstructionResult, InterpreterAction,
 };
 use core::cmp::min;
@@ -379,6 +379,13 @@ impl Interpreter {
             },
         }
     }
+
+    /// Resize the memory to the new size. Returns whether the gas was enough to resize the memory.
+    #[inline]
+    #[must_use]
+    pub fn resize_memory(&mut self, new_size: usize) -> bool {
+        resize_memory(&mut self.shared_memory, &mut self.gas, new_size)
+    }
 }
 
 impl InterpreterResult {
@@ -401,6 +408,22 @@ impl InterpreterResult {
     }
 }
 
+/// Resize the memory to the new size. Returns whether the gas was enough to resize the memory.
+#[inline(never)]
+#[cold]
+#[must_use]
+pub fn resize_memory(memory: &mut SharedMemory, gas: &mut Gas, new_size: usize) -> bool {
+    let new_words = num_words(new_size as u64);
+    let new_cost = gas::memory_gas(new_words);
+    let current_cost = memory.current_expansion_cost();
+    let cost = new_cost - current_cost;
+    let success = gas.record_cost(cost);
+    if success {
+        memory.resize((new_words as usize) * 32);
+    }
+    success
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

Diff for: crates/interpreter/src/interpreter/shared_memory.rs

@@ -1,10 +1,5 @@
+use core::{cmp::min, fmt, ops::Range};
 use revm_primitives::{B256, U256};
-
-use core::{
-    cmp::min,
-    fmt,
-    ops::{BitAnd, Not, Range},
-};
 use std::vec::Vec;
 
 /// A sequential memory shared between calls, which uses
@@ -128,6 +123,12 @@ impl SharedMemory {
         self.len() == 0
     }
 
+    /// Returns the gas cost for the current memory expansion.
+    #[inline]
+    pub fn current_expansion_cost(&self) -> u64 {
+        crate::gas::memory_gas_for_len(self.len())
+    }
+
     /// Resizes the memory in-place so that `len` is equal to `new_len`.
     #[inline]
     pub fn resize(&mut self, new_size: usize) {
@@ -145,21 +146,18 @@ impl SharedMemory {
         self.slice_range(offset..offset + size)
     }
 
+    /// Returns a byte slice of the memory region at the given offset.
+    ///
+    /// # Panics
+    ///
+    /// Panics on out of bounds.
     #[inline]
     #[cfg_attr(debug_assertions, track_caller)]
-    pub fn slice_range(&self, range: Range<usize>) -> &[u8] {
-        let last_checkpoint = self.last_checkpoint;
-
-        self.buffer
-            .get(last_checkpoint + range.start..last_checkpoint + range.end)
-            .unwrap_or_else(|| {
-                debug_unreachable!(
-                    "slice OOB: {}..{}; len: {}",
-                    range.start,
-                    range.end,
-                    self.len()
-                )
-            })
+    pub fn slice_range(&self, range @ Range { start, end }: Range<usize>) -> &[u8] {
+        match self.context_memory().get(range) {
+            Some(slice) => slice,
+            None => debug_unreachable!("slice OOB: {start}..{end}; len: {}", self.len()),
+        }
     }
 
     /// Returns a byte slice of the memory region at the given offset.
@@ -170,13 +168,11 @@ impl SharedMemory {
     #[inline]
     #[cfg_attr(debug_assertions, track_caller)]
     pub fn slice_mut(&mut self, offset: usize, size: usize) -> &mut [u8] {
-        let len = self.len();
         let end = offset + size;
-        let last_checkpoint = self.last_checkpoint;
-
-        self.buffer
-            .get_mut(last_checkpoint + offset..last_checkpoint + offset + size)
-            .unwrap_or_else(|| debug_unreachable!("slice OOB: {offset}..{end}; len: {}", len))
+        match self.context_memory_mut().get_mut(offset..end) {
+            Some(slice) => slice,
+            None => debug_unreachable!("slice OOB: {offset}..{end}"),
+        }
     }
 
     /// Returns the byte at the given offset.
@@ -312,37 +308,28 @@ impl SharedMemory {
     }
 }
 
-/// Rounds up `x` to the closest multiple of 32. If `x % 32 == 0` then `x` is returned. Note, if `x`
-/// is greater than `usize::MAX - 31` this will return `usize::MAX` which isn't a multiple of 32.
+/// Returns number of words what would fit to provided number of bytes,
+/// i.e. it rounds up the number bytes to number of words.
 #[inline]
-pub fn next_multiple_of_32(x: usize) -> usize {
-    let r = x.bitand(31).not().wrapping_add(1).bitand(31);
-    x.saturating_add(r)
+pub const fn num_words(len: u64) -> u64 {
+    len.saturating_add(31) / 32
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
-    fn test_next_multiple_of_32() {
-        // next_multiple_of_32 returns x when it is a multiple of 32
-        for i in 0..32 {
-            let x = i * 32;
-            assert_eq!(x, next_multiple_of_32(x));
-        }
-
-        // next_multiple_of_32 rounds up to the nearest multiple of 32 when `x % 32 != 0`
-        for x in 0..1024 {
-            if x % 32 == 0 {
-                continue;
-            }
-            let next_multiple = x + 32 - (x % 32);
-            assert_eq!(next_multiple, next_multiple_of_32(x));
-        }
-
-        // We expect large values to saturate and not overflow.
-        assert_eq!(usize::MAX, next_multiple_of_32(usize::MAX));
+    fn test_num_words() {
+        assert_eq!(num_words(0), 0);
+        assert_eq!(num_words(1), 1);
+        assert_eq!(num_words(31), 1);
+        assert_eq!(num_words(32), 1);
+        assert_eq!(num_words(33), 2);
+        assert_eq!(num_words(63), 2);
+        assert_eq!(num_words(64), 2);
+        assert_eq!(num_words(65), 3);
+        assert_eq!(num_words(u64::MAX), u64::MAX / 32);
     }
 
     #[test]

Diff for: crates/interpreter/src/lib.rs

@@ -34,7 +34,7 @@ pub use gas::Gas;
 pub use host::{DummyHost, Host, LoadAccountResult, SStoreResult, SelfDestructResult};
 pub use instruction_result::*;
 pub use interpreter::{
-    analysis, next_multiple_of_32, Contract, Interpreter, InterpreterResult, SharedMemory, Stack,
+    analysis, num_words, Contract, Interpreter, InterpreterResult, SharedMemory, Stack,
     EMPTY_SHARED_MEMORY, STACK_LIMIT,
 };
 pub use interpreter_action::{