Skip to content

Commit

Permalink
perf: better cast create2 (#6212)
Browse files Browse the repository at this point in the history 
* perf: better cast create2

* chore: clippy

* docs

* fix: enforce either code or code hash
  • Loading branch information
@DaniPopes
DaniPopes authored Nov 3, 2023
1 parent f0528ad commit fa6b39c
Showing 1 changed file with 122 additions and 118 deletions.
240 changes: 122 additions & 118 deletions crates/cast/bin/cmd/create2.rs
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,14 @@
use alloy_primitives::{keccak256, Address, B256, U256};
use clap::Parser;
use eyre::{Result, WrapErr};
use rayon::prelude::*;
use regex::RegexSetBuilder;
use std::{str::FromStr, time::Instant};
use std::{
sync::{
atomic::{AtomicBool, Ordering},
Arc,
},
time::Instant,
};

/// CLI arguments for `cast create2`.
#[derive(Debug, Clone, Parser)]
Expand Down Expand Up @@ -39,18 +44,22 @@ pub struct Create2Args {
deployer: Address,

/// Init code of the contract to be deployed.
#[clap(short, long, value_name = "HEX", default_value = "")]
init_code: String,
#[clap(short, long, value_name = "HEX")]
init_code: Option<String>,

/// Init code hash of the contract to be deployed.
#[clap(alias = "ch", long, value_name = "HASH")]
#[clap(alias = "ch", long, value_name = "HASH", required_unless_present = "init_code")]
init_code_hash: Option<String>,

/// Number of threads to use. Defaults to and caps at the number of logical cores.
#[clap(short, long)]
jobs: Option<usize>,
}

#[allow(dead_code)]
pub struct Create2Output {
pub address: Address,
pub salt: U256,
pub salt: B256,
}

impl Create2Args {
Expand All @@ -63,6 +72,7 @@ impl Create2Args {
deployer,
init_code,
init_code_hash,
jobs,
} = self;

let mut regexs = vec![];
Expand Down Expand Up @@ -105,43 +115,85 @@ impl Create2Args {
let regex = RegexSetBuilder::new(regexs).case_insensitive(!case_sensitive).build()?;

let init_code_hash = if let Some(init_code_hash) = init_code_hash {
let mut a: [u8; 32] = [0; 32];
let init_code_hash_bytes = hex::decode(init_code_hash)?;
assert!(init_code_hash_bytes.len() == 32, "init code hash should be 32 bytes long");
a.copy_from_slice(&init_code_hash_bytes);
a.into()
} else {
let mut hash: [u8; 32] = [0; 32];
hex::decode_to_slice(init_code_hash, &mut hash)?;
hash.into()
} else if let Some(init_code) = init_code {
keccak256(hex::decode(init_code)?)
} else {
unreachable!();
};

let mut n_threads = std::thread::available_parallelism().map_or(1, |n| n.get());
if let Some(jobs) = jobs {
n_threads = n_threads.min(jobs);
}
if cfg!(test) {
n_threads = n_threads.min(2);
}

println!("Starting to generate deterministic contract address...");
let mut handles = Vec::with_capacity(n_threads);
let found = Arc::new(AtomicBool::new(false));
let timer = Instant::now();
let (salt, addr) = std::iter::repeat(())
.par_bridge()
.map(|_| {
let salt = B256::random();

let addr = deployer.create2(salt, init_code_hash).to_checksum(None);

(salt, addr)
})
.find_any(move |(_, addr)| {
let addr = addr.to_string();
let addr = addr.strip_prefix("0x").unwrap();
regex.matches(addr).into_iter().count() == regex.patterns().len()
})
.unwrap();

let salt = U256::from_be_bytes(*salt);
let address = Address::from_str(&addr).unwrap();

println!(
"Successfully found contract address in {} seconds.\nAddress: {}\nSalt: {}",
timer.elapsed().as_secs(),
addr,
salt
);

// Loops through all possible salts in parallel until a result is found.
// Each thread iterates over `(i..).step_by(n_threads)`.
for i in 0..n_threads {
// Create local copies for the thread.
let increment = n_threads;
let regex = regex.clone();
let regex_len = regex.patterns().len();
let found = Arc::clone(&found);
handles.push(std::thread::spawn(move || {
// Read the first bytes of the salt as a usize to be able to increment it.
struct B256Aligned(B256, [usize; 0]);
let mut salt = B256Aligned(B256::ZERO, []);
// SAFETY: B256 is aligned to `usize`.
let salt_word = unsafe { &mut *salt.0.as_mut_ptr().cast::<usize>() };

// Important: set the salt to the start value, otherwise all threads loop over the
// same values.
*salt_word = i;

let mut checksum = [0; 42];
loop {
// Stop if a result was found in another thread.
if found.load(Ordering::Relaxed) {
break None;
}

// Calculate the `CREATE2` address.
#[allow(clippy::needless_borrows_for_generic_args)]
let addr = deployer.create2(&salt.0, init_code_hash);

// Check if the the regex matches the calculated address' checksum.
let _ = addr.to_checksum_raw(&mut checksum, None);
// SAFETY: stripping 2 ASCII bytes ("0x") off of an already valid UTF-8 string
// is safe.
let s = unsafe { std::str::from_utf8_unchecked(checksum.get_unchecked(2..)) };
if regex.matches(s).into_iter().count() == regex_len {
// Notify other threads that we found a result.
found.store(true, Ordering::Relaxed);
break Some((salt.0, addr));
}

// Increment the salt for the next iteration.
*salt_word += increment;
}
}));
}

let results = handles.into_iter().filter_map(|h| h.join().unwrap()).collect::<Vec<_>>();
println!("Successfully found contract address(es) in {:?}", timer.elapsed());
for (i, (salt, address)) in results.iter().enumerate() {
if i > 0 {
println!("---");
}
println!("Address: {address}\nSalt: {salt} ({})", U256::from_be_bytes(salt.0));
}

let (salt, address) = results.into_iter().next().unwrap();
Ok(Create2Output { address, salt })
}
}
Expand All @@ -156,58 +208,51 @@ fn get_regex_hex_string(s: String) -> Result<String> {
#[cfg(test)]
mod tests {
use super::*;
use std::str::FromStr;

const DEPLOYER: &str = "0x4e59b44847b379578588920ca78fbf26c0b4956c";

#[test]
fn basic_create2() {
let mk_args = |args: &[&str]| {
Create2Args::parse_from(["foundry-cli", "--init-code-hash=0x0000000000000000000000000000000000000000000000000000000000000000"].iter().chain(args))
};

// even hex chars
let args = Create2Args::parse_from(["foundry-cli", "--starts-with", "aa"]);
let args = mk_args(&["--starts-with", "aa"]);
let create2_out = args.run().unwrap();
let address = create2_out.address;
let address = format!("{address:x}");
assert!(format!("{:x}", create2_out.address).starts_with("aa"));

assert!(address.starts_with("aa"));
let args = mk_args(&["--ends-with", "bb"]);
let create2_out = args.run().unwrap();
assert!(format!("{:x}", create2_out.address).ends_with("bb"));

// odd hex chars
let args = Create2Args::parse_from(["foundry-cli", "--starts-with", "aaa"]);
let args = mk_args(&["--starts-with", "aaa"]);
let create2_out = args.run().unwrap();
let address = create2_out.address;
let address = format!("{address:x}");
assert!(format!("{:x}", create2_out.address).starts_with("aaa"));

assert!(address.starts_with("aaa"));
let args = mk_args(&["--ends-with", "bbb"]);
let create2_out = args.run().unwrap();
assert!(format!("{:x}", create2_out.address).ends_with("bbb"));

// even hex chars with 0x prefix
let args = Create2Args::parse_from(["foundry-cli", "--starts-with", "0xaa"]);
let args = mk_args(&["--starts-with", "0xaa"]);
let create2_out = args.run().unwrap();
let address = create2_out.address;
let address = format!("{address:x}");

assert!(address.starts_with("aa"));
assert!(format!("{:x}", create2_out.address).starts_with("aa"));

// odd hex chars with 0x prefix
let args = Create2Args::parse_from(["foundry-cli", "--starts-with", "0xaaa"]);
let args = mk_args(&["--starts-with", "0xaaa"]);
let create2_out = args.run().unwrap();
let address = create2_out.address;
let address = format!("{address:x}");

assert!(address.starts_with("aaa"));

// odd hex chars with 0x suffix
let args = Create2Args::parse_from(["foundry-cli", "--ends-with", "bb"]);
let create2_out = args.run().unwrap();
let address = create2_out.address;
let address = format!("{address:x}");

assert!(address.ends_with("bb"));
assert!(format!("{:x}", create2_out.address).starts_with("aaa"));

// check fails on wrong chars
let args = Create2Args::parse_from(["foundry-cli", "--starts-with", "0xerr"]);
let args = mk_args(&["--starts-with", "0xerr"]);
let create2_out = args.run();
assert!(create2_out.is_err());

// check fails on wrong x prefixed string provided
let args = Create2Args::parse_from(["foundry-cli", "--starts-with", "x00"]);
let args = mk_args(&["--starts-with", "x00"]);
let create2_out = args.run();
assert!(create2_out.is_err());
}
Expand All @@ -216,88 +261,47 @@ mod tests {
fn matches_pattern() {
let args = Create2Args::parse_from([
"foundry-cli",
"--matching",
"0xbbXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX",
"--init-code-hash=0x0000000000000000000000000000000000000000000000000000000000000000",
"--matching=0xbbXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX",
]);
let create2_out = args.run().unwrap();
let address = create2_out.address;
let address = format!("{address:x}");

assert!(address.starts_with("bb"));
assert!(format!("{address:x}").starts_with("bb"));
}

#[test]
fn create2_init_code() {
let init_code = "00";
let args = Create2Args::parse_from([
"foundry-cli",
"--starts-with",
"cc",
"--init-code",
init_code,
]);
let args =
Create2Args::parse_from(["foundry-cli", "--starts-with=cc", "--init-code", init_code]);
let create2_out = args.run().unwrap();
let address = create2_out.address;
let address_str = format!("{address:x}");
assert!(format!("{address:x}").starts_with("cc"));
let salt = create2_out.salt;
let deployer = Address::from_str(DEPLOYER).unwrap();

assert!(address_str.starts_with("cc"));
assert_eq!(address, verify_create2(deployer, salt, hex::decode(init_code).unwrap()));
assert_eq!(address, deployer.create2_from_code(salt, hex::decode(init_code).unwrap()));
}

#[test]
fn create2_init_code_hash() {
let init_code_hash = "bc36789e7a1e281436464229828f817d6612f7b477d66591ff96a9e064bcc98a";
let args = Create2Args::parse_from([
"foundry-cli",
"--starts-with",
"dd",
"--starts-with=dd",
"--init-code-hash",
init_code_hash,
]);
let create2_out = args.run().unwrap();
let address = create2_out.address;
let address_str = format!("{address:x}");
assert!(format!("{address:x}").starts_with("dd"));

let salt = create2_out.salt;
let deployer = Address::from_str(DEPLOYER).unwrap();

assert!(address_str.starts_with("dd"));
assert_eq!(
address,
verify_create2_hash(deployer, salt, hex::decode(init_code_hash).unwrap())
deployer
.create2(salt, B256::from_slice(hex::decode(init_code_hash).unwrap().as_slice()))
);
}

#[test]
fn verify_helpers() {
// https://eips.ethereum.org/EIPS/eip-1014
let eip_address = Address::from_str("0x4D1A2e2bB4F88F0250f26Ffff098B0b30B26BF38").unwrap();

let deployer = Address::from_str("0x0000000000000000000000000000000000000000").unwrap();
let salt =
U256::from_str("0x0000000000000000000000000000000000000000000000000000000000000000")
.unwrap();
let init_code = hex::decode("00").unwrap();
let address = verify_create2(deployer, salt, init_code);

assert_eq!(address, eip_address);

let init_code_hash =
hex::decode("bc36789e7a1e281436464229828f817d6612f7b477d66591ff96a9e064bcc98a")
.unwrap();
let address = verify_create2_hash(deployer, salt, init_code_hash);

assert_eq!(address, eip_address);
}

fn verify_create2(deployer: Address, salt: U256, init_code: Vec<u8>) -> Address {
let init_code_hash = keccak256(init_code);
deployer.create2(salt.to_be_bytes(), init_code_hash)
}

fn verify_create2_hash(deployer: Address, salt: U256, init_code_hash: Vec<u8>) -> Address {
let init_code_hash = B256::from_slice(&init_code_hash);
deployer.create2(salt.to_be_bytes(), init_code_hash)
}
}

0 comments on commit fa6b39c

Please sign in to comment.