feat(katana): initialize slot paymaster account

bin/katana/Cargo.toml

@@ -23,7 +23,6 @@ clap_complete.workspace = true
 comfy-table = "7.1.1"
 dojo-utils.workspace = true
 inquire = "0.7.5"
-lazy_static.workspace = true
 piltover = { git = "https://github.com/keep-starknet-strange/piltover.git", rev = "fb9d988" }
 rand.workspace = true
 shellexpand = "3.1.0"
@@ -41,7 +40,8 @@ rstest.workspace = true
 starknet.workspace = true
 
 [features]
-default = [ "jemalloc", "katana-cli/slot" ]
+default = [ "init-slot", "jemalloc", "katana-cli/slot" ]
 
 init-custom-settlement-chain = [  ]
+init-slot = [  ]
 jemalloc = [  ]

bin/katana/src/cli/init/mod.rs

@@ -13,7 +13,6 @@
 use katana_primitives::genesis::constant::DEFAULT_PREFUNDED_ACCOUNT_BALANCE;
 use katana_primitives::genesis::Genesis;
 use katana_primitives::{ContractAddress, Felt, U256};
-use lazy_static::lazy_static;
 use prompt::CARTRIDGE_SN_SEPOLIA_PROVIDER;
 use starknet::accounts::{ExecutionEncoding, SingleOwnerAccount};
 use starknet::core::utils::{cairo_short_string_to_felt, parse_cairo_short_string};
@@ -24,6 +23,8 @@
 
 mod deployment;
 mod prompt;
+#[cfg(feature = "init-slot")]
+mod slot;
 
 #[derive(Debug, Args)]
 pub struct InitArgs {
@@ -54,13 +55,17 @@
     /// Specify the path of the directory where the configuration files will be stored at.
     #[arg(long)]
     output_path: Option<PathBuf>,
+
+    #[cfg(feature = "init-slot")]
+    #[command(flatten)]
+    slot: slot::SlotArgs,
 }
 
 impl InitArgs {
     // TODO:
     // - deploy bridge contract
     pub(crate) async fn execute(self) -> anyhow::Result<()> {
-        let output = if let Some(output) = self.process_args().await {
+        let output = if let Some(output) = self.configure_from_args().await {
             output?
         } else {
             prompt::prompt().await?
@@ -75,7 +80,12 @@
         };
 
         let id = ChainId::parse(&output.id)?;
-        let genesis = GENESIS.clone();
+
+        #[cfg_attr(not(feature = "init-slot"), allow(unused_mut))]
+        let mut genesis = generate_genesis();
+        #[cfg(feature = "init-slot")]
+        slot::add_paymasters_to_genesis(&mut genesis, &output.slot_paymasters.unwrap_or_default());
+
         // At the moment, the fee token is limited to a predefined token.
         let fee_contract = FeeContract::default();
         let chain_spec = rollup::ChainSpec { id, genesis, settlement, fee_contract };
@@ -92,7 +102,7 @@
         Ok(())
     }
 
-    async fn process_args(&self) -> Option<anyhow::Result<Outcome>> {
+    async fn configure_from_args(&self) -> Option<anyhow::Result<Outcome>> {
         // Here we just check that if `id` is present, then all the other required* arguments must
         // be present as well. This is guaranteed by `clap`.
         if let Some(id) = self.id.clone() {
@@ -144,6 +154,8 @@
                 rpc_url: settlement_url,
                 account: settlement_account_address,
                 settlement_id: parse_cairo_short_string(&l1_chain_id).unwrap(),
+                #[cfg(feature = "init-slot")]
+                slot_paymasters: self.slot.paymaster_accounts.clone(),
             }))
         } else {
             None
@@ -167,16 +179,18 @@
     pub rpc_url: Url,
 
     pub deployment_outcome: DeploymentOutcome,
+
+    #[cfg(feature = "init-slot")]
+    pub slot_paymasters: Option<Vec<slot::PaymasterAccountArgs>>,
 }
 
-lazy_static! {
-    static ref GENESIS: Genesis = {
-        // master account
-        let accounts = DevAllocationsGenerator::new(1).with_balance(U256::from(DEFAULT_PREFUNDED_ACCOUNT_BALANCE)).generate();
-        let mut genesis = Genesis::default();
-        genesis.extend_allocations(accounts.into_iter().map(|(k, v)| (k, v.into())));
-        genesis
-    };
+fn generate_genesis() -> Genesis {
+    let accounts = DevAllocationsGenerator::new(1)
+        .with_balance(U256::from(DEFAULT_PREFUNDED_ACCOUNT_BALANCE))
+        .generate();
+    let mut genesis = Genesis::default();
+    genesis.extend_allocations(accounts.into_iter().map(|(k, v)| (k, v.into())));
+    genesis
 }
 
 #[derive(Debug, thiserror::Error)]

bin/katana/src/cli/init/prompt.rs

@@ -125,11 +125,24 @@
             DeploymentOutcome { contract_address: address, block_number }
         };
 
+    // Prompt for slot paymaster accounts
+    let mut slot_paymasters = Vec::new();
+
+    while Confirm::new("Add slot paymaster account?").with_default(true).prompt()? {
+        let public_key = CustomType::<Felt>::new("Paymaster public key")
+            .with_formatter(&|input: Felt| format!("{input:#x}"))
+            .prompt()?;
+
+        slot_paymasters.push(super::slot::PaymasterAccountArgs { public_key });
+    }
+
     Ok(Outcome {
         id: chain_id,
         deployment_outcome,
-        account: account_address,
         rpc_url: settlement_url,
+        account: account_address,
         settlement_id: parse_cairo_short_string(&l1_chain_id)?,
+        #[cfg(feature = "init-slot")]
+        slot_paymasters: Some(slot_paymasters),
     })
 }

