foundry_evm_traces/identifier/
signatures.rs1use alloy_json_abi::{Error, Event, Function, JsonAbi};
2use alloy_primitives::{B256, Selector, map::HashMap};
3use eyre::Result;
4use foundry_common::{
5 abi::{get_error, get_event, get_func},
6 fs,
7 selectors::{OpenChainClient, SelectorKind},
8};
9use foundry_config::Config;
10use serde::{Deserialize, Serialize};
11use std::{
12 collections::BTreeMap,
13 path::{Path, PathBuf},
14 sync::Arc,
15};
16use tokio::sync::RwLock;
17
18#[derive(Debug, Default, Deserialize)]
20#[serde(try_from = "SignaturesDiskCache")]
21pub struct SignaturesCache {
22 signatures: HashMap<SelectorKind, Option<String>>,
23}
24
25#[derive(Serialize, Deserialize)]
27struct SignaturesDiskCache {
28 functions: BTreeMap<Selector, String>,
29 errors: BTreeMap<Selector, String>,
30 events: BTreeMap<B256, String>,
31}
32
33impl From<SignaturesDiskCache> for SignaturesCache {
34 fn from(value: SignaturesDiskCache) -> Self {
35 let functions = value
36 .functions
37 .into_iter()
38 .map(|(selector, signature)| (SelectorKind::Function(selector), signature));
39 let errors = value
40 .errors
41 .into_iter()
42 .map(|(selector, signature)| (SelectorKind::Error(selector), signature));
43 let events = value
44 .events
45 .into_iter()
46 .map(|(selector, signature)| (SelectorKind::Event(selector), signature));
47 Self {
48 signatures: functions
49 .chain(errors)
50 .chain(events)
51 .map(|(sel, sig)| (sel, (!sig.is_empty()).then_some(sig)))
52 .collect(),
53 }
54 }
55}
56
57impl From<&SignaturesCache> for SignaturesDiskCache {
58 fn from(value: &SignaturesCache) -> Self {
59 let (functions, errors, events) = value.signatures.iter().fold(
60 (BTreeMap::new(), BTreeMap::new(), BTreeMap::new()),
61 |mut acc, (kind, signature)| {
62 let value = signature.clone().unwrap_or_default();
63 match *kind {
64 SelectorKind::Function(selector) => _ = acc.0.insert(selector, value),
65 SelectorKind::Error(selector) => _ = acc.1.insert(selector, value),
66 SelectorKind::Event(selector) => _ = acc.2.insert(selector, value),
67 }
68 acc
69 },
70 );
71 Self { functions, errors, events }
72 }
73}
74
75impl Serialize for SignaturesCache {
76 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
77 where
78 S: serde::Serializer,
79 {
80 SignaturesDiskCache::from(self).serialize(serializer)
81 }
82}
83
84impl SignaturesCache {
85 #[instrument(target = "evm::traces", name = "SignaturesCache::load")]
87 pub fn load(path: &Path) -> Self {
88 trace!(target: "evm::traces", ?path, "reading signature cache");
89 fs::read_json_file(path)
90 .inspect_err(
91 |err| warn!(target: "evm::traces", ?path, ?err, "failed to read cache file"),
92 )
93 .unwrap_or_default()
94 }
95
96 #[instrument(target = "evm::traces", name = "SignaturesCache::save", skip(self))]
98 pub fn save(&self, path: &Path) {
99 if let Some(parent) = path.parent()
100 && let Err(err) = std::fs::create_dir_all(parent)
101 {
102 warn!(target: "evm::traces", ?parent, %err, "failed to create cache");
103 }
104 if let Err(err) = fs::write_json_file(path, self) {
105 warn!(target: "evm::traces", %err, "failed to flush signature cache");
106 } else {
107 trace!(target: "evm::traces", "flushed signature cache")
108 }
109 }
110
111 pub fn extend_from_abi(&mut self, abi: &JsonAbi) {
113 self.extend(abi.items().filter_map(|item| match item {
114 alloy_json_abi::AbiItem::Function(f) => {
115 Some((SelectorKind::Function(f.selector()), f.signature()))
116 }
117 alloy_json_abi::AbiItem::Error(e) => {
118 Some((SelectorKind::Error(e.selector()), e.signature()))
119 }
120 alloy_json_abi::AbiItem::Event(e) => {
121 Some((SelectorKind::Event(e.selector()), e.full_signature()))
122 }
123 _ => None,
124 }));
125 }
126
127 pub fn insert(&mut self, key: SelectorKind, value: String) {
129 self.extend(std::iter::once((key, value)));
130 }
131
132 pub fn extend(&mut self, signatures: impl IntoIterator<Item = (SelectorKind, String)>) {
134 self.signatures
135 .extend(signatures.into_iter().map(|(k, v)| (k, (!v.is_empty()).then_some(v))));
136 }
137
138 pub fn get(&self, key: &SelectorKind) -> Option<Option<String>> {
140 self.signatures.get(key).cloned()
141 }
142
143 pub fn contains_key(&self, key: &SelectorKind) -> bool {
145 self.signatures.contains_key(key)
146 }
147}
148
149#[derive(Clone, Debug)]
152pub struct SignaturesIdentifier(Arc<SignaturesIdentifierInner>);
153
154#[derive(Debug)]
155struct SignaturesIdentifierInner {
156 cache: RwLock<SignaturesCache>,
158 cache_path: Option<PathBuf>,
160 client: Option<OpenChainClient>,
162}
163
164impl SignaturesIdentifier {
165 pub fn new(offline: bool) -> Result<Self> {
167 Self::new_with(Config::foundry_cache_dir().as_deref(), offline)
168 }
169
170 pub fn from_config(config: &Config) -> Result<Self> {
172 Self::new(config.offline)
173 }
174
175 pub fn new_with(cache_dir: Option<&Path>, offline: bool) -> Result<Self> {
180 let client = if !offline { Some(OpenChainClient::new()?) } else { None };
181 let (cache, cache_path) = if let Some(cache_dir) = cache_dir {
182 let path = cache_dir.join("signatures");
183 let cache = SignaturesCache::load(&path);
184 (cache, Some(path))
185 } else {
186 Default::default()
187 };
188 Ok(Self(Arc::new(SignaturesIdentifierInner {
189 cache: RwLock::new(cache),
190 cache_path,
191 client,
192 })))
193 }
194
195 pub fn save(&self) {
197 self.0.save();
198 }
199
200 pub async fn identify_functions(
202 &self,
203 identifiers: impl IntoIterator<Item = Selector>,
204 ) -> Vec<Option<Function>> {
205 self.identify_map(identifiers.into_iter().map(SelectorKind::Function), get_func).await
206 }
207
208 pub async fn identify_function(&self, identifier: Selector) -> Option<Function> {
210 self.identify_functions([identifier]).await.pop().unwrap()
211 }
212
213 pub async fn identify_events(
215 &self,
216 identifiers: impl IntoIterator<Item = B256>,
217 ) -> Vec<Option<Event>> {
218 self.identify_map(identifiers.into_iter().map(SelectorKind::Event), get_event).await
219 }
220
221 pub async fn identify_event(&self, identifier: B256) -> Option<Event> {
223 self.identify_events([identifier]).await.pop().unwrap()
224 }
225
226 pub async fn identify_errors(
228 &self,
229 identifiers: impl IntoIterator<Item = Selector>,
230 ) -> Vec<Option<Error>> {
231 self.identify_map(identifiers.into_iter().map(SelectorKind::Error), get_error).await
232 }
233
234 pub async fn identify_error(&self, identifier: Selector) -> Option<Error> {
236 self.identify_errors([identifier]).await.pop().unwrap()
237 }
238
239 pub async fn identify(&self, selectors: &[SelectorKind]) -> Vec<Option<String>> {
241 if selectors.is_empty() {
242 return vec![];
243 }
244 trace!(target: "evm::traces", ?selectors, "identifying selectors");
245
246 let mut cache_r = self.0.cache.read().await;
247 if let Some(client) = &self.0.client {
248 let query =
249 selectors.iter().copied().filter(|v| !cache_r.contains_key(v)).collect::<Vec<_>>();
250 if !query.is_empty() {
251 drop(cache_r);
252 let mut cache_w = self.0.cache.write().await;
253 if let Ok(res) = client.decode_selectors(&query).await {
254 for (selector, signatures) in std::iter::zip(query, res) {
255 cache_w.signatures.insert(selector, signatures.into_iter().next());
256 }
257 }
258 drop(cache_w);
259 cache_r = self.0.cache.read().await;
260 }
261 }
262 selectors.iter().map(|selector| cache_r.get(selector).unwrap_or_default()).collect()
263 }
264
265 async fn identify_map<T>(
266 &self,
267 selectors: impl IntoIterator<Item = SelectorKind>,
268 get_type: impl Fn(&str) -> Result<T>,
269 ) -> Vec<Option<T>> {
270 let results = self.identify(&Vec::from_iter(selectors)).await;
271 results.into_iter().map(|r| r.and_then(|r| get_type(&r).ok())).collect()
272 }
273}
274
275impl SignaturesIdentifierInner {
276 fn save(&self) {
277 if let Some(path) = &self.cache_path {
278 self.cache
279 .try_read()
280 .expect("SignaturesIdentifier cache is locked while attempting to save")
281 .save(path);
282 }
283 }
284}
285
286impl Drop for SignaturesIdentifierInner {
287 fn drop(&mut self) {
288 self.save();
289 }
290}