Skip to content

Commit 4a536e8

Browse files
ananas-blockananas-block
authored
chore: add print state tree xtask (#2007)
1 parent 9ff1cc1 commit 4a536e8

2 files changed

Lines changed: 285 additions & 0 deletions

File tree

‎xtask/src/main.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ mod export_photon_test_data;
1010
mod fee;
1111
mod hash_set;
1212
mod new_deployment;
13+
mod print_state_tree;
1314
mod type_sizes;
1415
mod utils;
1516
mod zero_bytes;
@@ -57,6 +58,10 @@ enum Command {
5758
InitNewDeployment(new_deployment::Options),
5859
/// cargo xtask create-update-protocol-config --slot-length <u64>
5960
CreateUpdateProtocolConfigIx(create_update_protocol_config_ix::Options),
61+
/// Print batched state tree metadata
62+
/// Example:
63+
/// cargo xtask print-state-tree --pubkey <PUBKEY> --network mainnet
64+
PrintStateTree(print_state_tree::Options),
6065
}
6166

6267
#[tokio::main]
@@ -89,5 +94,6 @@ async fn main() -> Result<(), anyhow::Error> {
8994
Command::CreateUpdateProtocolConfigIx(opts) => {
9095
create_update_protocol_config_ix::create_update_protocol_config_ix(opts).await
9196
}
97+
Command::PrintStateTree(opts) => print_state_tree::print_state_tree(opts).await,
9298
}
9399
}

‎xtask/src/print_state_tree.rs‎

