Skip to content

Commit 7b3dbd8

Browse files
ananas-blockananas-block
committed
refactor: system program relax output tree ordering
1 parent af746de commit 7b3dbd8

File tree

4 files changed

+124
-85
lines changed

4 files changed

+124
-85
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.

‎programs/system/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ pinocchio-pubkey = { workspace = true }
5656
solana-msg = { workspace = true }
5757
light-program-profiler = { workspace = true }
5858
light-heap = { workspace = true, optional = true }
59+
arrayvec = { workspace = true }
5960
[dev-dependencies]
6061
rand = { workspace = true }
6162
light-compressed-account = { workspace = true, features = [

‎programs/system/src/errors.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,10 @@ pub enum SystemProgramError {
140140
PackedAccountIndexOutOfBounds,
141141
#[error("Unimplemented.")]
142142
Unimplemented,
143+
#[error("Too many output V2 queues (max 30).")]
144+
TooManyOutputV2Queues,
145+
#[error("Too many output V1 trees (max 30).")]
146+
TooManyOutputV1Trees,
143147
#[error("Batched Merkle tree error {0}")]
144148
BatchedMerkleTreeError(#[from] BatchedMerkleTreeError),
145149
#[error("Concurrent Merkle tree error {0}")]
@@ -223,6 +227,8 @@ impl From<SystemProgramError> for u32 {
223227
SystemProgramError::Unimplemented => 6063,
224228
SystemProgramError::CpiContextDeactivated => 6064,
225229
SystemProgramError::InputMerkleTreeIndexOutOfBounds => 6065,
230+
SystemProgramError::TooManyOutputV2Queues => 6066,
231+
SystemProgramError::TooManyOutputV1Trees => 6067,
226232
SystemProgramError::BatchedMerkleTreeError(e) => e.into(),
227233
SystemProgramError::IndexedMerkleTreeError(e) => e.into(),
228234
SystemProgramError::ConcurrentMerkleTreeError(e) => e.into(),

‎programs/system/src/processor/create_outputs_cpi_data.rs‎

Lines changed: 116 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use arrayvec::ArrayVec;
12
use light_compressed_account::{
23
hash_to_bn254_field_size_be,
34
instruction_data::{
@@ -42,132 +43,160 @@ pub fn create_outputs_cpi_data<'a, 'info, T: InstructionData<'a>>(
4243
if inputs.output_len() == 0 {
4344
return Ok([0u8; 32]);
4445
}
46+
47+
// Track (tree_pubkey, next_leaf_index) for each unique tree
48+
let mut tree_leaf_indices: ArrayVec<(light_compressed_account::pubkey::Pubkey, u64), 30> =
49+
ArrayVec::new();
50+
51+
#[allow(unused_assignments)]
4552
let mut current_index: i16 = -1;
46-
let mut num_leaves_in_tree: u32 = 0;
47-
let mut mt_next_index: u32 = 0;
4853
let mut hashed_merkle_tree = [0u8; 32];
4954
cpi_ix_data.start_output_appends = context.account_indices.len() as u8;
5055
let mut index_merkle_tree_account_account = cpi_ix_data.start_output_appends;
5156
let mut index_merkle_tree_account = 0;
52-
let number_of_merkle_trees =
53-
inputs.output_accounts().last().unwrap().merkle_tree_index() as usize + 1;
5457

55-
let mut merkle_tree_pubkeys =
56-
Vec::<light_compressed_account::pubkey::Pubkey>::with_capacity(number_of_merkle_trees);
5758
let mut hash_chain = [0u8; 32];
5859
let mut rollover_fee = 0;
5960
let mut is_batched = true;
6061

6162
for (j, account) in inputs.output_accounts().enumerate() {
62-
// if mt index == current index Merkle tree account info has already been added.
63-
// if mt index != current index, Merkle tree account info is new, add it.
64-
#[allow(clippy::comparison_chain)]
65-
if account.merkle_tree_index() as i16 == current_index {
66-
// Do nothing, but it is the most common case.
67-
} else if account.merkle_tree_index() as i16 > current_index {
68-
current_index = account.merkle_tree_index().into();
63+
current_index = account.merkle_tree_index() as i16;
6964

70-
let pubkey = match &accounts
71-
.get(current_index as usize)
72-
.ok_or(SystemProgramError::OutputMerkleTreeIndexOutOfBounds)?
73-
{
74-
AcpAccount::OutputQueue(output_queue) => {
65+
// Get tree/queue pubkey and metadata
66+
let leaf_index = match &accounts
67+
.get(current_index as usize)
68+
.ok_or(SystemProgramError::OutputMerkleTreeIndexOutOfBounds)?
69+
{
70+
AcpAccount::OutputQueue(output_queue) => {
71+
let pubkey = *output_queue.pubkey();
72+
73+
// Check if we've seen this tree before
74+
if let Some(entry) = tree_leaf_indices.iter_mut().find(|(p, _)| *p == pubkey) {
75+
// Tree exists, use and increment its next_leaf_index
76+
let idx = entry.1;
77+
entry.1 += 1;
78+
idx
79+
} else {
80+
// New tree - set up metadata and add to tracking
7581
context.set_network_fee(
7682
output_queue.metadata.rollover_metadata.network_fee,
7783
current_index as u8,
7884
);
79-
8085
hashed_merkle_tree = output_queue.hashed_merkle_tree_pubkey;
8186
rollover_fee = output_queue.metadata.rollover_metadata.rollover_fee;
82-
mt_next_index = output_queue.batch_metadata.next_index as u32;
87+
is_batched = true;
88+
8389
cpi_ix_data.output_sequence_numbers[index_merkle_tree_account as usize] =
8490
MerkleTreeSequenceNumber {
8591
tree_pubkey: output_queue.metadata.associated_merkle_tree,
86-
queue_pubkey: *output_queue.pubkey(),
92+
queue_pubkey: pubkey,
8793
tree_type: (TreeType::StateV2 as u64).into(),
8894
seq: output_queue.batch_metadata.next_index.into(),
8995
};
90-
is_batched = true;
91-
*output_queue.pubkey()
96+
97+
context.get_index_or_insert(
98+
account.merkle_tree_index(),
99+
remaining_accounts,
100+
"Output queue for V2 state trees (Merkle tree for V1 state trees)",
101+
)?;
102+
103+
let next_idx = output_queue.batch_metadata.next_index;
104+
tree_leaf_indices
105+
.try_push((pubkey, next_idx + 1))
106+
.map_err(|_| SystemProgramError::TooManyOutputV2Queues)?;
107+
index_merkle_tree_account += 1;
108+
index_merkle_tree_account_account += 1;
109+
next_idx
92110
}
93-
AcpAccount::StateTree((pubkey, tree)) => {
111+
}
112+
AcpAccount::StateTree((pubkey, tree)) => {
113+
let pubkey = *pubkey;
114+
115+
// Check if we've seen this tree before
116+
if let Some(entry) = tree_leaf_indices.iter_mut().find(|(p, _)| *p == pubkey) {
117+
// Tree exists, use and increment its next_leaf_index
118+
let idx = entry.1;
119+
entry.1 += 1;
120+
idx
121+
} else {
122+
// New tree - set up metadata and add to tracking
94123
cpi_ix_data.output_sequence_numbers[index_merkle_tree_account as usize] =
95124
MerkleTreeSequenceNumber {
96-
tree_pubkey: *pubkey,
97-
queue_pubkey: *pubkey,
125+
tree_pubkey: pubkey,
126+
queue_pubkey: pubkey,
98127
tree_type: (TreeType::StateV1 as u64).into(),
99128
seq: (tree.sequence_number() as u64 + 1).into(),
100129
};
130+
101131
let merkle_context = context
102132
.get_legacy_merkle_context(current_index as u8)
103133
.unwrap();
104134
hashed_merkle_tree = merkle_context.hashed_pubkey;
105135
rollover_fee = merkle_context.rollover_fee;
106-
mt_next_index = tree.next_index() as u32;
107136
is_batched = false;
108-
*pubkey
109-
}
110-
AcpAccount::Unknown() => {
111-
msg!(
112-
format!("found batched unknown create outputs {} ", current_index).as_str()
113-
);
114137

115-
return Err(
116-
SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch.into(),
117-
);
118-
}
119-
AcpAccount::BatchedAddressTree(_) => {
120-
msg!(format!(
121-
"found batched address tree create outputs {} ",
122-
current_index
123-
)
124-
.as_str());
138+
context.get_index_or_insert(
139+
account.merkle_tree_index(),
140+
remaining_accounts,
141+
"Output queue for V2 state trees (Merkle tree for V1 state trees)",
142+
)?;
125143

126-
return Err(
127-
SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch.into(),
128-
);
144+
let next_idx = tree.next_index() as u64;
145+
tree_leaf_indices
146+
.try_push((pubkey, next_idx + 1))
147+
.map_err(|_| SystemProgramError::TooManyOutputV1Trees)?;
148+
index_merkle_tree_account += 1;
149+
index_merkle_tree_account_account += 1;
150+
next_idx
129151
}
130-
AcpAccount::BatchedStateTree(_) => {
131-
msg!(
132-
format!("found batched state tree create outputs {} ", current_index)
133-
.as_str()
134-
);
152+
}
153+
AcpAccount::Unknown() => {
154+
msg!(format!("found batched unknown create outputs {} ", current_index).as_str());
135155

136-
return Err(
137-
SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch.into(),
138-
);
139-
}
140-
_ => {
141-
msg!(format!("create outputs {} ", current_index).as_str());
156+
return Err(SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch.into());
157+
}
158+
AcpAccount::BatchedAddressTree(_) => {
159+
msg!(format!(
160+
"found batched address tree create outputs {} ",
161+
current_index
162+
)
163+
.as_str());
142164

143-
return Err(
144-
SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch.into(),
145-
);
146-
}
147-
};
148-
// check Merkle tree uniqueness
149-
if merkle_tree_pubkeys.contains(&pubkey) {
150-
return Err(SystemProgramError::OutputMerkleTreeNotUnique.into());
151-
} else {
152-
merkle_tree_pubkeys.push(pubkey);
165+
return Err(SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch.into());
153166
}
167+
AcpAccount::BatchedStateTree(_) => {
168+
msg!(
169+
format!("found batched state tree create outputs {} ", current_index).as_str()
170+
);
154171

155-
context.get_index_or_insert(
156-
account.merkle_tree_index(),
157-
remaining_accounts,
158-
"Output queue for V2 state trees (Merkle tree for V1 state trees)",
159-
)?;
160-
num_leaves_in_tree = 0;
161-
index_merkle_tree_account += 1;
162-
index_merkle_tree_account_account += 1;
163-
} else {
164-
// Check 2.
165-
// Output Merkle tree indices must be in order since we use the
166-
// number of leaves in a Merkle tree to determine the correct leaf
167-
// index. Since the leaf index is part of the hash this is security
168-
// critical.
169-
return Err(SystemProgramError::OutputMerkleTreeIndicesNotInOrder.into());
170-
}
172+
return Err(SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch.into());
173+
}
174+
_ => {
175+
msg!(format!("create outputs {} ", current_index).as_str());
176+
177+
return Err(SystemProgramError::StateMerkleTreeAccountDiscriminatorMismatch.into());
178+
}
179+
};
180+
181+
// // Check if this tree already exists in our tracking array
182+
// let leaf_index =
183+
// if let Some(entry) = tree_leaf_indices.iter_mut().find(|(p, _)| *p == pubkey) {
184+
// // Tree exists, use and increment its next_leaf_index
185+
// let idx = entry.1;
186+
// entry.1 += 1;
187+
// idx
188+
// } else {
189+
// context.get_index_or_insert(
190+
// account.merkle_tree_index(),
191+
// remaining_accounts,
192+
// "Output queue for V2 state trees (Merkle tree for V1 state trees)",
193+
// )?;
194+
195+
// tree_leaf_indices.push((pubkey, mt_next_index + 1));
196+
// index_merkle_tree_account += 1;
197+
// index_merkle_tree_account_account += 1;
198+
// mt_next_index
199+
// };
171200

172201
// Check 3.
173202
if let Some(address) = account.address() {
@@ -183,9 +212,11 @@ pub fn create_outputs_cpi_data<'a, 'info, T: InstructionData<'a>>(
183212
return Err(SystemProgramError::InvalidAddress.into());
184213
}
185214
}
186-
cpi_ix_data.output_leaf_indices[j] = (mt_next_index + num_leaves_in_tree).into();
187215

188-
num_leaves_in_tree += 1;
216+
// Use the tracked leaf index from our ArrayVec
217+
cpi_ix_data.output_leaf_indices[j] = u32::try_from(leaf_index)
218+
.map_err(|_| SystemProgramError::PackedAccountIndexOutOfBounds)?
219+
.into();
189220
if account.has_data() && context.invoking_program_id.is_none() {
190221
msg!("Invoking program is not provided.");
191222
msg!("Only program owned compressed accounts can have data.");

0 commit comments

Comments
 (0)