-
Notifications
You must be signed in to change notification settings - Fork 2.1k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
refactor(selector-generator): Polish the code (#2888)
## What ❔ - Changes procedural approach to more idiomatic for rust (introduces `Selectors` + `App` structures). `Selectors` are not FS aware, `App` is. - Introduces more strong typing -- when both selectors and function names are `String`, it's easy to misuse them (I did several times, while working on this PR). - Removes manual selector calculation and uses `ethabi` for that. - Makes the code async. - Adds docs for CLI (e.g. to see with `--help` argument). - Adds a test for `Selectors`. - Reduces number of panics and uses `anyhow` to propagate errors. - Adds new selectors (I guess I have more contracts compiled locally). ## Why ❔ Maintainability and testability, mostly. With the refactoring, it should be easier to add new functionality and tests for it. Also, the runtime has improved. Before ``` $ time cargo run --bin selector_generator -- contracts etc/selector-generator-data/selectors.json Compiling selector_generator v0.1.0 (/home/popzxc/workspace/current/zksync-era/core/bin/selector_generator) Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.54s Running `target/debug/selector_generator contracts etc/selector-generator-data/selectors.json` Analyzed 401 files. Added 452 selectors (before: 516 after: 968) cargo run --bin selector_generator -- contracts 4.48s user 7.02s system 100% cpu 11.486 total ``` After ``` $ time cargo run --bin selector_generator -- contracts etc/selector-generator-data/selectors.json Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.20s Running `target/debug/selector_generator contracts etc/selector-generator-data/selectors.json` Analyzed 1087 files. Added 0 selectors (before: 1023 after: 1023) cargo run --bin selector_generator -- contracts 1.24s user 0.12s system 100% cpu 1.360 total ``` ## Next steps I would love to see `--check` mode introduced (which checks that 0 new entries were added) and this tool is added to CI, so that we continously run it and prevent from rotting. @mm-zk leaving it to you.
- Loading branch information
Showing
6 changed files
with
1,240 additions
and
576 deletions.
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,105 @@ | ||
use std::path::PathBuf; | ||
|
||
use anyhow::Context; | ||
use glob::glob; | ||
use tokio::io::AsyncWriteExt as _; | ||
|
||
use crate::selectors::Selectors; | ||
|
||
#[derive(Debug, Default)] | ||
pub(crate) struct App { | ||
/// Selectors file. | ||
file_path: PathBuf, | ||
/// All the selectors. Initially, will be loaded from the file. | ||
/// All the discovered selectors will be merged into it. | ||
selectors: Selectors, | ||
/// Number of selectors before processing the files. | ||
/// Used for reporting. | ||
selectors_before: usize, | ||
/// Number of files analyzed. | ||
/// Used for reporting. | ||
analyzed_files: usize, | ||
} | ||
|
||
impl App { | ||
/// Loads the selectors from the file, or returns a new instance if the file doesn't exist. | ||
pub async fn load(file_path: impl Into<PathBuf>) -> anyhow::Result<Self> { | ||
let file_path = file_path.into(); | ||
// If doesn't exist, return default. | ||
if !file_path.exists() { | ||
return Ok(Self::default()); | ||
} | ||
|
||
let file = tokio::fs::read(&file_path) | ||
.await | ||
.context("Failed to read file")?; | ||
let selectors: Selectors = | ||
serde_json::from_slice(&file).context("Failed to deserialize file")?; | ||
let selectors_before = selectors.len(); | ||
Ok(Self { | ||
file_path, | ||
selectors, | ||
selectors_before, | ||
analyzed_files: 0, | ||
}) | ||
} | ||
|
||
/// Analyses all the JSON files, looking for 'abi' entries, and then computing the selectors for them. | ||
pub async fn process_files(&mut self, directory: &str) -> anyhow::Result<()> { | ||
for file_path in Self::load_file_paths(directory) { | ||
let Ok(new_selectors) = Selectors::load(&file_path).await.inspect_err(|e| { | ||
eprintln!("Error parsing file {file_path:?}: {e:?}"); | ||
}) else { | ||
continue; | ||
}; | ||
self.merge(new_selectors); | ||
} | ||
Ok(()) | ||
} | ||
|
||
/// Saves the selectors to the file. | ||
pub async fn save(self) -> anyhow::Result<()> { | ||
let mut file = tokio::fs::OpenOptions::new() | ||
.write(true) | ||
.create(true) | ||
.truncate(true) | ||
.open(self.file_path) | ||
.await | ||
.context("Failed to open file")?; | ||
let json = serde_json::to_string_pretty(&self.selectors)?; | ||
file.write_all(json.as_bytes()) | ||
.await | ||
.context("Failed to save file")?; | ||
Ok(()) | ||
} | ||
|
||
/// Merges the new selectors into the current ones. | ||
pub fn merge(&mut self, new: Selectors) { | ||
self.selectors.merge(new); | ||
self.analyzed_files += 1; | ||
} | ||
|
||
/// Reports the number of analyzed files and the number of added selectors. | ||
pub fn report(&self) { | ||
println!( | ||
"Analyzed {} files. Added {} selectors (before: {} after: {})", | ||
self.analyzed_files, | ||
self.selectors.len() - self.selectors_before, | ||
self.selectors_before, | ||
self.selectors.len() | ||
); | ||
} | ||
|
||
fn load_file_paths(dir: &str) -> Vec<PathBuf> { | ||
glob(&format!("{}/**/*.json", dir)) | ||
.expect("Failed to read glob pattern") | ||
.filter_map(|entry| match entry { | ||
Ok(path) => Some(path), | ||
Err(e) => { | ||
eprintln!("Error reading file: {:?}", e); | ||
None | ||
} | ||
}) | ||
.collect() | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,105 +1,33 @@ | ||
use std::{ | ||
collections::HashMap, | ||
fs::{File, OpenOptions}, | ||
io::{self}, | ||
}; | ||
|
||
use app::App; | ||
use clap::Parser; | ||
use glob::glob; | ||
use serde::{Deserialize, Serialize}; | ||
use sha3::{Digest, Keccak256}; | ||
|
||
#[derive(Debug, Serialize, Deserialize)] | ||
struct ABIEntry { | ||
#[serde(rename = "type")] | ||
entry_type: String, | ||
name: Option<String>, | ||
inputs: Option<Vec<ABIInput>>, | ||
} | ||
|
||
#[derive(Debug, Serialize, Deserialize)] | ||
struct ABIInput { | ||
#[serde(rename = "type")] | ||
input_type: String, | ||
} | ||
pub(crate) mod app; | ||
pub(crate) mod selectors; | ||
|
||
/// Selector generator tool. | ||
/// | ||
/// Generates a mapping of short (4-byte) function selectors to their corresponding function names. | ||
/// | ||
/// The generated JSON can be used to lookup function names by their selectors, when interacting | ||
/// with Ethereum contracts. | ||
#[derive(Debug, Parser)] | ||
#[command(author, version, about, long_about = None)] | ||
#[command(author, version, about, long_about)] | ||
struct Cli { | ||
/// Path to the directory with JSON files containing ABI. | ||
/// All JSON files in this directory will be processed. | ||
contracts_dir: String, | ||
/// Path to the output file. | ||
/// The file will contain the list of function selectors. | ||
/// If the file already exists, new selectors will be appended to it. | ||
output_file: String, | ||
} | ||
|
||
/// Computes solidity selector for a given method and arguments. | ||
fn compute_selector(name: &str, inputs: &[ABIInput]) -> String { | ||
let signature = format!( | ||
"{}({})", | ||
name, | ||
inputs | ||
.iter() | ||
.map(|i| i.input_type.clone()) | ||
.collect::<Vec<_>>() | ||
.join(",") | ||
); | ||
let mut hasher = Keccak256::new(); | ||
hasher.update(signature); | ||
format!("{:x}", hasher.finalize())[..8].to_string() | ||
} | ||
|
||
/// Analyses all the JSON files, looking for 'abi' entries, and then computing the selectors for them. | ||
fn process_files(directory: &str, output_file: &str) -> io::Result<()> { | ||
let mut selectors: HashMap<String, String> = match File::open(output_file) { | ||
Ok(file) => serde_json::from_reader(file).unwrap_or_default(), | ||
Err(_) => HashMap::new(), | ||
}; | ||
let selectors_before = selectors.len(); | ||
let mut analyzed_files = 0; | ||
|
||
for entry in glob(&format!("{}/**/*.json", directory)).expect("Failed to read glob pattern") { | ||
match entry { | ||
Ok(path) => { | ||
let file_path = path.clone(); | ||
let file = File::open(path)?; | ||
let json: Result<serde_json::Value, _> = serde_json::from_reader(file); | ||
|
||
if let Ok(json) = json { | ||
if let Some(abi) = json.get("abi").and_then(|v| v.as_array()) { | ||
analyzed_files += 1; | ||
for item in abi { | ||
let entry: ABIEntry = serde_json::from_value(item.clone()).unwrap(); | ||
if entry.entry_type == "function" { | ||
if let (Some(name), Some(inputs)) = (entry.name, entry.inputs) { | ||
let selector = compute_selector(&name, &inputs); | ||
selectors.entry(selector).or_insert(name); | ||
} | ||
} | ||
} | ||
} | ||
} else { | ||
eprintln!("Error parsing file: {:?} - ignoring.", file_path) | ||
} | ||
} | ||
Err(e) => eprintln!("Error reading file: {:?}", e), | ||
} | ||
} | ||
println!( | ||
"Analyzed {} files. Added {} selectors (before: {} after: {})", | ||
analyzed_files, | ||
selectors.len() - selectors_before, | ||
selectors_before, | ||
selectors.len() | ||
); | ||
|
||
let file = OpenOptions::new() | ||
.write(true) | ||
.create(true) | ||
.truncate(true) | ||
.open(output_file)?; | ||
serde_json::to_writer_pretty(file, &selectors)?; | ||
Ok(()) | ||
} | ||
|
||
fn main() -> io::Result<()> { | ||
#[tokio::main] | ||
async fn main() -> anyhow::Result<()> { | ||
let args = Cli::parse(); | ||
process_files(&args.contracts_dir, &args.output_file) | ||
let mut app = App::load(args.output_file).await?; | ||
app.process_files(&args.contracts_dir).await?; | ||
app.report(); | ||
app.save().await?; | ||
Ok(()) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,118 @@ | ||
use std::{collections::HashMap, path::PathBuf}; | ||
|
||
use anyhow::Context; | ||
use serde::{Deserialize, Serialize}; | ||
|
||
/// Short (4-byte) function selector. | ||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] | ||
#[serde(transparent)] | ||
struct Selector(String); | ||
|
||
/// Function name without parameters. | ||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] | ||
#[serde(transparent)] | ||
struct FunctionName(String); | ||
|
||
/// A set of function selectors and their corresponding function names. | ||
#[derive(Debug, Default, Serialize, Deserialize)] | ||
pub(crate) struct Selectors { | ||
#[serde(flatten)] | ||
selectors: HashMap<Selector, FunctionName>, | ||
} | ||
|
||
impl Selectors { | ||
/// Loads the selectors from the file, or returns a new instance if the file is a valid | ||
/// JSON, but doesn't contain `ABI` section. | ||
/// | ||
/// Will return an error if file doesn't exist or cannot be deserialized. | ||
pub async fn load(file_path: &PathBuf) -> anyhow::Result<Selectors> { | ||
let file = tokio::fs::read(file_path) | ||
.await | ||
.context("Failed to read file")?; | ||
let json: serde_json::Value = | ||
serde_json::from_slice(&file).context("Failed to deserialize file")?; | ||
let Some(abi) = json.get("abi").cloned() else { | ||
return Ok(Selectors::default()); | ||
}; | ||
|
||
let contract: ethabi::Contract = | ||
serde_json::from_value(abi).context("Failed to parse abi")?; | ||
Ok(Self::new(contract)) | ||
} | ||
|
||
/// Loads selectors from a given contract. | ||
pub fn new(contract: ethabi::Contract) -> Self { | ||
let selectors: HashMap<_, _> = contract | ||
.functions | ||
.into_values() | ||
.flatten() | ||
.map(|function| { | ||
let selector = hex::encode(function.short_signature()); | ||
(Selector(selector), FunctionName(function.name)) | ||
}) | ||
.collect(); | ||
Self { selectors } | ||
} | ||
|
||
/// Merges new selectors into the existing set. | ||
pub fn merge(&mut self, new: Self) { | ||
for (selector, name) in new.selectors { | ||
self.selectors | ||
.entry(selector.clone()) | ||
.and_modify(|e| { | ||
assert_eq!( | ||
e, &name, | ||
"Function name mismatch for selector '{:?}'", | ||
selector | ||
) | ||
}) | ||
.or_insert(name); | ||
} | ||
} | ||
|
||
pub fn len(&self) -> usize { | ||
self.selectors.len() | ||
} | ||
} | ||
|
||
#[cfg(test)] | ||
mod tests { | ||
use super::*; | ||
|
||
#[test] | ||
fn test_selectors() { | ||
let contract_json = r#"[ | ||
{ | ||
"type": "function", | ||
"name": "transfer", | ||
"inputs": [ | ||
{ "name": "to", "type": "address" }, | ||
{ "name": "value", "type": "uint256" } | ||
], | ||
"outputs": [], | ||
"stateMutability": "nonpayable" | ||
}, | ||
{ | ||
"type": "function", | ||
"name": "bar", | ||
"inputs": [], | ||
"outputs": [], | ||
"stateMutability": "nonpayable" | ||
} | ||
] | ||
"#; | ||
|
||
let contract: ethabi::Contract = serde_json::from_str(contract_json).unwrap(); | ||
let selectors = Selectors::new(contract); | ||
assert_eq!(selectors.len(), 2); | ||
|
||
// Check the generated selectors. | ||
assert_eq!( | ||
selectors | ||
.selectors | ||
.get(&Selector("a9059cbb".to_string())) | ||
.expect("No selector for transfer found"), | ||
&FunctionName("transfer".to_string()) | ||
); | ||
} | ||
} |
Oops, something went wrong.