Lines changed: 279 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,279 @@
1+
use std::{mem, str::FromStr};
2+
3+
use clap::Parser;
4+
use light_batched_merkle_tree::merkle_tree::BatchedMerkleTreeAccount;
5+
use light_client::rpc::{LightClient, LightClientConfig, Rpc};
6+
use light_compressed_account::{pubkey::Pubkey, TreeType};
7+
use light_concurrent_merkle_tree::zero_copy::ConcurrentMerkleTreeZeroCopy;
8+
use light_hasher::Poseidon;
9+
use light_merkle_tree_metadata::merkle_tree::MerkleTreeMetadata;
10+
use solana_sdk::{account::Account, bs58, pubkey::Pubkey as SolanaPubkey};
11+
12+
#[derive(Debug, Parser)]
13+
pub struct Options {
14+
/// The pubkey of the state Merkle tree to print
15+
#[clap(long)]
16+
pubkey: String,
17+
/// Network: mainnet, devnet, local, or custom URL
18+
#[clap(long, default_value = "mainnet")]
19+
network: String,
20+
}
21+
22+
pub async fn print_state_tree(options: Options) -> anyhow::Result<()> {
23+
let rpc_url = match options.network.as_str() {
24+
"local" => String::from("http://127.0.0.1:8899"),
25+
"devnet" => String::from("https://api.devnet.solana.com"),
26+
"mainnet" => String::from("https://api.mainnet-beta.solana.com"),
27+
_ => options.network.clone(),
28+
};
29+
30+
let rpc = LightClient::new(LightClientConfig {
31+
url: rpc_url.clone(),
32+
photon_url: None,
33+
commitment_config: None,
34+
fetch_active_tree: false,
35+
api_key: None,
36+
})
37+
.await?;
38+
39+
let pubkey = SolanaPubkey::from_str(&options.pubkey)?;
40+
let pubkey_bytes: [u8; 32] = pubkey.to_bytes();
41+
let light_pubkey = Pubkey::new_from_array(pubkey_bytes);
42+
43+
println!("Fetching account: {}", pubkey);
44+
println!("RPC URL: {}", rpc_url);
45+
println!();
46+
47+
let account = rpc
48+
.get_account(pubkey)
49+
.await?
50+
.ok_or_else(|| anyhow::anyhow!("Account not found"))?;
51+
let mut account_data = account.data.clone();
52+
53+
// Try v2 (batched) first
54+
if let Ok(tree) = BatchedMerkleTreeAccount::state_from_bytes(&mut account_data, &light_pubkey) {
55+
print_v2_tree(&tree, &pubkey, &account)?;
56+
} else {
57+
// Try v1 (concurrent)
58+
print_v1_tree(&account_data, &pubkey, &account)?;
59+
}
60+
61+
Ok(())
62+
}
63+
64+
fn print_v2_tree(
65+
tree: &BatchedMerkleTreeAccount,
66+
pubkey: &SolanaPubkey,
67+
account: &Account,
68+
) -> anyhow::Result<()> {
69+
println!("=== Batched State Merkle Tree (V2) Metadata ===");
70+
println!();
71+
println!("Pubkey: {}", pubkey);
72+
println!("Tree Type: {:?}", TreeType::from(tree.tree_type));
73+
println!("Height: {}", tree.height);
74+
println!("Capacity: {}", tree.capacity);
75+
println!();
76+
77+
println!("=== Tree State ===");
78+
println!("Next Index: {}", tree.next_index);
79+
println!("Sequence Number: {}", tree.sequence_number);
80+
println!("Nullifier Next Index: {}", tree.nullifier_next_index);
81+
println!();
82+
83+
println!("=== Root History ===");
84+
println!("Root History Capacity: {}", tree.root_history_capacity);
85+
println!("Current Root Index: {}", tree.get_root_index());
86+
if let Some(current_root) = tree.get_root() {
87+
println!("Current Root: {}", bs58::encode(current_root).into_string());
88+
}
89+
println!();
90+
91+
println!("=== Access Metadata ===");
92+
println!(
93+
"Owner: {}",
94+
SolanaPubkey::from(tree.metadata.access_metadata.owner.to_bytes())
95+
);
96+
let program_owner = SolanaPubkey::from(tree.metadata.access_metadata.program_owner.to_bytes());
97+
if program_owner != SolanaPubkey::default() {
98+
println!("Program Owner: {}", program_owner);
99+
} else {
100+
println!("Program Owner: None");
101+
}
102+
let forester = SolanaPubkey::from(tree.metadata.access_metadata.forester.to_bytes());
103+
if forester != SolanaPubkey::default() {
104+
println!("Forester: {}", forester);
105+
} else {
106+
println!("Forester: None");
107+
}
108+
println!();
109+
110+
println!("=== Rollover Metadata ===");
111+
println!("Index: {}", tree.metadata.rollover_metadata.index);
112+
println!(
113+
"Rollover Fee: {}",
114+
tree.metadata.rollover_metadata.rollover_fee
115+
);
116+
let threshold = tree.metadata.rollover_metadata.rollover_threshold;
117+
if threshold > 0 {
118+
println!("Rollover Threshold: {}", threshold);
119+
} else {
120+
println!("Rollover Threshold: None");
121+
}
122+
println!(
123+
"Network Fee: {}",
124+
tree.metadata.rollover_metadata.network_fee
125+
);
126+
let next_merkle_tree = SolanaPubkey::from(tree.metadata.next_merkle_tree.to_bytes());
127+
if next_merkle_tree != SolanaPubkey::default() {
128+
println!("Next Merkle Tree: {}", next_merkle_tree);
129+
} else {
130+
println!("Next Merkle Tree: None");
131+
}
132+
println!();
133+
134+
println!("=== Queue Configuration ===");
135+
let associated_queue = SolanaPubkey::from(tree.metadata.associated_queue.to_bytes());
136+
if associated_queue != SolanaPubkey::default() {
137+
println!("Associated Queue: {}", associated_queue);
138+
} else {
139+
println!("Associated Queue: None");
140+
}
141+
println!("Num Batches: {}", tree.queue_batches.num_batches);
142+
println!("Batch Size: {}", tree.queue_batches.batch_size);
143+
println!("ZKP Batch Size: {}", tree.queue_batches.zkp_batch_size);
144+
println!(
145+
"Bloom Filter Capacity: {}",
146+
tree.queue_batches.bloom_filter_capacity
147+
);
148+
println!(
149+
"Currently Processing Batch Index: {}",
150+
tree.queue_batches.currently_processing_batch_index
151+
);
152+
println!(
153+
"Pending Batch Index: {}",
154+
tree.queue_batches.pending_batch_index
155+
);
156+
println!("Next Index: {}", tree.queue_batches.next_index);
157+
println!();
158+
159+
println!("=== Batch States ===");
160+
for (i, batch) in tree.queue_batches.batches.iter().enumerate() {
161+
println!("Batch {}:", i);
162+
println!(" State: {:?}", batch.get_state());
163+
println!(
164+
" Num Inserted Elements: {}",
165+
batch.get_num_inserted_elements()
166+
);
167+
println!(" Num Inserted ZKPs: {}", batch.get_num_inserted_zkps());
168+
println!(" Bloom Filter Zeroed: {}", batch.bloom_filter_is_zeroed());
169+
println!(" Start Index: {}", batch.start_index);
170+
println!(" Start Slot: {}", batch.start_slot);
171+
println!(" Sequence Number: {}", batch.sequence_number);
172+
println!(" Root Index: {}", batch.root_index);
173+
println!();
174+
}
175+
176+
println!("=== Account Info ===");
177+
println!("Account Size: {} bytes", account.data.len());
178+
println!("Lamports: {}", account.lamports);
179+
println!("Owner: {}", account.owner);
180+
println!();
181+
182+
Ok(())
183+
}
184+
185+
fn print_v1_tree(
186+
account_data: &[u8],
187+
pubkey: &SolanaPubkey,
188+
account: &Account,
189+
) -> anyhow::Result<()> {
190+
// Skip discriminator (8 bytes)
191+
let metadata_offset = 8;
192+
let metadata_size = mem::size_of::<MerkleTreeMetadata>();
193+
let metadata_bytes = &account_data[metadata_offset..metadata_offset + metadata_size];
194+
195+
// Safety: MerkleTreeMetadata is Pod and we're reading exactly the right size
196+
let metadata: &MerkleTreeMetadata =
197+
unsafe { &*(metadata_bytes.as_ptr() as *const MerkleTreeMetadata) };
198+
199+
// Parse the concurrent merkle tree
200+
let tree_data = &account_data[8 + mem::size_of::<MerkleTreeMetadata>()..];
201+
let tree = ConcurrentMerkleTreeZeroCopy::<Poseidon, 26>::from_bytes_zero_copy(tree_data)?;
202+
203+
println!("=== State Merkle Tree (V1) Metadata ===");
204+
println!();
205+
println!("Pubkey: {}", pubkey);
206+
println!("Height: {}", tree.height);
207+
println!("Canopy Depth: {}", tree.canopy_depth);
208+
println!("Capacity: {}", 2_usize.pow(tree.height as u32));
209+
println!();
210+
211+
println!("=== Tree State ===");
212+
println!("Next Index: {}", tree.next_index());
213+
println!("Sequence Number: {}", tree.sequence_number());
214+
println!();
215+
216+
println!("=== Root History ===");
217+
println!("Roots Capacity: {}", tree.roots.capacity());
218+
println!("Current Root Index: {}", tree.root_index());
219+
println!("Current Root: {}", bs58::encode(tree.root()).into_string());
220+
println!();
221+
222+
println!("=== Changelog ===");
223+
println!("Changelog Capacity: {}", tree.changelog.capacity());
224+
println!();
225+
226+
println!("=== Access Metadata ===");
227+
println!(
228+
"Owner: {}",
229+
SolanaPubkey::from(metadata.access_metadata.owner.to_bytes())
230+
);
231+
let program_owner = SolanaPubkey::from(metadata.access_metadata.program_owner.to_bytes());
232+
if program_owner != SolanaPubkey::default() {
233+
println!("Program Owner: {}", program_owner);
234+
} else {
235+
println!("Program Owner: None");
236+
}
237+
let forester = SolanaPubkey::from(metadata.access_metadata.forester.to_bytes());
238+
if forester != SolanaPubkey::default() {
239+
println!("Forester: {}", forester);
240+
} else {
241+
println!("Forester: None");
242+
}
243+
println!();
244+
245+
println!("=== Rollover Metadata ===");
246+
println!("Index: {}", metadata.rollover_metadata.index);
247+
println!("Rollover Fee: {}", metadata.rollover_metadata.rollover_fee);
248+
let threshold = metadata.rollover_metadata.rollover_threshold;
249+
if threshold > 0 {
250+
println!("Rollover Threshold: {}", threshold);
251+
} else {
252+
println!("Rollover Threshold: None");
253+
}
254+
println!("Network Fee: {}", metadata.rollover_metadata.network_fee);
255+
let next_merkle_tree = SolanaPubkey::from(metadata.next_merkle_tree.to_bytes());
256+
if next_merkle_tree != SolanaPubkey::default() {
257+
println!("Next Merkle Tree: {}", next_merkle_tree);
258+
} else {
259+
println!("Next Merkle Tree: None");
260+
}
261+
println!();
262+
263+
println!("=== Queue Configuration ===");
264+
let associated_queue = SolanaPubkey::from(metadata.associated_queue.to_bytes());
265+
if associated_queue != SolanaPubkey::default() {
266+
println!("Associated Queue: {}", associated_queue);
267+
} else {
268+
println!("Associated Queue: None");
269+
}
270+
println!();
271+
272+
println!("=== Account Info ===");
273+
println!("Account Size: {} bytes", account.data.len());
274+
println!("Lamports: {}", account.lamports);
275+
println!("Owner: {}", account.owner);
276+
println!();
277+
278+
Ok(())
279+
}

0 commit comments

Comments
 (0)