revm/context/
context_precompiles.rs

1use super::InnerEvmContext;
2use crate::{
3    precompile::{Precompile, PrecompileResult},
4    primitives::{db::Database, Address, Bytes, HashMap, HashSet},
5};
6use dyn_clone::DynClone;
7use revm_precompile::{PrecompileSpecId, PrecompileWithAddress, Precompiles};
8use std::{boxed::Box, sync::Arc};
9
10/// A single precompile handler.
11pub enum ContextPrecompile<DB: Database> {
12    /// Ordinary precompiles
13    Ordinary(Precompile),
14    /// Stateful precompile that is Arc over [`ContextStatefulPrecompile`] trait.
15    /// It takes a reference to input, gas limit and Context.
16    ContextStateful(ContextStatefulPrecompileArc<DB>),
17    /// Mutable stateful precompile that is Box over [`ContextStatefulPrecompileMut`] trait.
18    /// It takes a reference to input, gas limit and context.
19    ContextStatefulMut(ContextStatefulPrecompileBox<DB>),
20}
21
22impl<DB: Database> Clone for ContextPrecompile<DB> {
23    fn clone(&self) -> Self {
24        match self {
25            Self::Ordinary(p) => Self::Ordinary(p.clone()),
26            Self::ContextStateful(p) => Self::ContextStateful(p.clone()),
27            Self::ContextStatefulMut(p) => Self::ContextStatefulMut(p.clone()),
28        }
29    }
30}
31
32enum PrecompilesCow<DB: Database> {
33    /// Default precompiles, returned by `Precompiles::new`. Used to fast-path the default case.
34    StaticRef(&'static Precompiles),
35    Owned(HashMap<Address, ContextPrecompile<DB>>),
36}
37
38impl<DB: Database> Clone for PrecompilesCow<DB> {
39    fn clone(&self) -> Self {
40        match *self {
41            PrecompilesCow::StaticRef(p) => PrecompilesCow::StaticRef(p),
42            PrecompilesCow::Owned(ref inner) => PrecompilesCow::Owned(inner.clone()),
43        }
44    }
45}
46
47/// Precompiles context.
48pub struct ContextPrecompiles<DB: Database> {
49    inner: PrecompilesCow<DB>,
50}
51
52impl<DB: Database> Clone for ContextPrecompiles<DB> {
53    fn clone(&self) -> Self {
54        Self {
55            inner: self.inner.clone(),
56        }
57    }
58}
59
60impl<DB: Database> ContextPrecompiles<DB> {
61    /// Creates a new precompiles context at the given spec ID.
62    ///
63    /// This is a cheap operation that does not allocate by reusing the global precompiles.
64    #[inline]
65    pub fn new(spec_id: PrecompileSpecId) -> Self {
66        Self::from_static_precompiles(Precompiles::new(spec_id))
67    }
68
69    /// Creates a new precompiles context from the given static precompiles.
70    ///
71    /// NOTE: The internal precompiles must not be `StatefulMut` or `call` will panic.
72    /// This is done because the default precompiles are not stateful.
73    #[inline]
74    pub fn from_static_precompiles(precompiles: &'static Precompiles) -> Self {
75        Self {
76            inner: PrecompilesCow::StaticRef(precompiles),
77        }
78    }
79
80    /// Creates a new precompiles context from the given precompiles.
81    #[inline]
82    pub fn from_precompiles(precompiles: HashMap<Address, ContextPrecompile<DB>>) -> Self {
83        Self {
84            inner: PrecompilesCow::Owned(precompiles),
85        }
86    }
87
88    /// Returns precompiles addresses as a HashSet.
89    pub fn addresses_set(&self) -> HashSet<Address> {
90        match self.inner {
91            PrecompilesCow::StaticRef(inner) => inner.addresses_set().clone(),
92            PrecompilesCow::Owned(ref inner) => inner.keys().cloned().collect(),
93        }
94    }
95
96    /// Returns precompiles addresses.
97    #[inline]
98    pub fn addresses<'a>(&'a self) -> Box<dyn ExactSizeIterator<Item = &'a Address> + 'a> {
99        match self.inner {
100            PrecompilesCow::StaticRef(inner) => Box::new(inner.addresses()),
101            PrecompilesCow::Owned(ref inner) => Box::new(inner.keys()),
102        }
103    }
104
105    /// Returns `true` if the precompiles contains the given address.
106    #[inline]
107    pub fn contains(&self, address: &Address) -> bool {
108        match self.inner {
109            PrecompilesCow::StaticRef(inner) => inner.contains(address),
110            PrecompilesCow::Owned(ref inner) => inner.contains_key(address),
111        }
112    }
113
114    /// Call precompile and executes it. Returns the result of the precompile execution.
115    ///
116    /// Returns `None` if the precompile does not exist.
117    #[inline]
118    pub fn call(
119        &mut self,
120        address: &Address,
121        bytes: &Bytes,
122        gas_limit: u64,
123        evmctx: &mut InnerEvmContext<DB>,
124    ) -> Option<PrecompileResult> {
125        Some(match self.inner {
126            PrecompilesCow::StaticRef(p) => p.get(address)?.call_ref(bytes, gas_limit, &evmctx.env),
127            PrecompilesCow::Owned(ref mut owned) => match owned.get_mut(address)? {
128                ContextPrecompile::Ordinary(p) => p.call(bytes, gas_limit, &evmctx.env),
129                ContextPrecompile::ContextStateful(p) => p.call(bytes, gas_limit, evmctx),
130                ContextPrecompile::ContextStatefulMut(p) => p.call_mut(bytes, gas_limit, evmctx),
131            },
132        })
133    }
134
135    /// Returns a mutable reference to the precompiles map.
136    ///
137    /// Clones the precompiles map if it is shared.
138    #[inline]
139    pub fn to_mut(&mut self) -> &mut HashMap<Address, ContextPrecompile<DB>> {
140        if let PrecompilesCow::StaticRef(_) = self.inner {
141            self.mutate_into_owned();
142        }
143
144        let PrecompilesCow::Owned(inner) = &mut self.inner else {
145            unreachable!("self is mutated to Owned.")
146        };
147        inner
148    }
149
150    /// Mutates Self into Owned variant, or do nothing if it is already Owned.
151    /// Mutation will clone all precompiles.
152    #[cold]
153    fn mutate_into_owned(&mut self) {
154        let PrecompilesCow::StaticRef(precompiles) = self.inner else {
155            return;
156        };
157        self.inner = PrecompilesCow::Owned(
158            precompiles
159                .inner()
160                .iter()
161                .map(|(k, v)| (*k, v.clone().into()))
162                .collect(),
163        );
164    }
165}
166
167impl<DB: Database> Extend<(Address, ContextPrecompile<DB>)> for ContextPrecompiles<DB> {
168    fn extend<T: IntoIterator<Item = (Address, ContextPrecompile<DB>)>>(&mut self, iter: T) {
169        self.to_mut().extend(iter.into_iter().map(Into::into))
170    }
171}
172
173impl<DB: Database> Extend<PrecompileWithAddress> for ContextPrecompiles<DB> {
174    fn extend<T: IntoIterator<Item = PrecompileWithAddress>>(&mut self, iter: T) {
175        self.to_mut().extend(iter.into_iter().map(|precompile| {
176            let (address, precompile) = precompile.into();
177            (address, precompile.into())
178        }));
179    }
180}
181
182impl<DB: Database> Default for ContextPrecompiles<DB> {
183    fn default() -> Self {
184        Self {
185            inner: Default::default(),
186        }
187    }
188}
189
190impl<DB: Database> Default for PrecompilesCow<DB> {
191    fn default() -> Self {
192        Self::Owned(Default::default())
193    }
194}
195
196/// Context aware stateful precompile trait. It is used to create
197/// a arc precompile in [`ContextPrecompile`].
198pub trait ContextStatefulPrecompile<DB: Database>: Sync + Send {
199    fn call(
200        &self,
201        bytes: &Bytes,
202        gas_limit: u64,
203        evmctx: &mut InnerEvmContext<DB>,
204    ) -> PrecompileResult;
205}
206
207/// Context aware mutable stateful precompile trait. It is used to create
208/// a boxed precompile in [`ContextPrecompile`].
209pub trait ContextStatefulPrecompileMut<DB: Database>: DynClone + Send + Sync {
210    fn call_mut(
211        &mut self,
212        bytes: &Bytes,
213        gas_limit: u64,
214        evmctx: &mut InnerEvmContext<DB>,
215    ) -> PrecompileResult;
216}
217
218dyn_clone::clone_trait_object!(<DB> ContextStatefulPrecompileMut<DB>);
219
220/// Arc over context stateful precompile.
221pub type ContextStatefulPrecompileArc<DB> = Arc<dyn ContextStatefulPrecompile<DB>>;
222
223/// Box over context mutable stateful precompile
224pub type ContextStatefulPrecompileBox<DB> = Box<dyn ContextStatefulPrecompileMut<DB>>;
225
226impl<DB: Database> From<Precompile> for ContextPrecompile<DB> {
227    fn from(p: Precompile) -> Self {
228        ContextPrecompile::Ordinary(p)
229    }
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235    use crate::db::EmptyDB;
236
237    #[test]
238    fn test_precompiles_context() {
239        let custom_address = Address::with_last_byte(0xff);
240
241        let mut precompiles = ContextPrecompiles::<EmptyDB>::new(PrecompileSpecId::HOMESTEAD);
242        assert_eq!(precompiles.addresses().count(), 4);
243        assert!(matches!(precompiles.inner, PrecompilesCow::StaticRef(_)));
244        assert!(!precompiles.contains(&custom_address));
245
246        let precompile = Precompile::Standard(|_, _| panic!());
247        precompiles.extend([(custom_address, precompile.into())]);
248        assert_eq!(precompiles.addresses().count(), 5);
249        assert!(matches!(precompiles.inner, PrecompilesCow::Owned(_)));
250        assert!(precompiles.contains(&custom_address));
251    }
252}