Skip to content

Commit 159dd67

Browse files
ananas-blockananas-block
authored
chore: batched Merkle tree cleanup (#1561)
* cleanup * rename hashchain -> hash_chain * chore: improve test coverage * chore: cleanup, rename index -> tree_account_index
1 parent c143c24 commit 159dd67

File tree

67 files changed

+1760
-1027
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

67 files changed

+1760
-1027
lines changed

‎Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎forester-utils/src/address_merkle_tree_config.rs

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -169,8 +169,11 @@ pub async fn state_tree_ready_for_rollover<R: RpcConnection>(
169169
(tree.next_index(), tree_meta_data, 26)
170170
}
171171
BatchedMerkleTreeAccount::DISCRIMINATOR => {
172-
let tree_meta_data =
173-
BatchedMerkleTreeAccount::state_from_bytes(account.data.as_mut_slice()).unwrap();
172+
let tree_meta_data = BatchedMerkleTreeAccount::state_from_bytes(
173+
account.data.as_mut_slice(),
174+
&merkle_tree.into(),
175+
)
176+
.unwrap();
174177

175178
(
176179
tree_meta_data.next_index as usize,

‎forester-utils/src/instructions.rs

Lines changed: 33 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -59,22 +59,24 @@ where
5959
})?
6060
.unwrap();
6161

62-
let (leaves_hashchain, start_index, current_root, batch_size, full_batch_index) = {
63-
let merkle_tree =
64-
BatchedMerkleTreeAccount::address_from_bytes(merkle_tree_account.data.as_mut_slice())
65-
.unwrap();
62+
let (leaves_hash_chain, start_index, current_root, batch_size, full_batch_index) = {
63+
let merkle_tree = BatchedMerkleTreeAccount::address_from_bytes(
64+
merkle_tree_account.data.as_mut_slice(),
65+
&merkle_tree_pubkey.into(),
66+
)
67+
.unwrap();
6668

67-
let full_batch_index = merkle_tree.queue_metadata.next_full_batch_index;
68-
let batch = &merkle_tree.queue_metadata.batches[full_batch_index as usize];
69+
let full_batch_index = merkle_tree.queue_batches.pending_batch_index;
70+
let batch = &merkle_tree.queue_batches.batches[full_batch_index as usize];
6971
let zkp_batch_index = batch.get_num_inserted_zkps();
70-
let leaves_hashchain =
72+
let leaves_hash_chain =
7173
merkle_tree.hash_chain_stores[full_batch_index as usize][zkp_batch_index as usize];
7274
let start_index = merkle_tree.next_index;
7375
let current_root = *merkle_tree.root_history.last().unwrap();
7476
let batch_size = batch.zkp_batch_size as usize;
7577

7678
(
77-
leaves_hashchain,
79+
leaves_hash_chain,
7880
start_index,
7981
current_root,
8082
batch_size,
@@ -161,7 +163,7 @@ where
161163
low_element_proofs,
162164
addresses,
163165
subtrees,
164-
leaves_hashchain,
166+
leaves_hash_chain,
165167
batch_start_index,
166168
batch_size,
167169
)
@@ -214,35 +216,37 @@ pub async fn create_append_batch_ix_data<R: RpcConnection, I: Indexer<R>>(
214216
) -> Result<InstructionDataBatchAppendInputs, ForesterUtilsError> {
215217
let (merkle_tree_next_index, current_root) = {
216218
let mut merkle_tree_account = rpc.get_account(merkle_tree_pubkey).await.unwrap().unwrap();
217-
let merkle_tree =
218-
BatchedMerkleTreeAccount::state_from_bytes(merkle_tree_account.data.as_mut_slice())
219-
.unwrap();
219+
let merkle_tree = BatchedMerkleTreeAccount::state_from_bytes(
220+
merkle_tree_account.data.as_mut_slice(),
221+
&merkle_tree_pubkey.into(),
222+
)
223+
.unwrap();
220224
(
221225
merkle_tree.next_index,
222226
*merkle_tree.root_history.last().unwrap(),
223227
)
224228
};
225229

226-
let (zkp_batch_size, full_batch_index, num_inserted_zkps, leaves_hashchain) = {
230+
let (zkp_batch_size, full_batch_index, num_inserted_zkps, leaves_hash_chain) = {
227231
let mut output_queue_account = rpc.get_account(output_queue_pubkey).await.unwrap().unwrap();
228232
let output_queue =
229233
BatchedQueueAccount::output_from_bytes(output_queue_account.data.as_mut_slice())
230234
.unwrap();
231235

232-
let full_batch_index = output_queue.batch_metadata.next_full_batch_index;
236+
let full_batch_index = output_queue.batch_metadata.pending_batch_index;
233237
let zkp_batch_size = output_queue.batch_metadata.zkp_batch_size;
234238

235239
let num_inserted_zkps =
236240
output_queue.batch_metadata.batches[full_batch_index as usize].get_num_inserted_zkps();
237241

238-
let leaves_hashchain =
242+
let leaves_hash_chain =
239243
output_queue.hash_chain_stores[full_batch_index as usize][num_inserted_zkps as usize];
240244

241245
(
242246
zkp_batch_size,
243247
full_batch_index,
244248
num_inserted_zkps,
245-
leaves_hashchain,
249+
leaves_hash_chain,
246250
)
247251
};
248252
let start = num_inserted_zkps as usize * zkp_batch_size as usize;
@@ -282,7 +286,7 @@ pub async fn create_append_batch_ix_data<R: RpcConnection, I: Indexer<R>>(
282286
current_root,
283287
merkle_tree_next_index as u32,
284288
leaves,
285-
leaves_hashchain,
289+
leaves_hash_chain,
286290
old_leaves,
287291
merkle_proofs,
288292
zkp_batch_size as u32,
@@ -336,17 +340,20 @@ pub async fn create_nullify_batch_ix_data<R: RpcConnection, I: Indexer<R>>(
336340
indexer: &mut I,
337341
merkle_tree_pubkey: Pubkey,
338342
) -> Result<InstructionDataBatchNullifyInputs, ForesterUtilsError> {
339-
let (zkp_batch_size, old_root, leaves_hashchain) = {
343+
let (zkp_batch_size, old_root, leaves_hash_chain) = {
340344
let mut account = rpc.get_account(merkle_tree_pubkey).await.unwrap().unwrap();
341-
let merkle_tree =
342-
BatchedMerkleTreeAccount::state_from_bytes(account.data.as_mut_slice()).unwrap();
343-
let batch_idx = merkle_tree.queue_metadata.next_full_batch_index as usize;
344-
let zkp_size = merkle_tree.queue_metadata.zkp_batch_size;
345-
let batch = &merkle_tree.queue_metadata.batches[batch_idx];
345+
let merkle_tree = BatchedMerkleTreeAccount::state_from_bytes(
346+
account.data.as_mut_slice(),
347+
&merkle_tree_pubkey.into(),
348+
)
349+
.unwrap();
350+
let batch_idx = merkle_tree.queue_batches.pending_batch_index as usize;
351+
let zkp_size = merkle_tree.queue_batches.zkp_batch_size;
352+
let batch = &merkle_tree.queue_batches.batches[batch_idx];
346353
let zkp_idx = batch.get_num_inserted_zkps();
347-
let hashchain = merkle_tree.hash_chain_stores[batch_idx][zkp_idx as usize];
354+
let hash_chain = merkle_tree.hash_chain_stores[batch_idx][zkp_idx as usize];
348355
let root = *merkle_tree.root_history.last().unwrap();
349-
(zkp_size, root, hashchain)
356+
(zkp_size, root, hash_chain)
350357
};
351358

352359
let leaf_indices_tx_hashes =
@@ -383,7 +390,7 @@ pub async fn create_nullify_batch_ix_data<R: RpcConnection, I: Indexer<R>>(
383390
old_root,
384391
tx_hashes,
385392
leaves.to_vec(),
386-
leaves_hashchain,
393+
leaves_hash_chain,
387394
old_leaves,
388395
merkle_proofs,
389396
path_indices,

‎forester/src/batch_processor/common.rs

Lines changed: 17 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -120,13 +120,14 @@ impl<R: RpcConnection, I: Indexer<R> + IndexerType<R>> BatchProcessor<R, I> {
120120
}
121121

122122
fn calculate_completion_from_tree(data: &mut [u8]) -> f64 {
123-
let tree = match BatchedMerkleTreeAccount::state_from_bytes(data) {
123+
let tree = match BatchedMerkleTreeAccount::state_from_bytes(data, &Pubkey::default().into())
124+
{
124125
Ok(tree) => tree,
125126
Err(_) => return 0.0,
126127
};
127128

128-
let batch_index = tree.queue_metadata.next_full_batch_index;
129-
match tree.queue_metadata.batches.get(batch_index as usize) {
129+
let batch_index = tree.queue_batches.pending_batch_index;
130+
match tree.queue_batches.batches.get(batch_index as usize) {
130131
Some(batch) => Self::calculate_completion(batch),
131132
None => 0.0,
132133
}
@@ -138,7 +139,7 @@ impl<R: RpcConnection, I: Indexer<R> + IndexerType<R>> BatchProcessor<R, I> {
138139
Err(_) => return 0.0,
139140
};
140141

141-
let batch_index = queue.batch_metadata.next_full_batch_index;
142+
let batch_index = queue.batch_metadata.pending_batch_index;
142143
match queue.batch_metadata.batches.get(batch_index as usize) {
143144
Some(batch) => Self::calculate_completion(batch),
144145
None => 0.0,
@@ -177,7 +178,7 @@ impl<R: RpcConnection, I: Indexer<R> + IndexerType<R>> BatchProcessor<R, I> {
177178
BatchedQueueAccount::output_from_bytes(output_queue_account.data.as_mut_slice())
178179
.map_err(|e| BatchProcessError::QueueParsing(e.to_string()))?;
179180

180-
let batch_index = output_queue.batch_metadata.next_full_batch_index;
181+
let batch_index = output_queue.batch_metadata.pending_batch_index;
181182
let zkp_batch_size = output_queue.batch_metadata.zkp_batch_size;
182183

183184
(
@@ -195,19 +196,21 @@ impl<R: RpcConnection, I: Indexer<R> + IndexerType<R>> BatchProcessor<R, I> {
195196
};
196197

197198
let merkle_tree = match self.tree_type {
198-
TreeType::BatchedAddress => {
199-
BatchedMerkleTreeAccount::address_from_bytes(account.data.as_mut_slice())
200-
}
201-
TreeType::BatchedState => {
202-
BatchedMerkleTreeAccount::state_from_bytes(account.data.as_mut_slice())
203-
}
199+
TreeType::BatchedAddress => BatchedMerkleTreeAccount::address_from_bytes(
200+
account.data.as_mut_slice(),
201+
&self.context.merkle_tree.into(),
202+
),
203+
TreeType::BatchedState => BatchedMerkleTreeAccount::state_from_bytes(
204+
account.data.as_mut_slice(),
205+
&self.context.merkle_tree.into(),
206+
),
204207
_ => return false,
205208
};
206209

207210
if let Ok(tree) = merkle_tree {
208-
let batch_index = tree.queue_metadata.next_full_batch_index;
211+
let batch_index = tree.queue_batches.pending_batch_index;
209212
let full_batch = tree
210-
.queue_metadata
213+
.queue_batches
211214
.batches
212215
.get(batch_index as usize)
213216
.unwrap();
@@ -233,7 +236,7 @@ impl<R: RpcConnection, I: Indexer<R> + IndexerType<R>> BatchProcessor<R, I> {
233236
};
234237

235238
if let Ok(queue) = output_queue {
236-
let batch_index = queue.batch_metadata.next_full_batch_index;
239+
let batch_index = queue.batch_metadata.pending_batch_index;
237240
let full_batch = queue
238241
.batch_metadata
239242
.batches

‎forester/src/batch_processor/state.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,8 +110,9 @@ async fn get_batch_index<R: RpcConnection, I: Indexer<R>>(
110110
let merkle_tree =
111111
light_batched_merkle_tree::merkle_tree::BatchedMerkleTreeAccount::state_from_bytes(
112112
account.data.as_mut_slice(),
113+
&context.merkle_tree.into(),
113114
)
114115
.map_err(|e| BatchProcessError::MerkleTreeParsing(e.to_string()))?;
115116

116-
Ok(merkle_tree.queue_metadata.next_full_batch_index as usize)
117+
Ok(merkle_tree.queue_batches.pending_batch_index as usize)
117118
}

‎forester/src/rollover/operations.rs

Lines changed: 20 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -107,36 +107,37 @@ pub async fn get_tree_fullness<R: RpcConnection>(
107107
TreeType::BatchedState => {
108108
let mut account = rpc.get_account(tree_pubkey).await?.unwrap();
109109
let merkle_tree =
110-
BatchedMerkleTreeAccount::state_from_bytes(&mut account.data).unwrap();
110+
BatchedMerkleTreeAccount::state_from_bytes(&mut account.data, &tree_pubkey.into())
111+
.unwrap();
111112
println!(
112113
"merkle_tree.get_account().queue.batch_size: {:?}",
113-
merkle_tree.queue_metadata.batch_size
114+
merkle_tree.queue_batches.batch_size
114115
);
115116

116117
println!(
117118
"queue currently_processing_batch_index: {:?}",
118-
merkle_tree.queue_metadata.currently_processing_batch_index as usize
119+
merkle_tree.queue_batches.currently_processing_batch_index as usize
119120
);
120121

121122
println!(
122123
"queue batch_size: {:?}",
123-
merkle_tree.queue_metadata.batch_size
124+
merkle_tree.queue_batches.batch_size
124125
);
125126
println!(
126127
"queue zkp_batch_size: {:?}",
127-
merkle_tree.queue_metadata.zkp_batch_size
128+
merkle_tree.queue_batches.zkp_batch_size
128129
);
129130
println!(
130-
"queue next_full_batch_index: {:?}",
131-
merkle_tree.queue_metadata.next_full_batch_index
131+
"queue pending_batch_index: {:?}",
132+
merkle_tree.queue_batches.pending_batch_index
132133
);
133134
println!(
134135
"queue bloom_filter_capacity: {:?}",
135-
merkle_tree.queue_metadata.bloom_filter_capacity
136+
merkle_tree.queue_batches.bloom_filter_capacity
136137
);
137138
println!(
138139
"queue num_batches: {:?}",
139-
merkle_tree.queue_metadata.num_batches
140+
merkle_tree.queue_batches.num_batches
140141
);
141142

142143
println!("tree next_index: {:?}", merkle_tree.next_index);
@@ -157,36 +158,37 @@ pub async fn get_tree_fullness<R: RpcConnection>(
157158
TreeType::BatchedAddress => {
158159
let mut account = rpc.get_account(tree_pubkey).await?.unwrap();
159160
let merkle_tree =
160-
BatchedMerkleTreeAccount::state_from_bytes(&mut account.data).unwrap();
161+
BatchedMerkleTreeAccount::state_from_bytes(&mut account.data, &tree_pubkey.into())
162+
.unwrap();
161163
println!(
162164
"merkle_tree.get_account().queue.batch_size: {:?}",
163-
merkle_tree.queue_metadata.batch_size
165+
merkle_tree.queue_batches.batch_size
164166
);
165167

166168
println!(
167169
"queue currently_processing_batch_index: {:?}",
168-
merkle_tree.queue_metadata.currently_processing_batch_index as usize
170+
merkle_tree.queue_batches.currently_processing_batch_index as usize
169171
);
170172

171173
println!(
172174
"queue batch_size: {:?}",
173-
merkle_tree.queue_metadata.batch_size
175+
merkle_tree.queue_batches.batch_size
174176
);
175177
println!(
176178
"queue zkp_batch_size: {:?}",
177-
merkle_tree.queue_metadata.zkp_batch_size
179+
merkle_tree.queue_batches.zkp_batch_size
178180
);
179181
println!(
180-
"queue next_full_batch_index: {:?}",
181-
merkle_tree.queue_metadata.next_full_batch_index
182+
"queue pending_batch_index: {:?}",
183+
merkle_tree.queue_batches.pending_batch_index
182184
);
183185
println!(
184186
"queue bloom_filter_capacity: {:?}",
185-
merkle_tree.queue_metadata.bloom_filter_capacity
187+
merkle_tree.queue_batches.bloom_filter_capacity
186188
);
187189
println!(
188190
"queue num_batches: {:?}",
189-
merkle_tree.queue_metadata.num_batches
191+
merkle_tree.queue_batches.num_batches
190192
);
191193

192194
println!("tree next_index: {:?}", merkle_tree.next_index);

‎forester/src/tree_data_sync.rs

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -53,11 +53,11 @@ fn process_address_account(account: &Account, pubkey: Pubkey) -> Result<TreeAcco
5353

5454
fn process_batch_state_account(account: &mut Account, pubkey: Pubkey) -> Result<TreeAccounts> {
5555
let tree_account =
56-
BatchedMerkleTreeAccount::state_from_bytes(&mut account.data).map_err(|e| {
57-
AccountDeserializationError::BatchStateMerkleTree {
56+
BatchedMerkleTreeAccount::state_from_bytes(&mut account.data, &pubkey.into()).map_err(
57+
|e| AccountDeserializationError::BatchStateMerkleTree {
5858
error: e.to_string(),
59-
}
60-
})?;
59+
},
60+
)?;
6161
Ok(create_tree_accounts(
6262
pubkey,
6363
&tree_account.metadata,
@@ -67,11 +67,11 @@ fn process_batch_state_account(account: &mut Account, pubkey: Pubkey) -> Result<
6767

6868
fn process_batch_address_account(account: &mut Account, pubkey: Pubkey) -> Result<TreeAccounts> {
6969
let tree_account =
70-
BatchedMerkleTreeAccount::address_from_bytes(&mut account.data).map_err(|e| {
71-
AccountDeserializationError::BatchAddressMerkleTree {
70+
BatchedMerkleTreeAccount::address_from_bytes(&mut account.data, &pubkey.into()).map_err(
71+
|e| AccountDeserializationError::BatchAddressMerkleTree {
7272
error: e.to_string(),
73-
}
74-
})?;
73+
},
74+
)?;
7575
Ok(create_tree_accounts(
7676
pubkey,
7777
&tree_account.metadata,

0 commit comments

Comments
 (0)