1use alloy_primitives::hex;
2use clap::Parser;
3use comfy_table::{modifiers::UTF8_ROUND_CORNERS, Table};
4use eyre::Result;
5use foundry_cli::{
6 opts::{BuildOpts, CompilerOpts, ProjectPathOpts},
7 utils::{cache_local_signatures, FoundryPathExt},
8};
9use foundry_common::{
10 compile::{compile_target, PathOrContractInfo, ProjectCompiler},
11 selectors::{import_selectors, SelectorImportData},
12};
13use foundry_compilers::{artifacts::output_selection::ContractOutputSelection, info::ContractInfo};
14use std::fs::canonicalize;
15
16#[derive(Clone, Debug, Parser)]
18pub enum SelectorsSubcommands {
19 #[command(visible_alias = "co")]
21 Collision {
22 first_contract: ContractInfo,
25
26 second_contract: ContractInfo,
29
30 #[command(flatten)]
31 build: Box<BuildOpts>,
32 },
33
34 #[command(visible_alias = "up")]
36 Upload {
37 #[arg(required_unless_present = "all")]
40 contract: Option<PathOrContractInfo>,
41
42 #[arg(long, required_unless_present = "contract")]
44 all: bool,
45
46 #[command(flatten)]
47 project_paths: ProjectPathOpts,
48 },
49
50 #[command(visible_alias = "ls")]
52 List {
53 #[arg(help = "The name of the contract to list selectors for.")]
55 contract: Option<String>,
56
57 #[command(flatten)]
58 project_paths: ProjectPathOpts,
59 },
60
61 #[command(visible_alias = "f")]
63 Find {
64 #[arg(help = "The selector to search for (with or without 0x prefix)")]
66 selector: String,
67
68 #[command(flatten)]
69 project_paths: ProjectPathOpts,
70 },
71
72 #[command(visible_alias = "c")]
74 Cache {
75 #[command(flatten)]
76 project_paths: ProjectPathOpts,
77 },
78}
79
80impl SelectorsSubcommands {
81 pub async fn run(self) -> Result<()> {
82 match self {
83 Self::Cache { project_paths } => {
84 sh_println!("Caching selectors for contracts in the project...")?;
85 let build_args = BuildOpts {
86 project_paths,
87 compiler: CompilerOpts {
88 extra_output: vec![ContractOutputSelection::Abi],
89 ..Default::default()
90 },
91 ..Default::default()
92 };
93
94 let project = build_args.project()?;
96 let outcome = ProjectCompiler::new().quiet(true).compile(&project)?;
97 cache_local_signatures(&outcome)?;
98 }
99 Self::Upload { contract, all, project_paths } => {
100 let build_args = BuildOpts {
101 project_paths: project_paths.clone(),
102 compiler: CompilerOpts {
103 extra_output: vec![ContractOutputSelection::Abi],
104 ..Default::default()
105 },
106 ..Default::default()
107 };
108
109 let project = build_args.project()?;
110 let output = if let Some(contract_info) = &contract {
111 let Some(contract_name) = contract_info.name() else {
112 eyre::bail!("No contract name provided.")
113 };
114
115 let target_path = contract_info
116 .path()
117 .map(Ok)
118 .unwrap_or_else(|| project.find_contract_path(contract_name))?;
119 compile_target(&target_path, &project, false)?
120 } else {
121 ProjectCompiler::new().compile(&project)?
122 };
123 let artifacts = if all {
124 output
125 .into_artifacts_with_files()
126 .filter(|(file, _, _)| {
127 let is_sources_path = file.starts_with(&project.paths.sources);
128 let is_test = file.is_sol_test();
129
130 is_sources_path && !is_test
131 })
132 .map(|(_, contract, artifact)| (contract, artifact))
133 .collect()
134 } else {
135 let contract_info = contract.unwrap();
136 let contract = contract_info.name().unwrap().to_string();
137
138 let found_artifact = if let Some(path) = contract_info.path() {
139 output.find(project.root().join(path).as_path(), &contract)
140 } else {
141 output.find_first(&contract)
142 };
143
144 let artifact = found_artifact
145 .ok_or_else(|| {
146 eyre::eyre!(
147 "Could not find artifact `{contract}` in the compiled artifacts"
148 )
149 })?
150 .clone();
151 vec![(contract, artifact)]
152 };
153
154 let mut artifacts = artifacts.into_iter().peekable();
155 while let Some((contract, artifact)) = artifacts.next() {
156 let abi = artifact.abi.ok_or_else(|| eyre::eyre!("Unable to fetch abi"))?;
157 if abi.functions.is_empty() && abi.events.is_empty() && abi.errors.is_empty() {
158 continue
159 }
160
161 sh_println!("Uploading selectors for {contract}...")?;
162
163 import_selectors(SelectorImportData::Abi(vec![abi])).await?.describe();
165
166 if artifacts.peek().is_some() {
167 sh_println!()?
168 }
169 }
170 }
171 Self::Collision { mut first_contract, mut second_contract, build } => {
172 let project = build.project()?;
174 let mut compiler = ProjectCompiler::new().quiet(true);
175
176 if let Some(contract_path) = &mut first_contract.path {
177 let target_path = canonicalize(&*contract_path)?;
178 *contract_path = target_path.to_string_lossy().to_string();
179 compiler = compiler.files([target_path]);
180 }
181 if let Some(contract_path) = &mut second_contract.path {
182 let target_path = canonicalize(&*contract_path)?;
183 *contract_path = target_path.to_string_lossy().to_string();
184 compiler = compiler.files([target_path]);
185 }
186
187 let output = compiler.compile(&project)?;
188
189 let methods = |contract: &ContractInfo| -> eyre::Result<_> {
191 let artifact = output
192 .find_contract(contract)
193 .ok_or_else(|| eyre::eyre!("Could not find artifact for {contract}"))?;
194 artifact.method_identifiers.as_ref().ok_or_else(|| {
195 eyre::eyre!("Could not find method identifiers for {contract}")
196 })
197 };
198 let first_method_map = methods(&first_contract)?;
199 let second_method_map = methods(&second_contract)?;
200
201 let colliding_methods: Vec<(&String, &String, &String)> = first_method_map
202 .iter()
203 .filter_map(|(k1, v1)| {
204 second_method_map
205 .iter()
206 .find_map(|(k2, v2)| if **v2 == *v1 { Some((k2, v2)) } else { None })
207 .map(|(k2, v2)| (v2, k1, k2))
208 })
209 .collect();
210
211 if colliding_methods.is_empty() {
212 sh_println!("No colliding method selectors between the two contracts.")?;
213 } else {
214 let mut table = Table::new();
215 table.apply_modifier(UTF8_ROUND_CORNERS);
216 table.set_header([
217 String::from("Selector"),
218 first_contract.name,
219 second_contract.name,
220 ]);
221 for method in &colliding_methods {
222 table.add_row([method.0, method.1, method.2]);
223 }
224 sh_println!("{} collisions found:", colliding_methods.len())?;
225 sh_println!("\n{table}\n")?;
226 }
227 }
228 Self::List { contract, project_paths } => {
229 sh_println!("Listing selectors for contracts in the project...")?;
230 let build_args = BuildOpts {
231 project_paths,
232 compiler: CompilerOpts {
233 extra_output: vec![ContractOutputSelection::Abi],
234 ..Default::default()
235 },
236 ..Default::default()
237 };
238
239 let project = build_args.project()?;
241 let outcome = ProjectCompiler::new().quiet(true).compile(&project)?;
242 let artifacts = if let Some(contract) = contract {
243 let found_artifact = outcome.find_first(&contract);
244 let artifact = found_artifact
245 .ok_or_else(|| {
246 let candidates = outcome
247 .artifacts()
248 .map(|(name, _,)| name)
249 .collect::<Vec<_>>();
250 let suggestion = if let Some(suggestion) = foundry_cli::utils::did_you_mean(&contract, candidates).pop() {
251 format!("\nDid you mean `{suggestion}`?")
252 } else {
253 String::new()
254 };
255 eyre::eyre!(
256 "Could not find artifact `{contract}` in the compiled artifacts{suggestion}",
257 )
258 })?
259 .clone();
260 vec![(contract, artifact)]
261 } else {
262 outcome
263 .into_artifacts_with_files()
264 .filter(|(file, _, _)| {
265 let is_sources_path = file.starts_with(&project.paths.sources);
266 let is_test = file.is_sol_test();
267
268 is_sources_path && !is_test
269 })
270 .map(|(_, contract, artifact)| (contract, artifact))
271 .collect()
272 };
273
274 let mut artifacts = artifacts.into_iter().peekable();
275
276 while let Some((contract, artifact)) = artifacts.next() {
277 let abi = artifact.abi.ok_or_else(|| eyre::eyre!("Unable to fetch abi"))?;
278 if abi.functions.is_empty() && abi.events.is_empty() && abi.errors.is_empty() {
279 continue
280 }
281
282 sh_println!("{contract}")?;
283
284 let mut table = Table::new();
285 table.apply_modifier(UTF8_ROUND_CORNERS);
286
287 table.set_header(["Type", "Signature", "Selector"]);
288
289 for func in abi.functions() {
290 let sig = func.signature();
291 let selector = func.selector();
292 table.add_row(["Function", &sig, &hex::encode_prefixed(selector)]);
293 }
294
295 for event in abi.events() {
296 let sig = event.signature();
297 let selector = event.selector();
298 table.add_row(["Event", &sig, &hex::encode_prefixed(selector)]);
299 }
300
301 for error in abi.errors() {
302 let sig = error.signature();
303 let selector = error.selector();
304 table.add_row(["Error", &sig, &hex::encode_prefixed(selector)]);
305 }
306
307 sh_println!("\n{table}\n")?;
308
309 if artifacts.peek().is_some() {
310 sh_println!()?
311 }
312 }
313 }
314
315 Self::Find { selector, project_paths } => {
316 sh_println!("Searching for selector {selector:?} in the project...")?;
317
318 let build_args = BuildOpts {
319 project_paths,
320 compiler: CompilerOpts {
321 extra_output: vec![ContractOutputSelection::Abi],
322 ..Default::default()
323 },
324 ..Default::default()
325 };
326
327 let project = build_args.project()?;
328 let outcome = ProjectCompiler::new().quiet(true).compile(&project)?;
329 let artifacts = outcome
330 .into_artifacts_with_files()
331 .filter(|(file, _, _)| {
332 let is_sources_path = file.starts_with(&project.paths.sources);
333 let is_test = file.is_sol_test();
334 is_sources_path && !is_test
335 })
336 .collect::<Vec<_>>();
337
338 let mut table = Table::new();
339 table.apply_modifier(UTF8_ROUND_CORNERS);
340
341 table.set_header(["Type", "Signature", "Selector", "Contract"]);
342
343 for (_file, contract, artifact) in artifacts {
344 let abi = artifact.abi.ok_or_else(|| eyre::eyre!("Unable to fetch abi"))?;
345
346 let selector_bytes =
347 hex::decode(selector.strip_prefix("0x").unwrap_or(&selector))?;
348
349 for func in abi.functions() {
350 if func.selector().as_slice().starts_with(selector_bytes.as_slice()) {
351 table.add_row([
352 "Function",
353 &func.signature(),
354 &hex::encode_prefixed(func.selector()),
355 contract.as_str(),
356 ]);
357 }
358 }
359
360 for event in abi.events() {
361 if event.selector().as_slice().starts_with(selector_bytes.as_slice()) {
362 table.add_row([
363 "Event",
364 &event.signature(),
365 &hex::encode_prefixed(event.selector()),
366 contract.as_str(),
367 ]);
368 }
369 }
370
371 for error in abi.errors() {
372 if error.selector().as_slice().starts_with(selector_bytes.as_slice()) {
373 table.add_row([
374 "Error",
375 &error.signature(),
376 &hex::encode_prefixed(error.selector()),
377 contract.as_str(),
378 ]);
379 }
380 }
381 }
382
383 if table.row_count() > 0 {
384 sh_println!("\nFound {} instance(s)...", table.row_count())?;
385 sh_println!("\n{table}\n")?;
386 } else {
387 return Err(eyre::eyre!("\nSelector not found in the project."));
388 }
389 }
390 }
391 Ok(())
392 }
393}