refactor(selector-generator): Polish the code

core/bin/selector_generator/Cargo.toml

@@ -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"] }

core/bin/selector_generator/src/app.rs

@@ -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()
+    }
+}

core/bin/selector_generator/src/main.rs

@@ -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(())
 }

core/bin/selector_generator/src/selectors.rs

@@ -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())
+        );
+    }
+}