bin/katana/src/cli/init/slot.rs

@@ -0,0 +1,92 @@
+use std::str::FromStr;
+
+use anyhow::Result;
+use clap::Args;
+use katana_primitives::genesis::allocation::{
+    GenesisAccount, GenesisAccountAlloc, GenesisAllocation,
+};
+use katana_primitives::genesis::constant::{
+    DEFAULT_ACCOUNT_CLASS_HASH, DEFAULT_PREFUNDED_ACCOUNT_BALANCE,
+};
+use katana_primitives::genesis::Genesis;
+use katana_primitives::{ContractAddress, Felt, U256};
+
+#[derive(Debug, Args)]
+#[command(next_help_heading = "Slot options")]
+pub struct SlotArgs {
+    #[arg(long)]
+    pub slot: bool,
+
+    #[arg(requires_all = ["id", "slot"])]
+    #[arg(long = "slot.paymasters", value_delimiter = ',')]
+    pub paymaster_accounts: Option<Vec<PaymasterAccountArgs>>,
+}
+
+#[derive(Debug, Clone)]
+pub struct PaymasterAccountArgs {
+    /// The public key of the paymaster account.
+    pub public_key: Felt,
+}
+
+impl FromStr for PaymasterAccountArgs {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self> {
+        Ok(PaymasterAccountArgs { public_key: Felt::from_str(s)? })
+    }
+}
+
+pub fn add_paymasters_to_genesis(
+    genesis: &mut Genesis,
+    slot_paymasters: &[PaymasterAccountArgs],
+) -> Vec<ContractAddress> {
+    let mut accounts = Vec::with_capacity(slot_paymasters.len());
+
+    for paymaster in slot_paymasters {
+        let public_key = paymaster.public_key;
+        let class_hash = DEFAULT_ACCOUNT_CLASS_HASH;
+        let balance = U256::from(DEFAULT_PREFUNDED_ACCOUNT_BALANCE);
+
+        let (addr, account) = GenesisAccount::new_with_balance(public_key, class_hash, balance);
+        let account = GenesisAllocation::Account(GenesisAccountAlloc::Account(account));
+        accounts.push((addr, account));
+    }
+
+    let addresses: Vec<ContractAddress> = accounts.iter().map(|(addr, ..)| *addr).collect();
+    genesis.extend_allocations(accounts);
+
+    addresses
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_add_paymasters_to_genesis() {
+        let mut genesis = Genesis::default();
+        let mut paymasters = Vec::new();
+
+        for i in 0..3 {
+            paymasters.push(PaymasterAccountArgs { public_key: Felt::from(i) });
+        }
+
+        let expected_addresses = add_paymasters_to_genesis(&mut genesis, &paymasters);
+
+        for (i, addr) in expected_addresses.iter().enumerate() {
+            let account = genesis.allocations.get(addr).expect("account missing");
+            match account {
+                GenesisAllocation::Account(GenesisAccountAlloc::Account(account)) => {
+                    assert_eq!(account.public_key, Felt::from(i));
+                    assert_eq!(account.class_hash, DEFAULT_ACCOUNT_CLASS_HASH);
+                    assert_eq!(
+                        account.balance,
+                        Some(U256::from(DEFAULT_PREFUNDED_ACCOUNT_BALANCE))
+                    );
+                }
+
+                _ => panic!("Expected GenesisAccountAlloc::Account"),
+            }
+        }
+    }
+}