From 126575aa5e653ad57e70b76da639ec5a9d28eec0 Mon Sep 17 00:00:00 2001
From: ananas-block <58553958+ananas-block@users.noreply.github.com>
Date: Sat, 18 Jan 2025 17:39:46 +0000
Subject: [PATCH] chore: add create state Merkle tree script (#1485)

* chore: add create state Merkle tree script

* fix: lint
---
 scripts/create-state-trees.sh  | 38 ++++++++++++++++++++++++++++++++++
 xtask/src/create_state_tree.rs | 22 +++++---------------
 xtask/src/main.rs              |  2 +-
 3 files changed, 44 insertions(+), 18 deletions(-)
 create mode 100755 scripts/create-state-trees.sh

diff --git a/scripts/create-state-trees.sh b/scripts/create-state-trees.sh
new file mode 100755
index 000000000..baca42fce
--- /dev/null
+++ b/scripts/create-state-trees.sh
@@ -0,0 +1,38 @@
+#!/bin/bash
+
+# Base directory for keypairs
+KEYPAIR_DIR="./target/tree-keypairs"
+
+# Command template
+CMD_TEMPLATE="cargo xtask create-state-tree --mt-pubkey {SMT} --nfq-pubkey {NFQ} --cpi-pubkey {CPI} --index {INDEX} --network devnet"
+
+# Collect sorted key files for each type
+SMT_KEYS=($(ls $KEYPAIR_DIR/smt*.json | sort))
+NFQ_KEYS=($(ls $KEYPAIR_DIR/nfq*.json | sort))
+CPI_KEYS=($(ls $KEYPAIR_DIR/cpi*.json | sort))
+
+# Ensure equal number of keys for each type
+if [[ ${#SMT_KEYS[@]} -ne ${#NFQ_KEYS[@]} || ${#NFQ_KEYS[@]} -ne ${#CPI_KEYS[@]} ]]; then
+    echo "Error: Mismatched number of SMT, NFQ, and CPI key files."
+    exit 1
+fi
+
+# Execute the command for each triple
+for i in "${!SMT_KEYS[@]}"; do
+    SMT_KEY="${SMT_KEYS[i]}"
+    NFQ_KEY="${NFQ_KEYS[i]}"
+    CPI_KEY="${CPI_KEYS[i]}"
+    INDEX=$((i + 2))
+
+    # Replace placeholders in the command template
+    CMD=${CMD_TEMPLATE//\{SMT\}/"$SMT_KEY"}
+    CMD=${CMD//\{NFQ\}/"$NFQ_KEY"}
+    CMD=${CMD//\{CPI\}/"$CPI_KEY"}
+    CMD=${CMD//\{INDEX\}/"$INDEX"}
+
+    echo "Executing: $CMD"
+    eval "$CMD"
+
+done
+
+echo "All commands executed."
diff --git a/xtask/src/create_state_tree.rs b/xtask/src/create_state_tree.rs
index 957090550..9967d0a1b 100644
--- a/xtask/src/create_state_tree.rs
+++ b/xtask/src/create_state_tree.rs
@@ -12,8 +12,6 @@ pub struct Options {
     #[clap(long)]
     payer: Option<PathBuf>,
     #[clap(long)]
-    path: PathBuf,
-    #[clap(long)]
     mt_pubkey: String,
     #[clap(long)]
     nfq_pubkey: String,
@@ -40,24 +38,14 @@ pub async fn create_state_tree(options: Options) -> anyhow::Result<()> {
         String::from("https://api.mainnet-beta.solana.com")
     };
     let mut rpc = SolanaRpcConnection::new(rpc_url, None);
-    let path = if let Some(path) = options.path.to_str() {
-        String::from(path)
-    } else {
-        String::from("./target/state-trees/")
-    };
 
     let mut mt_keypairs: Vec<Keypair> = vec![];
     let mut nfq_keypairs: Vec<Keypair> = vec![];
     let mut cpi_keypairs: Vec<Keypair> = vec![];
 
-    println!(
-        "path mt_pubkey {:?}",
-        format!("{}-{}.json", path, options.mt_pubkey)
-    );
-
-    let mt_keypair = read_keypair_file(format!("{}{}.json", path, options.mt_pubkey)).unwrap();
-    let nfq_keypair = read_keypair_file(format!("{}{}.json", path, options.nfq_pubkey)).unwrap();
-    let cpi_keypair = read_keypair_file(format!("{}{}.json", path, options.cpi_pubkey)).unwrap();
+    let mt_keypair = read_keypair_file(options.mt_pubkey).unwrap();
+    let nfq_keypair = read_keypair_file(options.nfq_pubkey).unwrap();
+    let cpi_keypair = read_keypair_file(options.cpi_pubkey).unwrap();
     println!("read mt: {:?}", mt_keypair.pubkey());
     println!("read nfq: {:?}", nfq_keypair.pubkey());
     println!("read cpi: {:?}", cpi_keypair.pubkey());
@@ -65,8 +53,8 @@ pub async fn create_state_tree(options: Options) -> anyhow::Result<()> {
     nfq_keypairs.push(nfq_keypair);
     cpi_keypairs.push(cpi_keypair);
 
-    let payer = if let Some(payer) = options.payer {
-        read_keypair_file(&payer).unwrap_or_else(|_| panic!("{}{}.json", path, options.cpi_pubkey))
+    let payer = if let Some(payer) = options.payer.as_ref() {
+        read_keypair_file(payer).unwrap_or_else(|_| panic!("{:?}", options.payer))
     } else {
         // Construct the path to the keypair file in the user's home directory
         let keypair_path: PathBuf = home_dir()
diff --git a/xtask/src/main.rs b/xtask/src/main.rs
index 00baa24b0..baf8b2ef9 100644
--- a/xtask/src/main.rs
+++ b/xtask/src/main.rs
@@ -41,7 +41,7 @@ enum Command {
     HashSet(hash_set::HashSetOptions),
     /// Create state tree
     /// Example:
-    /// cargo xtask create-state-tree --path ./target/state-trees/  --mt-pubkey smtAvYA5UbTRyKAkAj5kHs1CmrA42t6WkVLi4c6mA1f --nfq-pubkey nfqAroCRkcZBgsAJDNkptKpsSWyM6cgB9XpWNNiCEC4 --cpi-pubkey cpiAb2eNFf6MQeqMWEyEjSN3VJcD5hghujhmtdcMuZp --index 10
+    /// cargo xtask create-state-tree --mt-pubkey ./target/tree-keypairs/smtAvYA5UbTRyKAkAj5kHs1CmrA42t6WkVLi4c6mA1f.json --nfq-pubkey ./target/tree-keypairs/nfqAroCRkcZBgsAJDNkptKpsSWyM6cgB9XpWNNiCEC4.json --cpi-pubkey ./target/tree-keypairs/cpiAb2eNFf6MQeqMWEyEjSN3VJcD5hghujhmtdcMuZp.json --index 10 --network local
     CreateStateTree(create_state_tree::Options),
 }