update create-merkle-tree to discover next airdrop version for mint

Author: wphan

cli/src/bin/cli.rs

@@ -124,6 +124,9 @@ pub enum Commands {
     TotalClaim(TotalClaimAgrs),
 
     SetClawbackReceiver(ClawbackReceiverArgs),
+
+    /// Iteratively find airdrop version (PDA that don't exist yet)
+    FindAirdropVersion(FindAirdropVersionArgs),
 }
 
 #[derive(Parser, Debug)]
@@ -177,6 +180,13 @@ pub struct VerifyArgs {
     pub skip_verify_amount: bool,
 }
 
+// NewDistributor subcommand args
+#[derive(Parser, Debug)]
+pub struct FindAirdropVersionArgs {
+    #[clap(long, env)]
+    pub start_airdrop_version: Option<u64>,
+}
+
 // NewDistributor subcommand args
 #[derive(Parser, Debug)]
 pub struct NewDistributorArgs {
@@ -236,6 +246,10 @@ pub struct CreateMerkleTreeArgs {
     pub amount: u64,
     #[clap(long, env)]
     pub decimals: u32,
+    
+    /// Optional starting airdrop version. If not provided, will auto-detect next available version
+    #[clap(long, env)]
+    pub start_airdrop_version: Option<u64>,
 }
 
 #[derive(Parser, Debug)]
@@ -423,7 +437,7 @@ fn main() {
         }
         Commands::Clawback(clawback_args) => process_clawback(&args, clawback_args),
         Commands::CreateMerkleTree(merkle_tree_args) => {
-            process_create_merkle_tree(merkle_tree_args);
+            process_create_merkle_tree(&args, merkle_tree_args);
         }
         Commands::SetAdmin(set_admin_args) => {
             process_set_admin(&args, set_admin_args);
@@ -473,6 +487,9 @@ fn main() {
         Commands::SetClawbackReceiver(set_clawback_receiver_argrs) => {
             process_set_clawback_receiver(&args, set_clawback_receiver_argrs)
         }
+        Commands::FindAirdropVersion(find_airdrop_version_args) => {
+            process_find_airdrop_version(&args, find_airdrop_version_args);
+        }
     }
 }
 

cli/src/bin/instructions/mod.rs

@@ -42,3 +42,5 @@ pub mod verify_kv_proof;
 pub use verify_kv_proof::*;
 pub mod process_set_clawback_receiver;
 pub use process_set_clawback_receiver::*;
+pub mod process_find_airdrop_version;
+pub use process_find_airdrop_version::*;

cli/src/bin/instructions/process_create_merkle_tree.rs

@@ -2,7 +2,42 @@ use std::collections::HashMap;
 
 use crate::*;
 
-pub fn process_create_merkle_tree(merkle_tree_args: &CreateMerkleTreeArgs) {
+pub fn process_create_merkle_tree(args: &Args, merkle_tree_args: &CreateMerkleTreeArgs) {
+    // Find the next available airdrop version
+    let mut start_airdrop_version = 0;
+    if args.mint == Pubkey::default() {
+        println!("Mint is not set, will start creating merkle trees from version 0");
+    }
+    if args.mint != Pubkey::default() && merkle_tree_args.start_airdrop_version.is_none() {
+        println!("Finding next available airdrop version for mint: {}", args.mint);
+        
+        let program_client = args.get_program_client();
+        let rpc_client = program_client.rpc();
+        
+        let mut current_version = 0;
+        while current_version < 1000000000000000000 {
+            let (distributor_pubkey, _bump) =
+                get_merkle_distributor_pda(&args.program_id, &args.mint, current_version);
+            match rpc_client.get_account_data(&distributor_pubkey) {
+                Ok(_) => {
+                    println!("Airdrop version {} exists", current_version);
+                }
+                Err(e) => {
+                    if e.to_string().contains("AccountNotFound") {
+                        println!("Found next available airdrop version: {}", current_version);
+                        start_airdrop_version = current_version;
+                        break;
+                    } else {
+                        println!("Failed to get PDA, Error: {}", e.to_string());
+                    }
+                }
+            }
+            current_version += 1;
+        }
+    } else if let Some(version) = merkle_tree_args.start_airdrop_version {
+        start_airdrop_version = version;
+        println!("Using provided start airdrop version: {}", start_airdrop_version);
+    }
     let mut csv_entries = CsvEntry::new_from_file(&merkle_tree_args.csv_path).unwrap();
 
     // exclude test address if have
@@ -30,28 +65,28 @@ pub fn process_create_merkle_tree(merkle_tree_args: &CreateMerkleTreeArgs) {
     let max_nodes_per_tree = merkle_tree_args.max_nodes_per_tree as usize;
 
     let base_path = &merkle_tree_args.merkle_tree_path;
-    let mut index = 0;
+    let mut airdrop_version = start_airdrop_version;
     while csv_entries.len() > 0 {
         let last_index = max_nodes_per_tree.min(csv_entries.len());
         let sub_tree = csv_entries[0..last_index].to_vec();
         csv_entries = csv_entries[last_index..csv_entries.len()].to_vec();
 
-        // use index as version
+        // use airdrop_version as version
         let merkle_tree =
-            AirdropMerkleTree::new_from_entries(sub_tree, index, merkle_tree_args.decimals)
+            AirdropMerkleTree::new_from_entries(sub_tree, airdrop_version, merkle_tree_args.decimals)
                 .unwrap();
 
         let base_path_clone = base_path.clone();
         let path = base_path_clone
             .as_path()
-            .join(format!("tree_{}.json", index));
+            .join(format!("tree_{}.json", airdrop_version));
 
         merkle_tree.write_to_file(&path);
-        index += 1;
+        airdrop_version += 1;
     }
 
     if merkle_tree_args.should_include_test_list {
-        println!("create merkle tree for test claming index {}", index);
+        println!("create merkle tree for test claming version {}", airdrop_version);
         let test_list = get_test_list()
             .into_iter()
             .map(|x| CsvEntry {
@@ -62,13 +97,17 @@ pub fn process_create_merkle_tree(merkle_tree_args: &CreateMerkleTreeArgs) {
             .collect::<Vec<CsvEntry>>();
 
         let merkle_tree =
-            AirdropMerkleTree::new_from_entries(test_list, index, merkle_tree_args.decimals as u32)
+            AirdropMerkleTree::new_from_entries(test_list, airdrop_version, merkle_tree_args.decimals as u32)
                 .unwrap();
         let base_path_clone = base_path.clone();
         let path = base_path_clone
             .as_path()
-            .join(format!("tree_{}.json", index));
+            .join(format!("tree_{}.json", airdrop_version));
 
         merkle_tree.write_to_file(&path);
+        airdrop_version += 1;
     }
+    
+    // Output the last airdrop version used
+    println!("Last airdrop version created: {}", airdrop_version - 1);
 }

cli/src/bin/instructions/process_create_test_list.rs

@@ -18,6 +18,7 @@ pub fn process_create_test_list(args: &Args, create_test_list_args: &CreateTestL
         should_include_test_list: false,
         amount: create_test_list_args.amount,
         decimals: create_test_list_args.decimals,
+        start_airdrop_version: None,
     };
-    process_create_merkle_tree(merkle_tree_args);
+    process_create_merkle_tree(args, merkle_tree_args);
 }

cli/src/bin/instructions/process_find_airdrop_version.rs

@@ -0,0 +1,30 @@
+use crate::*;
+
+pub fn process_find_airdrop_version(args: &Args, new_distributor_args: &FindAirdropVersionArgs) {
+    let program_client = args.get_program_client();
+    let rpc_client = program_client.rpc();
+
+    let start_version_check = new_distributor_args.start_airdrop_version.unwrap_or(0);
+    let mut current_version = start_version_check;
+
+    println!("Starting to find airdrop version for mint: {}, starting from: {}", args.mint, start_version_check);
+
+    while current_version < 1000000000000000000 {
+        let (distributor_pubkey, _bump) =
+            get_merkle_distributor_pda(&args.program_id, &args.mint, current_version);
+        match rpc_client.get_account_data(&distributor_pubkey) {
+            Ok(_) => {
+                println!("Airdrop version {} exists, account: {}", current_version, distributor_pubkey);
+            }
+            Err(e) => {
+                if e.to_string().contains("AccountNotFound") {
+                    println!("Airdrop version does not exist: {} <- next airdrop_version", current_version);
+                    break;
+                } else {
+                    println!("Failed to get PDA, Error: {}", e.to_string());
+                }
+            }
+        }
+        current_version += 1;
+    }
+}