fix: end a section at SSTORE too (EIP-1706) (#21)

* fix: end a section at SSTORE too (EIP-1706)

* comment

crates/revm-jit-builtins/src/gas.rs

@@ -1,3 +1,124 @@
 //! Gas calculation utilities.
 
+use revm_primitives::{SpecId, U256};
+
 pub use revm_interpreter::gas::*;
+
+/// `const` Option `?`.
+#[allow(unused_macros)]
+macro_rules! tri {
+    ($e:expr) => {
+        match $e {
+            Some(v) => v,
+            None => return None,
+        }
+    };
+}
+
+// These are overridden to only account for the dynamic cost.
+
+/// `EXP` opcode cost calculation.
+#[inline]
+pub fn dyn_exp_cost(spec_id: SpecId, power: U256) -> Option<u64> {
+    #[inline]
+    const fn log2floor(value: U256) -> u64 {
+        let mut l: u64 = 256;
+        let mut i = 3;
+        loop {
+            if value.as_limbs()[i] == 0u64 {
+                l -= 64;
+            } else {
+                l -= value.as_limbs()[i].leading_zeros() as u64;
+                if l == 0 {
+                    return l;
+                } else {
+                    return l - 1;
+                }
+            }
+            if i == 0 {
+                break;
+            }
+            i -= 1;
+        }
+        l
+    }
+
+    if power == U256::ZERO {
+        Some(0)
+    } else {
+        // EIP-160: EXP cost increase
+        let gas_byte =
+            U256::from(if spec_id.is_enabled_in(SpecId::SPURIOUS_DRAGON) { 50 } else { 10 });
+        let gas = gas_byte.checked_mul(U256::from(log2floor(power) / 8 + 1))?;
+        u64::try_from(gas).ok()
+    }
+}
+
+/// `LOG` opcode cost calculation.
+#[inline]
+pub const fn dyn_log_cost(len: u64) -> Option<u64> {
+    LOGDATA.checked_mul(len)
+}
+
+/// `KECCAK256` opcode cost calculation.
+#[inline]
+pub const fn dyn_keccak256_cost(len: u64) -> Option<u64> {
+    cost_per_word(len, KECCAK256WORD)
+}
+
+/// `*COPY` opcodes cost calculation.
+#[inline]
+pub const fn dyn_verylowcopy_cost(len: u64) -> Option<u64> {
+    cost_per_word(len, COPY)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn exp_cost() {
+        for (spec_id, power) in [
+            (SpecId::CANCUN, U256::from(0)),
+            (SpecId::CANCUN, U256::from(1)),
+            (SpecId::CANCUN, U256::from(69)),
+        ] {
+            assert_eq!(
+                super::exp_cost(spec_id, power).unwrap(),
+                EXP + dyn_exp_cost(spec_id, power).unwrap(),
+            );
+        }
+    }
+
+    #[test]
+    fn log_cost() {
+        for n_topics in [0, 1, 2] {
+            for len in [0, 1, 69] {
+                assert_eq!(
+                    super::log_cost(n_topics, len).unwrap(),
+                    LOG + (LOGTOPIC * n_topics as u64) + dyn_log_cost(len).unwrap(),
+                );
+            }
+        }
+    }
+
+    #[test]
+    fn keccak256_cost() {
+        for len in [0, 1, 69] {
+            assert_eq!(
+                super::keccak256_cost(len).unwrap(),
+                KECCAK256 + dyn_keccak256_cost(len).unwrap(),
+            );
+        }
+    }
+
+    #[test]
+    fn verylowcopy_cost() {
+        for len in [0, 1, 69] {
+            assert_eq!(
+                super::verylowcopy_cost(len).unwrap(),
+                VERYLOW + dyn_verylowcopy_cost(len).unwrap(),
+            );
+        }
+    }
+}

crates/revm-jit-builtins/src/lib.rs

@@ -91,7 +91,7 @@ pub unsafe extern "C" fn __revm_jit_builtin_exp(
     spec_id: SpecId,
 ) -> InstructionResult {
     let exponent = exponent_ptr.to_u256();
-    gas_opt!(ecx, gas::exp_cost(spec_id, exponent).map(|g| g - gas::EXP));
+    gas_opt!(ecx, gas::dyn_exp_cost(spec_id, exponent));
     *exponent_ptr = base.to_u256().pow(exponent).into();
     InstructionResult::Continue
 }
@@ -105,7 +105,7 @@ pub unsafe extern "C" fn __revm_jit_builtin_keccak256(
     *len_ptr = EvmWord::from_be_bytes(if len == 0 {
         KECCAK_EMPTY.0
     } else {
-        gas_opt!(ecx, gas::keccak256_cost(len as u64).map(|g| g - gas::KECCAK256));
+        gas_opt!(ecx, gas::dyn_keccak256_cost(len as u64));
         let offset = try_into_usize!(offset);
         resize_memory!(ecx, offset, len);
         let data = ecx.memory.slice(offset, len);
@@ -209,7 +209,7 @@ pub unsafe extern "C" fn __revm_jit_builtin_returndatacopy(
     rev![memory_offset, offset, len]: &mut [EvmWord; 3],
 ) -> InstructionResult {
     let len = try_into_usize!(len);
-    gas_opt!(ecx, gas::verylowcopy_cost(len as u64).map(|g| g - gas::VERYLOW));
+    gas_opt!(ecx, gas::dyn_verylowcopy_cost(len as u64));
     let data_offset = offset.to_u256();
     let data_offset = as_usize_saturated!(data_offset);
     let (data_end, overflow) = data_offset.overflowing_add(len);
@@ -388,7 +388,7 @@ pub unsafe extern "C" fn __revm_jit_builtin_mcopy(
     rev![dst, src, len]: &mut [EvmWord; 3],
 ) -> InstructionResult {
     let len = try_into_usize!(len);
-    gas_opt!(ecx, gas::verylowcopy_cost(len as u64).map(|g| g - gas::VERYLOW));
+    gas_opt!(ecx, gas::dyn_verylowcopy_cost(len as u64));
     if len != 0 {
         let dst = try_into_usize!(dst);
         let src = try_into_usize!(src);
@@ -408,7 +408,7 @@ pub unsafe extern "C" fn __revm_jit_builtin_log(
     let sp = sp.add(n as usize);
     read_words!(sp, offset, len);
     let len = try_into_usize!(len);
-    gas_opt!(ecx, gas::LOGDATA.checked_mul(len as u64));
+    gas_opt!(ecx, gas::dyn_log_cost(len as u64));
     let data = if len != 0 {
         let offset = try_into_usize!(offset);
         resize_memory!(ecx, offset, len);

crates/revm-jit-builtins/src/utils.rs

@@ -39,7 +39,7 @@ pub(crate) unsafe fn copy_operation(
     data: &[u8],
 ) -> InstructionResult {
     let len = try_into_usize!(len);
-    gas_opt!(ecx, gas::verylowcopy_cost(len as u64).map(|g| g - gas::VERYLOW));
+    gas_opt!(ecx, gas::dyn_verylowcopy_cost(len as u64));
     if len == 0 {
         return InstructionResult::Continue;
     }

crates/revm-jit/src/bytecode/mod.rs

@@ -155,6 +155,7 @@ impl<'a> Bytecode<'a> {
     /// Returns a mutable reference the instruction at the given instruction counter.
     #[inline]
     #[track_caller]
+    #[allow(dead_code)]
     pub(crate) fn inst_mut(&mut self, inst: Inst) -> &mut InstData {
         &mut self.insts[inst]
     }
@@ -167,6 +168,15 @@ impl<'a> Bytecode<'a> {
         self.iter_all_insts().filter(|(_, data)| !data.is_dead_code())
     }
 
+    /// Returns an iterator over the instructions.
+    #[inline]
+    #[allow(dead_code)]
+    pub(crate) fn iter_mut_insts(
+        &mut self,
+    ) -> impl DoubleEndedIterator<Item = (usize, &mut InstData)> + '_ {
+        self.iter_mut_all_insts().filter(|(_, data)| !data.is_dead_code())
+    }
+
     /// Returns an iterator over all the instructions, including dead code.
     #[inline]
     pub(crate) fn iter_all_insts(
@@ -175,6 +185,15 @@ impl<'a> Bytecode<'a> {
         self.insts.iter().enumerate()
     }
 
+    /// Returns an iterator over all the instructions, including dead code.
+    #[inline]
+    #[allow(dead_code)]
+    pub(crate) fn iter_mut_all_insts(
+        &mut self,
+    ) -> impl DoubleEndedIterator<Item = (usize, &mut InstData)> + ExactSizeIterator + '_ {
+        self.insts.iter_mut().enumerate()
+    }
+
     /// Runs a list of analysis passes on the instructions.
     #[instrument(level = "debug", skip_all)]
     pub(crate) fn analyze(&mut self) -> Result<()> {
@@ -305,6 +324,18 @@ impl<'a> Bytecode<'a> {
         analysis.finish(self);
     }
 
+    /// Constructs the sections in the bytecode.
+    #[instrument(name = "sections", level = "debug", skip_all)]
+    #[cfg(any())]
+    fn construct_sections_default(&mut self) {
+        for inst in &mut self.insts {
+            let (inp, out) = inst.stack_io();
+            let stack_diff = out as i16 - inp as i16;
+            inst.section =
+                Section { gas_cost: inst.base_gas as _, inputs: inp as _, max_growth: stack_diff }
+        }
+    }
+
     /// Returns the immediate value of the given instruction data, if any.
     pub(crate) fn get_imm_of(&self, instr_data: &InstData) -> Option<&'a [u8]> {
         (instr_data.imm_len() > 0).then(|| self.get_imm(instr_data.data))
@@ -482,6 +513,15 @@ impl InstData {
         self.flags.contains(InstFlags::DEAD_CODE)
     }
 
+    /// Returns `true` if this instruction requires to know `gasleft()`.
+    /// Note that this does not include CALL and CREATE.
+    #[inline]
+    pub(crate) fn requires_gasleft(&self, spec_id: SpecId) -> bool {
+        // For SSTORE, see `revm_interpreter::gas::sstore_cost`.
+        self.opcode == op::GAS
+            || (self.opcode == op::SSTORE && spec_id.is_enabled_in(SpecId::ISTANBUL))
+    }
+
     /// Returns `true` if we know that this instruction will branch or stop execution.
     #[inline]
     pub(crate) const fn is_branching(&self, is_eof: bool) -> bool {

crates/revm-jit/src/bytecode/sections.rs

@@ -1,6 +1,9 @@
 use super::Bytecode;
 use core::fmt;
-use revm_interpreter::opcode as op;
+
+// TODO: Separate gas sections from stack length sections.
+// E.g. `GAS` should stop only a gas section because it requires `gasleft`, and execution will
+// continue with the next instruction.
 
 /// A section is a sequence of instructions that are executed sequentially without any jumps or
 /// branches.
@@ -68,8 +71,12 @@ impl SectionAnalysis {
 
         self.gas_cost += data.base_gas as u64;
 
-        // Branching instructions end a section, starting a new one on the next instruction, if any.
-        if data.opcode == op::GAS || data.is_branching(bytecode.is_eof()) || data.will_suspend() {
+        // Instructions that require `gasleft` and branching instructions end a section, starting a
+        // new one on the next instruction, if any.
+        if data.requires_gasleft(bytecode.spec_id)
+            || data.is_branching(bytecode.is_eof())
+            || data.will_suspend()
+        {
             let next = inst + 1;
             self.save_to(bytecode, next);
             self.reset(next);
@@ -109,7 +116,10 @@ impl SectionAnalysis {
                 ?section,
                 "saving"
             );
-            bytecode.inst_mut(self.start_inst).section = section;
+            let mut insts = bytecode.insts[self.start_inst..].iter_mut();
+            if let Some(inst) = insts.find(|inst| !inst.is_dead_code()) {
+                inst.section = section;
+            }
         }
     }