Refactor: Flattens TransactionContext::instruction_trace

program-runtime/src/invoke_context.rs

@@ -325,7 +325,7 @@ impl<'a> InvokeContext<'a> {
                 .get_instruction_context_stack_height())
                 .any(|level| {
                     self.transaction_context
-                        .get_instruction_context_at(level)
+                        .get_instruction_context_at_nesting_level(level)
                         .and_then(|instruction_context| {
                             instruction_context
                                 .try_borrow_last_program_account(self.transaction_context)

programs/bpf_loader/src/serialization.rs

@@ -497,21 +497,20 @@ mod tests {
                 &program_indices,
             )
             .instruction_accounts;
+            let instruction_data = vec![];
 
-            let transaction_context =
+            let mut transaction_context =
                 TransactionContext::new(transaction_accounts, Some(Rent::default()), 1, 1);
-            let instruction_data = vec![];
-            let instruction_context = InstructionContext::new(
-                0,
-                0,
-                &program_indices,
-                &instruction_accounts,
-                &instruction_data,
-            );
+            transaction_context
+                .push(&program_indices, &instruction_accounts, &instruction_data)
+                .unwrap();
+            let instruction_context = transaction_context
+                .get_instruction_context_at_index_in_trace(0)
+                .unwrap();
 
             let serialization_result = serialize_parameters(
                 &transaction_context,
-                &instruction_context,
+                instruction_context,
                 should_cap_ix_accounts,
             );
             assert_eq!(

programs/bpf_loader/src/syscalls/mod.rs

@@ -1709,34 +1709,37 @@ declare_syscall!(
             result
         );
 
+        // Reverse iterate through the instruction trace,
+        // ignoring anything except instructions on the same level
         let stack_height = invoke_context.get_stack_height();
-        let instruction_trace = invoke_context.transaction_context.get_instruction_trace();
-        let instruction_context = if stack_height == TRANSACTION_LEVEL_STACK_HEIGHT {
-            // pick one of the top-level instructions
-            instruction_trace
-                .len()
-                .checked_sub(2)
-                .and_then(|result| result.checked_sub(index as usize))
-                .and_then(|index| instruction_trace.get(index))
-                .and_then(|instruction_list| instruction_list.first())
-        } else {
-            // Walk the last list of inner instructions
-            instruction_trace.last().and_then(|inners| {
-                let mut current_index = 0;
-                inners.iter().rev().skip(1).find(|instruction_context| {
-                    if stack_height == instruction_context.get_stack_height() {
-                        if index == current_index {
-                            return true;
-                        } else {
-                            current_index = current_index.saturating_add(1);
-                        }
-                    }
-                    false
-                })
-            })
-        };
+        let instruction_trace_length = invoke_context
+            .transaction_context
+            .get_instruction_trace_length();
+        let mut reverse_index_at_stack_height = 0;
+        let mut found_instruction_context = None;
+        for index_in_trace in (0..instruction_trace_length).rev() {
+            let instruction_context = question_mark!(
+                invoke_context
+                    .transaction_context
+                    .get_instruction_context_at_index_in_trace(index_in_trace)
+                    .map_err(SyscallError::InstructionError),
+                result
+            );
+            if instruction_context.get_stack_height() == TRANSACTION_LEVEL_STACK_HEIGHT
+                && stack_height > TRANSACTION_LEVEL_STACK_HEIGHT
+            {
+                break;
+            }
+            if instruction_context.get_stack_height() == stack_height {
+                if index.saturating_add(1) == reverse_index_at_stack_height {
+                    found_instruction_context = Some(instruction_context);
+                    break;
+                }
+                reverse_index_at_stack_height = reverse_index_at_stack_height.saturating_add(1);
+            }
+        }
 
-        if let Some(instruction_context) = instruction_context {
+        if let Some(instruction_context) = found_instruction_context {
             let ProcessedSiblingInstruction {
                 data_len,
                 accounts_len,

runtime/src/bank.rs

@@ -120,7 +120,7 @@ use {
         hash::{extend_and_hash, hashv, Hash},
         incinerator,
         inflation::Inflation,
-        instruction::CompiledInstruction,
+        instruction::{CompiledInstruction, TRANSACTION_LEVEL_STACK_HEIGHT},
         lamports::LamportsError,
         message::{AccountKeys, SanitizedMessage},
         native_loader,
@@ -143,8 +143,7 @@ use {
             TransactionVerificationMode, VersionedTransaction, MAX_TX_ACCOUNT_LOCKS,
         },
         transaction_context::{
-            ExecutionRecord, InstructionTrace, TransactionAccount, TransactionContext,
-            TransactionReturnData,
+            ExecutionRecord, TransactionAccount, TransactionContext, TransactionReturnData,
         },
     },
     solana_stake_program::stake_state::{
@@ -762,40 +761,50 @@ pub type InnerInstructions = Vec<CompiledInstruction>;
 /// a transaction
 pub type InnerInstructionsList = Vec<InnerInstructions>;
 
-/// Convert from an InstructionTrace to InnerInstructionsList
+/// Extract the InnerInstructionsList from a TransactionContext
 pub fn inner_instructions_list_from_instruction_trace(
-    instruction_trace: &InstructionTrace,
+    transaction_context: &TransactionContext,
 ) -> InnerInstructionsList {
-    instruction_trace
-        .iter()
-        .map(|inner_instructions_trace| {
-            inner_instructions_trace
-                .iter()
-                .skip(1)
-                .map(|instruction_context| {
-                    CompiledInstruction::new_from_raw_parts(
-                        instruction_context
-                            .get_index_of_program_account_in_transaction(
-                                instruction_context
-                                    .get_number_of_program_accounts()
-                                    .saturating_sub(1),
-                            )
-                            .unwrap_or_default() as u8,
-                        instruction_context.get_instruction_data().to_vec(),
-                        (0..instruction_context.get_number_of_instruction_accounts())
-                            .map(|instruction_account_index| {
-                                instruction_context
-                                    .get_index_of_instruction_account_in_transaction(
-                                        instruction_account_index,
-                                    )
-                                    .unwrap_or_default() as u8
-                            })
-                            .collect(),
-                    )
-                })
-                .collect()
-        })
-        .collect()
+    debug_assert!(transaction_context
+        .get_instruction_context_at_index_in_trace(0)
+        .map(|instruction_context| instruction_context.get_stack_height()
+            == TRANSACTION_LEVEL_STACK_HEIGHT)
+        .unwrap_or(true));
+    let mut outer_instructions = Vec::new();
+    for index_in_trace in 0..transaction_context.get_instruction_trace_length() {
+        if let Ok(instruction_context) =
+            transaction_context.get_instruction_context_at_index_in_trace(index_in_trace)
+        {
+            if instruction_context.get_stack_height() == TRANSACTION_LEVEL_STACK_HEIGHT {
+                outer_instructions.push(Vec::new());
+            } else if let Some(inner_instructions) = outer_instructions.last_mut() {
+                inner_instructions.push(CompiledInstruction::new_from_raw_parts(
+                    instruction_context
+                        .get_index_of_program_account_in_transaction(
+                            instruction_context
+                                .get_number_of_program_accounts()
+                                .saturating_sub(1),
+                        )
+                        .unwrap_or_default() as u8,
+                    instruction_context.get_instruction_data().to_vec(),
+                    (0..instruction_context.get_number_of_instruction_accounts())
+                        .map(|instruction_account_index| {
+                            instruction_context
+                                .get_index_of_instruction_account_in_transaction(
+                                    instruction_account_index,
+                                )
+                                .unwrap_or_default() as u8
+                        })
+                        .collect(),
+                ));
+            } else {
+                debug_assert!(false);
+            }
+        } else {
+            debug_assert!(false);
+        }
+    }
+    outer_instructions
 }
 
 /// A list of log messages emitted during a transaction
@@ -4430,9 +4439,16 @@ impl Bank {
                     .ok()
             });
 
+        let inner_instructions = if enable_cpi_recording {
+            Some(inner_instructions_list_from_instruction_trace(
+                &transaction_context,
+            ))
+        } else {
+            None
+        };
+
         let ExecutionRecord {
             accounts,
-            instruction_trace,
             mut return_data,
             changed_account_count,
             total_size_of_all_accounts,
@@ -4460,14 +4476,6 @@ impl Bank {
             accounts_data_len_delta = status.as_ref().map_or(0, |_| accounts_resize_delta);
         }
 
-        let inner_instructions = if enable_cpi_recording {
-            Some(inner_instructions_list_from_instruction_trace(
-                &instruction_trace,
-            ))
-        } else {
-            None
-        };
-
         let return_data = if enable_return_data_recording {
             if let Some(end_index) = return_data.data.iter().rposition(|&x| x != 0) {
                 let end_index = end_index.saturating_add(1);
@@ -8045,7 +8053,6 @@ pub(crate) mod tests {
             system_program,
             timing::duration_as_s,
             transaction::MAX_TX_ACCOUNT_LOCKS,
-            transaction_context::InstructionContext,
         },
         solana_vote_program::{
             vote_instruction,
@@ -18837,26 +18844,24 @@ pub(crate) mod tests {
 
     #[test]
     fn test_inner_instructions_list_from_instruction_trace() {
-        let instruction_trace = vec![
-            vec![
-                InstructionContext::new(0, 0, &[], &[], &[1]),
-                InstructionContext::new(1, 0, &[], &[], &[2]),
-            ],
-            vec![],
-            vec![
-                InstructionContext::new(0, 0, &[], &[], &[3]),
-                InstructionContext::new(1, 0, &[], &[], &[4]),
-                InstructionContext::new(2, 0, &[], &[], &[5]),
-                InstructionContext::new(1, 0, &[], &[], &[6]),
-            ],
-        ];
-
-        let inner_instructions = inner_instructions_list_from_instruction_trace(&instruction_trace);
+        let mut transaction_context = TransactionContext::new(vec![], None, 3, 3);
+        for (index_in_trace, stack_height) in [1, 2, 1, 1, 2, 3, 2].into_iter().enumerate() {
+            while stack_height <= transaction_context.get_instruction_context_stack_height() {
+                transaction_context.pop().unwrap();
+            }
+            if stack_height > transaction_context.get_instruction_context_stack_height() {
+                transaction_context
+                    .push(&[], &[], &[index_in_trace as u8])
+                    .unwrap();
+            }
+        }
+        let inner_instructions =
+            inner_instructions_list_from_instruction_trace(&transaction_context);
 
         assert_eq!(
             inner_instructions,
             vec![
-                vec![CompiledInstruction::new_from_raw_parts(0, vec![2], vec![])],
+                vec![CompiledInstruction::new_from_raw_parts(0, vec![1], vec![])],
                 vec![],
                 vec![
                     CompiledInstruction::new_from_raw_parts(0, vec![4], vec![]),

runtime/src/message_processor.rs

@@ -680,6 +680,6 @@ mod tests {
                 InstructionError::Custom(0xbabb1e)
             ))
         );
-        assert_eq!(transaction_context.get_instruction_trace().len(), 2);
+        assert_eq!(transaction_context.get_instruction_trace_length(), 2);
     }
 }