fix e2e test

Author: ananas-block

program-tests/e2e-test/tests/test.rs

@@ -115,8 +115,8 @@ async fn test_batched_only() {
         None,
     )
     .await;
-    // remove concurrent Merkle trees
-    env.indexer.state_merkle_trees.remove(0);
+    // remove the two concurrent Merkle trees
+    env.indexer.state_merkle_trees.drain(..2);
     env.indexer.address_merkle_trees.remove(0);
     println!(
         "address_merkle_trees {:?}",

program-tests/utils/src/e2e_test_env.rs

@@ -2701,11 +2701,36 @@ where
             root_indices.as_slice(),
             &mut remaining_accounts,
         );
-        let output_compressed_accounts = pack_output_compressed_accounts(
+
+        // Pack output accounts and track original creation order
+        let packed_outputs_with_index = pack_output_compressed_accounts(
             output_accounts.as_slice(),
             output_merkle_trees.as_slice(),
             &mut remaining_accounts,
+        )
+        .into_iter()
+        .enumerate()
+        .collect::<Vec<_>>();
+
+        // Sort by merkle_tree_index while tracking original indices
+        let mut sorted_outputs = packed_outputs_with_index.clone();
+        sorted_outputs.sort_by_key(|(_, account)| account.merkle_tree_index);
+
+        // Create mapping: creation_order_index -> sorted_index
+        let mut creation_to_sorted_index = vec![0usize; sorted_outputs.len()];
+        for (sorted_idx, (creation_idx, _)) in sorted_outputs.iter().enumerate() {
+            creation_to_sorted_index[*creation_idx] = sorted_idx;
+        }
+        println!(
+            "creation_to_sorted_index mapping: {:?}",
+            creation_to_sorted_index
         );
+
+        let output_compressed_accounts = sorted_outputs
+            .into_iter()
+            .map(|(_, account)| account)
+            .collect::<Vec<_>>();
+
         println!(
             "output_compressed_accounts: {:?}",
             output_compressed_accounts
@@ -2734,13 +2759,15 @@ where
         );
         let mut new_address_params = Vec::new();
         for (index, new_address_param) in invoke_cpi.new_address_params.into_iter().enumerate() {
-            if let Some((_, account_index)) = new_address_indices
+            if let Some((_, creation_order_idx)) = new_address_indices
                 .iter()
                 .find(|(address_index, _)| *address_index == index)
             {
+                // Map creation order index to sorted index
+                let sorted_idx = creation_to_sorted_index[*creation_order_idx as usize];
                 new_address_params.push(NewAddressParamsAssignedPacked::new(
                     new_address_param,
-                    Some(*account_index),
+                    Some(sorted_idx.try_into().unwrap()),
                 ));
             } else {
                 new_address_params

sdk-libs/program-test/src/logging/decoder.rs

@@ -255,6 +255,7 @@ fn parse_invoke_instruction(data: &[u8], accounts: &[AccountMeta]) -> Instructio
                     merkle_tree_index: Some(param.address_merkle_tree_account_index),
                     root_index: Some(param.address_merkle_tree_root_index),
                     derived_address: None,
+                    assigned_account_index: super::types::AddressAssignment::V1,
                 })
                 .collect(),
         )
@@ -399,6 +400,7 @@ fn parse_invoke_cpi_instruction(data: &[u8], accounts: &[AccountMeta]) -> Instru
                     merkle_tree_index: Some(param.address_merkle_tree_account_index),
                     root_index: Some(param.address_merkle_tree_root_index),
                     derived_address: None,
+                    assigned_account_index: super::types::AddressAssignment::V1,
                 })
                 .collect(),
         )
@@ -538,6 +540,11 @@ fn parse_invoke_cpi_readonly_instruction(
             merkle_tree_index: Some(param.address_merkle_tree_account_index),
             root_index: Some(param.address_merkle_tree_root_index),
             derived_address: None,
+            assigned_account_index: if param.assigned_to_account {
+                super::types::AddressAssignment::AssignedIndex(param.assigned_account_index)
+            } else {
+                super::types::AddressAssignment::None
+            },
         });
     }
 
@@ -549,6 +556,7 @@ fn parse_invoke_cpi_readonly_instruction(
             merkle_tree_index: Some(readonly_addr.address_merkle_tree_account_index),
             root_index: Some(readonly_addr.address_merkle_tree_root_index),
             derived_address: Some(readonly_addr.address),
+            assigned_account_index: super::types::AddressAssignment::None,
         });
     }
 
@@ -684,6 +692,11 @@ fn parse_invoke_cpi_account_info_instruction(
             merkle_tree_index: Some(param.address_merkle_tree_account_index),
             root_index: Some(param.address_merkle_tree_root_index),
             derived_address: None,
+            assigned_account_index: if param.assigned_to_account {
+                super::types::AddressAssignment::AssignedIndex(param.assigned_account_index)
+            } else {
+                super::types::AddressAssignment::None
+            },
         });
     }
 
@@ -695,6 +708,7 @@ fn parse_invoke_cpi_account_info_instruction(
             merkle_tree_index: Some(readonly_addr.address_merkle_tree_account_index),
             root_index: Some(readonly_addr.address_merkle_tree_root_index),
             derived_address: Some(readonly_addr.address),
+            assigned_account_index: super::types::AddressAssignment::None,
         });
     }
 

sdk-libs/program-test/src/logging/formatter.rs

@@ -922,6 +922,26 @@ impl TransactionFormatter {
                                 self.colors.reset
                             )?;
                         }
+                        let assignment_str = match addr_param.assigned_account_index {
+                            super::types::AddressAssignment::AssignedIndex(idx) => {
+                                format!("{}", idx)
+                            }
+                            super::types::AddressAssignment::None => {
+                                "none".to_string()
+                            }
+                            super::types::AddressAssignment::V1 => {
+                                "n/a (v1)".to_string()
+                            }
+                        };
+                        writeln!(
+                            output,
+                            "{}      {}assigned_to_output_account: {}{}{}",
+                            indent,
+                            self.colors.gray,
+                            self.colors.yellow,
+                            assignment_str,
+                            self.colors.reset
+                        )?;
                     }
                 }
 

sdk-libs/program-test/src/logging/types.rs

@@ -106,6 +106,17 @@ pub struct FeeSummary {
     pub compression_fee: Option<u64>,
 }
 
+/// Address assignment state
+#[derive(Debug, Clone)]
+pub enum AddressAssignment {
+    /// V1 address param (no assignment tracking)
+    V1,
+    /// Not assigned to any output account
+    None,
+    /// Assigned to output account at index
+    AssignedIndex(u8),
+}
+
 /// Address parameter information
 #[derive(Debug, Clone)]
 pub struct AddressParam {
@@ -114,6 +125,7 @@ pub struct AddressParam {
     pub merkle_tree_index: Option<u8>,
     pub root_index: Option<u16>,
     pub derived_address: Option<[u8; 32]>,
+    pub assigned_account_index: AddressAssignment,
 }
 
 /// Input account data