perf: avoid cloning precompiles (#1486)

* perf: avoid cloning precompiles twice

* perf: cow it up

* pedantic changes, unsafe removed, nit renames

* into string

* nits

* Make to_mut call cold fn on owned mutation

* doc update

crates/precompile/src/bls12_381/utils.rs

@@ -1,8 +1,7 @@
-use core::cmp::Ordering;
-
 use blst::{
     blst_bendian_from_fp, blst_fp, blst_fp_from_bendian, blst_scalar, blst_scalar_from_bendian,
 };
+use core::cmp::Ordering;
 use revm_primitives::PrecompileError;
 
 /// Number of bits used in the BLS12-381 curve finite field elements.

crates/precompile/src/lib.rs

@@ -28,9 +28,9 @@ use once_cell::race::OnceBox;
 pub use revm_primitives as primitives;
 pub use revm_primitives::{
     precompile::{PrecompileError as Error, *},
-    Address, Bytes, HashMap, Log, B256,
+    Address, Bytes, HashMap, HashSet, Log, B256,
 };
-use std::boxed::Box;
+use std::{boxed::Box, vec::Vec};
 
 pub fn calc_linear_cost_u32(len: usize, base: u64, word: u64) -> u64 {
     (len as u64 + 32 - 1) / 32 * word + base
@@ -39,7 +39,9 @@ pub fn calc_linear_cost_u32(len: usize, base: u64, word: u64) -> u64 {
 #[derive(Clone, Default, Debug)]
 pub struct Precompiles {
     /// Precompiles.
-    pub inner: HashMap<Address, Precompile>,
+    inner: HashMap<Address, Precompile>,
+    /// Addresses of precompile.
+    addresses: HashSet<Address>,
 }
 
 impl Precompiles {
@@ -71,6 +73,11 @@ impl Precompiles {
         })
     }
 
+    /// Returns inner HashMap of precompiles.
+    pub fn inner(&self) -> &HashMap<Address, Precompile> {
+        &self.inner
+    }
+
     /// Returns precompiles for Byzantium spec.
     pub fn byzantium() -> &'static Self {
         static INSTANCE: OnceBox<Precompiles> = OnceBox::new();
@@ -168,13 +175,13 @@ impl Precompiles {
 
     /// Returns an iterator over the precompiles addresses.
     #[inline]
-    pub fn addresses(&self) -> impl Iterator<Item = &Address> {
+    pub fn addresses(&self) -> impl ExactSizeIterator<Item = &Address> {
         self.inner.keys()
     }
 
     /// Consumes the type and returns all precompile addresses.
     #[inline]
-    pub fn into_addresses(self) -> impl Iterator<Item = Address> {
+    pub fn into_addresses(self) -> impl ExactSizeIterator<Item = Address> {
         self.inner.into_keys()
     }
 
@@ -206,11 +213,19 @@ impl Precompiles {
         self.inner.len()
     }
 
+    /// Returns the precompiles addresses as a set.
+    pub fn addresses_set(&self) -> &HashSet<Address> {
+        &self.addresses
+    }
+
     /// Extends the precompiles with the given precompiles.
     ///
     /// Other precompiles with overwrite existing precompiles.
+    #[inline]
     pub fn extend(&mut self, other: impl IntoIterator<Item = PrecompileWithAddress>) {
-        self.inner.extend(other.into_iter().map(Into::into));
+        let items = other.into_iter().collect::<Vec<_>>();
+        self.addresses.extend(items.iter().map(|p| *p.address()));
+        self.inner.extend(items.into_iter().map(Into::into));
     }
 }
 
@@ -229,6 +244,20 @@ impl From<PrecompileWithAddress> for (Address, Precompile) {
     }
 }
 
+impl PrecompileWithAddress {
+    /// Returns reference of address.
+    #[inline]
+    pub fn address(&self) -> &Address {
+        &self.0
+    }
+
+    /// Returns reference of precompile.
+    #[inline]
+    pub fn precompile(&self) -> &Precompile {
+        &self.1
+    }
+}
+
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Ord, PartialOrd)]
 pub enum PrecompileSpecId {
     HOMESTEAD,

crates/primitives/src/precompile.rs

@@ -110,11 +110,25 @@ impl Precompile {
 
     /// Call the precompile with the given input and gas limit and return the result.
     pub fn call(&mut self, bytes: &Bytes, gas_price: u64, env: &Env) -> PrecompileResult {
-        match self {
+        match *self {
+            Precompile::Standard(p) => p(bytes, gas_price),
+            Precompile::Env(p) => p(bytes, gas_price, env),
+            Precompile::Stateful(ref p) => p.call(bytes, gas_price, env),
+            Precompile::StatefulMut(ref mut p) => p.call_mut(bytes, gas_price, env),
+        }
+    }
+
+    /// Call the precompile with the given input and gas limit and return the result.
+    ///
+    /// Returns an error if the precompile is mutable.
+    pub fn call_ref(&self, bytes: &Bytes, gas_price: u64, env: &Env) -> PrecompileResult {
+        match *self {
             Precompile::Standard(p) => p(bytes, gas_price),
             Precompile::Env(p) => p(bytes, gas_price, env),
-            Precompile::Stateful(p) => p.call(bytes, gas_price, env),
-            Precompile::StatefulMut(p) => p.call_mut(bytes, gas_price, env),
+            Precompile::Stateful(ref p) => p.call(bytes, gas_price, env),
+            Precompile::StatefulMut(_) => Err(PrecompileErrors::Fatal {
+                msg: "call_ref on mutable stateful precompile".into(),
+            }),
         }
     }
 }

crates/revm/src/context/context_precompiles.rs

@@ -1,15 +1,13 @@
+use super::InnerEvmContext;
 use crate::{
     precompile::{Precompile, PrecompileResult},
-    primitives::{db::Database, Address, Bytes, HashMap},
+    primitives::{db::Database, Address, Bytes, HashMap, HashSet},
 };
-use core::ops::{Deref, DerefMut};
 use dyn_clone::DynClone;
-use revm_precompile::Precompiles;
+use revm_precompile::{PrecompileSpecId, PrecompileWithAddress, Precompiles};
 use std::{boxed::Box, sync::Arc};
 
-use super::InnerEvmContext;
-
-/// Precompile and its handlers.
+/// A single precompile handler.
 pub enum ContextPrecompile<DB: Database> {
     /// Ordinary precompiles
     Ordinary(Precompile),
@@ -24,53 +22,160 @@ pub enum ContextPrecompile<DB: Database> {
 impl<DB: Database> Clone for ContextPrecompile<DB> {
     fn clone(&self) -> Self {
         match self {
-            Self::Ordinary(arg0) => Self::Ordinary(arg0.clone()),
-            Self::ContextStateful(arg0) => Self::ContextStateful(arg0.clone()),
-            Self::ContextStatefulMut(arg0) => Self::ContextStatefulMut(arg0.clone()),
+            Self::Ordinary(p) => Self::Ordinary(p.clone()),
+            Self::ContextStateful(p) => Self::ContextStateful(p.clone()),
+            Self::ContextStatefulMut(p) => Self::ContextStatefulMut(p.clone()),
         }
     }
 }
 
-#[derive(Clone)]
+enum PrecompilesCow<DB: Database> {
+    /// Default precompiles, returned by `Precompiles::new`. Used to fast-path the default case.
+    StaticRef(&'static Precompiles),
+    Owned(HashMap<Address, ContextPrecompile<DB>>),
+}
+
+impl<DB: Database> Clone for PrecompilesCow<DB> {
+    fn clone(&self) -> Self {
+        match *self {
+            PrecompilesCow::StaticRef(p) => PrecompilesCow::StaticRef(p),
+            PrecompilesCow::Owned(ref inner) => PrecompilesCow::Owned(inner.clone()),
+        }
+    }
+}
+
+/// Precompiles context.
 pub struct ContextPrecompiles<DB: Database> {
-    inner: HashMap<Address, ContextPrecompile<DB>>,
+    inner: PrecompilesCow<DB>,
+}
+
+impl<DB: Database> Clone for ContextPrecompiles<DB> {
+    fn clone(&self) -> Self {
+        Self {
+            inner: self.inner.clone(),
+        }
+    }
 }
 
 impl<DB: Database> ContextPrecompiles<DB> {
-    /// Returns precompiles addresses.
+    /// Creates a new precompiles context at the given spec ID.
+    ///
+    /// This is a cheap operation that does not allocate by reusing the global precompiles.
     #[inline]
-    pub fn addresses(&self) -> impl Iterator<Item = &Address> {
-        self.inner.keys()
+    pub fn new(spec_id: PrecompileSpecId) -> Self {
+        Self::from_static_precompiles(Precompiles::new(spec_id))
     }
 
-    /// Extends the precompiles with the given precompiles.
+    /// Creates a new precompiles context from the given static precompiles.
     ///
-    /// Other precompiles with overwrite existing precompiles.
+    /// NOTE: The internal precompiles must not be `StatefulMut` or `call` will panic.
+    /// This is done because the default precompiles are not stateful.
     #[inline]
-    pub fn extend(
-        &mut self,
-        other: impl IntoIterator<Item = impl Into<(Address, ContextPrecompile<DB>)>>,
-    ) {
-        self.inner.extend(other.into_iter().map(Into::into));
+    pub fn from_static_precompiles(precompiles: &'static Precompiles) -> Self {
+        Self {
+            inner: PrecompilesCow::StaticRef(precompiles),
+        }
+    }
+
+    /// Creates a new precompiles context from the given precompiles.
+    #[inline]
+    pub fn from_precompiles(precompiles: HashMap<Address, ContextPrecompile<DB>>) -> Self {
+        Self {
+            inner: PrecompilesCow::Owned(precompiles),
+        }
+    }
+
+    /// Returns precompiles addresses as a HashSet.
+    pub fn addresses_set(&self) -> HashSet<Address> {
+        match self.inner {
+            PrecompilesCow::StaticRef(inner) => inner.addresses_set().clone(),
+            PrecompilesCow::Owned(ref inner) => inner.keys().cloned().collect(),
+        }
+    }
+
+    /// Returns precompiles addresses.
+    #[inline]
+    pub fn addresses<'a>(&'a self) -> Box<dyn ExactSizeIterator<Item = &Address> + 'a> {
+        match self.inner {
+            PrecompilesCow::StaticRef(inner) => Box::new(inner.addresses()),
+            PrecompilesCow::Owned(ref inner) => Box::new(inner.keys()),
+        }
+    }
+
+    /// Returns `true` if the precompiles contains the given address.
+    #[inline]
+    pub fn contains(&self, address: &Address) -> bool {
+        match self.inner {
+            PrecompilesCow::StaticRef(inner) => inner.contains(address),
+            PrecompilesCow::Owned(ref inner) => inner.contains_key(address),
+        }
     }
 
     /// Call precompile and executes it. Returns the result of the precompile execution.
-    /// None if the precompile does not exist.
+    ///
+    /// Returns `None` if the precompile does not exist.
     #[inline]
     pub fn call(
         &mut self,
-        addess: Address,
+        address: &Address,
         bytes: &Bytes,
         gas_price: u64,
         evmctx: &mut InnerEvmContext<DB>,
     ) -> Option<PrecompileResult> {
-        let precompile = self.inner.get_mut(&addess)?;
+        Some(match self.inner {
+            PrecompilesCow::StaticRef(p) => p.get(address)?.call_ref(bytes, gas_price, &evmctx.env),
+            PrecompilesCow::Owned(ref mut owned) => match owned.get_mut(address)? {
+                ContextPrecompile::Ordinary(p) => p.call(bytes, gas_price, &evmctx.env),
+                ContextPrecompile::ContextStateful(p) => p.call(bytes, gas_price, evmctx),
+                ContextPrecompile::ContextStatefulMut(p) => p.call_mut(bytes, gas_price, evmctx),
+            },
+        })
+    }
 
-        match precompile {
-            ContextPrecompile::Ordinary(p) => Some(p.call(bytes, gas_price, &evmctx.env)),
-            ContextPrecompile::ContextStatefulMut(p) => Some(p.call_mut(bytes, gas_price, evmctx)),
-            ContextPrecompile::ContextStateful(p) => Some(p.call(bytes, gas_price, evmctx)),
+    /// Returns a mutable reference to the precompiles map.
+    ///
+    /// Clones the precompiles map if it is shared.
+    #[inline]
+    pub fn to_mut(&mut self) -> &mut HashMap<Address, ContextPrecompile<DB>> {
+        if let PrecompilesCow::StaticRef(_) = self.inner {
+            self.mutate_into_owned();
         }
+
+        let PrecompilesCow::Owned(inner) = &mut self.inner else {
+            unreachable!("self is mutated to Owned.")
+        };
+        inner
+    }
+
+    /// Mutates Self into Owned variant, or do nothing if it is already Owned.
+    /// Mutation will clone all precompiles.
+    #[cold]
+    fn mutate_into_owned(&mut self) {
+        let PrecompilesCow::StaticRef(precompiles) = self.inner else {
+            return;
+        };
+        self.inner = PrecompilesCow::Owned(
+            precompiles
+                .inner()
+                .iter()
+                .map(|(k, v)| (*k, v.clone().into()))
+                .collect(),
+        );
+    }
+}
+
+impl<DB: Database> Extend<(Address, ContextPrecompile<DB>)> for ContextPrecompiles<DB> {
+    fn extend<T: IntoIterator<Item = (Address, ContextPrecompile<DB>)>>(&mut self, iter: T) {
+        self.to_mut().extend(iter.into_iter().map(Into::into))
+    }
+}
+
+impl<DB: Database> Extend<PrecompileWithAddress> for ContextPrecompiles<DB> {
+    fn extend<T: IntoIterator<Item = PrecompileWithAddress>>(&mut self, iter: T) {
+        self.to_mut().extend(iter.into_iter().map(|precompile| {
+            let (address, precompile) = precompile.into();
+            (address, precompile.into())
+        }));
     }
 }
 
@@ -82,17 +187,9 @@ impl<DB: Database> Default for ContextPrecompiles<DB> {
     }
 }
 
-impl<DB: Database> Deref for ContextPrecompiles<DB> {
-    type Target = HashMap<Address, ContextPrecompile<DB>>;
-
-    fn deref(&self) -> &Self::Target {
-        &self.inner
-    }
-}
-
-impl<DB: Database> DerefMut for ContextPrecompiles<DB> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.inner
+impl<DB: Database> Default for PrecompilesCow<DB> {
+    fn default() -> Self {
+        Self::Owned(Default::default())
     }
 }
 
@@ -132,22 +229,24 @@ impl<DB: Database> From<Precompile> for ContextPrecompile<DB> {
     }
 }
 
-impl<DB: Database> From<Precompiles> for ContextPrecompiles<DB> {
-    fn from(p: Precompiles) -> Self {
-        ContextPrecompiles {
-            inner: p.inner.into_iter().map(|(k, v)| (k, v.into())).collect(),
-        }
-    }
-}
-
-impl<DB: Database> From<&Precompiles> for ContextPrecompiles<DB> {
-    fn from(p: &Precompiles) -> Self {
-        ContextPrecompiles {
-            inner: p
-                .inner
-                .iter()
-                .map(|(&k, v)| (k, v.clone().into()))
-                .collect(),
-        }
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::db::EmptyDB;
+
+    #[test]
+    fn test_precompiles_context() {
+        let custom_address = Address::with_last_byte(0xff);
+
+        let mut precompiles = ContextPrecompiles::<EmptyDB>::new(PrecompileSpecId::HOMESTEAD);
+        assert_eq!(precompiles.addresses().count(), 4);
+        assert!(matches!(precompiles.inner, PrecompilesCow::StaticRef(_)));
+        assert!(!precompiles.contains(&custom_address));
+
+        let precompile = Precompile::Standard(|_, _| panic!());
+        precompiles.extend([(custom_address, precompile.into())]);
+        assert_eq!(precompiles.addresses().count(), 5);
+        assert!(matches!(precompiles.inner, PrecompilesCow::Owned(_)));
+        assert!(precompiles.contains(&custom_address));
     }
 }