feat(katana): genesis state configuration (#1490)

* Define `katana` genesis configuration format (#1472)

* feat(katana): integrate genesis config (#1483)

* refactor(katana): update genesis with forked block (#1488)

* fix(katana): set forked parent hash

* Use `BTreeMap` for genesis allocations for sorted entries

bin/katana/Cargo.toml

@@ -7,6 +7,7 @@ repository.workspace = true
 version.workspace = true
 
 [dependencies]
+anyhow.workspace = true
 clap.workspace = true
 clap_complete.workspace = true
 console.workspace = true
@@ -22,6 +23,7 @@ tokio.workspace = true
 tracing-subscriber.workspace = true
 tracing.workspace = true
 url.workspace = true
+shellexpand = "3.1.0"
 
 [dev-dependencies]
 assert_matches = "1.5.0"

bin/katana/src/args.rs

@@ -21,13 +21,18 @@ use katana_core::constants::{
 };
 use katana_core::sequencer::SequencerConfig;
 use katana_primitives::chain::ChainId;
+use katana_primitives::genesis::allocation::DevAllocationsGenerator;
+use katana_primitives::genesis::constant::DEFAULT_PREFUNDED_ACCOUNT_BALANCE;
+use katana_primitives::genesis::Genesis;
 use katana_rpc::config::ServerConfig;
 use katana_rpc_api::ApiKind;
 use metrics::utils::parse_socket_address;
 use tracing::Subscriber;
 use tracing_subscriber::{fmt, EnvFilter};
 use url::Url;
 
+use crate::utils::{parse_genesis, parse_seed};
+
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
 #[command(propagate_version = true)]
@@ -147,6 +152,11 @@ pub struct StarknetOptions {
     #[command(flatten)]
     #[command(next_help_heading = "Environment options")]
     pub environment: EnvironmentOptions,
+
+    #[arg(long)]
+    #[arg(value_parser = parse_genesis)]
+    #[arg(conflicts_with_all(["rpc_url", "seed", "total_accounts"]))]
+    pub genesis: Option<Genesis>,
 }
 
 #[derive(Debug, Args, Clone)]
@@ -217,9 +227,21 @@ impl KatanaArgs {
     }
 
     pub fn starknet_config(&self) -> StarknetConfig {
+        let genesis = match self.starknet.genesis.clone() {
+            Some(genesis) => genesis,
+            None => {
+                let accounts = DevAllocationsGenerator::new(self.starknet.total_accounts)
+                    .with_seed(parse_seed(&self.starknet.seed))
+                    .with_balance(DEFAULT_PREFUNDED_ACCOUNT_BALANCE)
+                    .generate();
+
+                let mut genesis = Genesis::default();
+                genesis.extend_allocations(accounts.into_iter().map(|(k, v)| (k, v.into())));
+                genesis
+            }
+        };
+
         StarknetConfig {
-            total_accounts: self.starknet.total_accounts,
-            seed: parse_seed(&self.starknet.seed),
             disable_fee: self.starknet.disable_fee,
             disable_validate: self.starknet.disable_validate,
             fork_rpc_url: self.rpc_url.clone(),
@@ -239,22 +261,11 @@ impl KatanaArgs {
                     .unwrap_or(DEFAULT_VALIDATE_MAX_STEPS),
             },
             db_dir: self.db_dir.clone(),
+            genesis,
         }
     }
 }
 
-fn parse_seed(seed: &str) -> [u8; 32] {
-    let seed = seed.as_bytes();
-
-    if seed.len() >= 32 {
-        unsafe { *(seed[..32].as_ptr() as *const [u8; 32]) }
-    } else {
-        let mut actual_seed = [0u8; 32];
-        seed.iter().enumerate().for_each(|(i, b)| actual_seed[i] = *b);
-        actual_seed
-    }
-}
-
 #[cfg(test)]
 mod test {
     use super::*;

bin/katana/src/main.rs

@@ -4,16 +4,20 @@ use std::sync::Arc;
 use clap::{CommandFactory, Parser};
 use clap_complete::{generate, Shell};
 use console::Style;
-use katana_core::constants::{
-    ERC20_CONTRACT_CLASS_HASH, FEE_TOKEN_ADDRESS, UDC_ADDRESS, UDC_CLASS_HASH,
-};
 use katana_core::sequencer::KatanaSequencer;
+use katana_primitives::contract::ContractAddress;
+use katana_primitives::genesis::allocation::GenesisAccountAlloc;
+use katana_primitives::genesis::constant::{
+    DEFAULT_FEE_TOKEN_ADDRESS, DEFAULT_LEGACY_ERC20_CONTRACT_CLASS_HASH,
+    DEFAULT_LEGACY_UDC_CLASS_HASH, DEFAULT_UDC_ADDRESS,
+};
 use katana_rpc::{spawn, NodeHandle};
 use metrics::prometheus_exporter;
 use tokio::signal::ctrl_c;
 use tracing::info;
 
 mod args;
+mod utils;
 
 use args::Commands::Completions;
 use args::KatanaArgs;
@@ -40,8 +44,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let NodeHandle { addr, handle, .. } = spawn(Arc::clone(&sequencer), server_config).await?;
 
     if !config.silent {
-        let mut accounts = sequencer.backend.accounts.iter().peekable();
-        let account_class_hash = accounts.peek().unwrap().class_hash;
+        let mut accounts = sequencer.backend().config.genesis.accounts().peekable();
+        let account_class_hash = accounts.peek().unwrap().1.class_hash();
 
         if config.json_log {
             info!(
@@ -53,7 +57,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 })
             )
         } else {
-            let accounts = accounts.map(|a| format!("{a}")).collect::<Vec<_>>().join("\n");
             print_intro(
                 accounts,
                 config.starknet.seed.clone(),
@@ -91,14 +94,19 @@ fn print_completion(shell: Shell) {
     generate(shell, &mut command, name, &mut io::stdout());
 }
 
-fn print_intro(accounts: String, seed: String, address: String, account_class_hash: String) {
+fn print_intro<'a>(
+    accounts: impl Iterator<Item = (&'a ContractAddress, &'a GenesisAccountAlloc)>,
+    seed: String,
+    address: String,
+    account_class_hash: String,
+) {
     println!(
         "{}",
         Style::new().red().apply_to(
             r"
 
 
-██╗  ██╗ █████╗ ████████╗ █████╗ ███╗   ██╗ █████╗ 
+██╗  ██╗ █████╗ ████████╗ █████╗ ███╗   ██╗ █████╗
 ██║ ██╔╝██╔══██╗╚══██╔══╝██╔══██╗████╗  ██║██╔══██╗
 █████╔╝ ███████║   ██║   ███████║██╔██╗ ██║███████║
 ██╔═██╗ ██╔══██║   ██║   ██╔══██║██║╚██╗██║██╔══██║
@@ -110,35 +118,24 @@ fn print_intro(accounts: String, seed: String, address: String, account_class_ha
 
     println!(
         r"
-PREDEPLOYED CONTRACTS  
-================== 
+PREDEPLOYED CONTRACTS
+==================
 
 | Contract        | Fee Token
-| Address         | {}
-| Class Hash      | {:#064x}
+| Address         | {DEFAULT_FEE_TOKEN_ADDRESS}
+| Class Hash      | {DEFAULT_LEGACY_ERC20_CONTRACT_CLASS_HASH:#064x}
 
 | Contract        | Universal Deployer
-| Address         | {}
-| Class Hash      | {:#064x}
+| Address         | {DEFAULT_UDC_ADDRESS}
+| Class Hash      | {DEFAULT_LEGACY_UDC_CLASS_HASH:#064x}
 
 | Contract        | Account Contract
-| Class Hash      | {}
-    ",
-        *FEE_TOKEN_ADDRESS,
-        *ERC20_CONTRACT_CLASS_HASH,
-        *UDC_ADDRESS,
-        *UDC_CLASS_HASH,
-        account_class_hash
-    );
-
-    println!(
-        r"        
-PREFUNDED ACCOUNTS
-==================
-{accounts}
+| Class Hash      | {account_class_hash}
     "
     );
 
+    print_genesis_accounts(accounts);
+
     println!(
         r"
 ACCOUNTS SEED
@@ -149,3 +146,33 @@ ACCOUNTS SEED
 
     println!("\n{address}\n\n");
 }
+
+fn print_genesis_accounts<'a, Accounts>(accounts: Accounts)
+where
+    Accounts: Iterator<Item = (&'a ContractAddress, &'a GenesisAccountAlloc)>,
+{
+    println!(
+        r"
+PREFUNDED ACCOUNTS
+=================="
+    );
+
+    for (addr, account) in accounts {
+        if let Some(pk) = account.private_key() {
+            println!(
+                r"
+| Account address |  {addr}
+| Private key     |  {pk:#x}
+| Public key      |  {:#x}",
+                account.public_key()
+            )
+        } else {
+            println!(
+                r"
+| Account address |  {addr}
+| Public key      |  {:#x}",
+                account.public_key()
+            )
+        }
+    }
+}

bin/katana/src/utils.rs

@@ -0,0 +1,22 @@
+use std::path::PathBuf;
+
+use katana_primitives::genesis::json::GenesisJsonWithBasePath;
+use katana_primitives::genesis::Genesis;
+
+pub fn parse_seed(seed: &str) -> [u8; 32] {
+    let seed = seed.as_bytes();
+
+    if seed.len() >= 32 {
+        unsafe { *(seed[..32].as_ptr() as *const [u8; 32]) }
+    } else {
+        let mut actual_seed = [0u8; 32];
+        seed.iter().enumerate().for_each(|(i, b)| actual_seed[i] = *b);
+        actual_seed
+    }
+}
+
+pub fn parse_genesis(path: &str) -> Result<Genesis, anyhow::Error> {
+    let path = PathBuf::from(shellexpand::full(path)?.into_owned());
+    let genesis = Genesis::try_from(GenesisJsonWithBasePath::new(path)?)?;
+    Ok(genesis)
+}

crates/dojo-test-utils/src/sequencer.rs

@@ -51,9 +51,11 @@ impl TestSequencer {
 
         let url = Url::parse(&format!("http://{}", handle.addr)).expect("Failed to parse URL");
 
-        let account = sequencer.backend.accounts[0].clone();
-        let account =
-            TestAccount { private_key: account.private_key, account_address: account.address };
+        let account = sequencer.backend.config.genesis.accounts().next().unwrap();
+        let account = TestAccount {
+            private_key: account.1.private_key().unwrap(),
+            account_address: (*account.0).into(),
+        };
 
         TestSequencer { sequencer, account, handle, url }
     }