perf(interpreter): rewrite gas accounting for memory expansion

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,16 @@ impl Gas {
         Self {
             limit,
             remaining: limit,
-            remaining_nomem: limit,
-            memory: 0,
+            refunded: 0,
+        }
+    }
+
+    /// Creates a new `Gas` struct with the given gas limit, but without any gas remaining.
+    #[inline]
+    pub const fn new_spent(limit: u64) -> Self {
+        Self {
+            limit,
+            remaining: 0,
             refunded: 0,
         }
     }
@@ -43,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.
@@ -75,10 +82,15 @@ 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;
+    }
+
     /// Records a refund value.
     ///
     /// `refund` can be negative but `self.refunded` should always be positive
@@ -109,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
     }
 }

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.

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::MemoryLimitOOG;
                 return $ret;
             }
-            $interp.shared_memory.resize(rounded_size);
         }
     };
 }

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,21 @@ impl InterpreterResult {
     }
 }
 
+/// Resize the memory to the new size. Returns whether the gas was enough to resize the memory.
+#[inline]
+#[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::*;

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) {
@@ -312,37 +313,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]

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::{

crates/revm/src/builder.rs

@@ -448,7 +448,7 @@ mod test {
         },
         Context, ContextPrecompile, ContextStatefulPrecompile, Evm, InMemoryDB, InnerEvmContext,
     };
-    use revm_interpreter::{Host, Interpreter};
+    use revm_interpreter::{gas, Host, Interpreter};
     use std::{cell::RefCell, rc::Rc, sync::Arc};
 
     /// Custom evm context
@@ -513,7 +513,7 @@ mod test {
         const EXPECTED_RESULT_GAS: u64 = INITIAL_TX_GAS + CUSTOM_INSTRUCTION_COST;
         fn custom_instruction(interp: &mut Interpreter, _host: &mut impl Host) {
             // just spend some gas
-            interp.gas.record_cost(CUSTOM_INSTRUCTION_COST);
+            gas!(interp, CUSTOM_INSTRUCTION_COST);
         }
 
         let code = Bytecode::new_raw([0xEF, 0x00].into());

crates/revm/src/handler/mainnet/execution.rs

@@ -24,8 +24,7 @@ pub fn frame_return_with_refund_flag<SPEC: Spec>(
     let refunded = gas.refunded();
 
     // Spend the gas limit. Gas is reimbursed when the tx returns successfully.
-    *gas = Gas::new(env.tx.gas_limit);
-    gas.record_cost(env.tx.gas_limit);
+    *gas = Gas::new_spent(env.tx.gas_limit);
 
     match instruction_result {
         return_ok!() => {

crates/revm/src/inspector/gas.rs

@@ -60,10 +60,7 @@ impl<DB: Database> Inspector<DB> for GasInspector {
         mut outcome: CallOutcome,
     ) -> CallOutcome {
         if outcome.result.result.is_error() {
-            outcome
-                .result
-                .gas
-                .record_cost(outcome.result.gas.remaining());
+            outcome.result.gas.spend_all();
             self.gas_remaining = 0;
         }
         outcome
@@ -76,10 +73,7 @@ impl<DB: Database> Inspector<DB> for GasInspector {
         mut outcome: CreateOutcome,
     ) -> CreateOutcome {
         if outcome.result.result.is_error() {
-            outcome
-                .result
-                .gas
-                .record_cost(outcome.result.gas.remaining());
+            outcome.result.gas.spend_all();
             self.gas_remaining = 0;
         }
         outcome

crates/revm/src/optimism/handler_register.rs

@@ -82,8 +82,7 @@ pub fn last_frame_return<SPEC: Spec, EXT, DB: Database>(
     let remaining = gas.remaining();
     let refunded = gas.refunded();
     // Spend the gas limit. Gas is reimbursed when the tx returns successfully.
-    *gas = Gas::new(tx_gas_limit);
-    gas.record_cost(tx_gas_limit);
+    *gas = Gas::new_spent(tx_gas_limit);
 
     match instruction_result {
         return_ok!() => {