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
10pub enum ContextPrecompile<DB: Database> {
12 Ordinary(Precompile),
14 ContextStateful(ContextStatefulPrecompileArc<DB>),
17 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 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
47pub 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 #[inline]
65 pub fn new(spec_id: PrecompileSpecId) -> Self {
66 Self::from_static_precompiles(Precompiles::new(spec_id))
67 }
68
69 #[inline]
74 pub fn from_static_precompiles(precompiles: &'static Precompiles) -> Self {
75 Self {
76 inner: PrecompilesCow::StaticRef(precompiles),
77 }
78 }
79
80 #[inline]
82 pub fn from_precompiles(precompiles: HashMap<Address, ContextPrecompile<DB>>) -> Self {
83 Self {
84 inner: PrecompilesCow::Owned(precompiles),
85 }
86 }
87
88 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 #[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 #[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 #[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 #[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 #[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
196pub 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
207pub 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
220pub type ContextStatefulPrecompileArc<DB> = Arc<dyn ContextStatefulPrecompile<DB>>;
222
223pub 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}