Skip to content

Commit

Permalink
refactor(selector-generator): Polish the code (#2888)
Browse files Browse the repository at this point in the history
## 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
popzxc authored Sep 16, 2024
1 parent af614fe commit 9d40704
Show file tree
Hide file tree
Showing 6 changed files with 1,240 additions and 576 deletions.
5 changes: 4 additions & 1 deletion Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

7 changes: 5 additions & 2 deletions core/bin/selector_generator/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,11 @@ categories.workspace = true
publish = false

[dependencies]
anyhow.workspace = true
serde = { workspace = true, features = ["derive"] }
serde_json.workspace = true
sha3.workspace = true
glob.workspace = true
clap = { workspace = true, features = ["derive"] }
clap = { workspace = true, features = ["derive"] }
ethabi.workspace = true
hex.workspace = true
tokio = { workspace = true, features = ["full"] }
105 changes: 105 additions & 0 deletions core/bin/selector_generator/src/app.rs
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()
}
}
116 changes: 22 additions & 94 deletions core/bin/selector_generator/src/main.rs
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(())
}
118 changes: 118 additions & 0 deletions core/bin/selector_generator/src/selectors.rs
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())
);
}
}
Loading

0 comments on commit 9d40704

Please sign in to comment.