feat: add post accounts in simulation

magicblock-aperture/src/requests/http/simulate_transaction.rs

@@ -1,3 +1,6 @@
+use std::collections::HashMap;
+
+use magicblock_core::link::transactions::TransactionSimulationResult;
 use solana_message::inner_instruction::InnerInstructions;
 use solana_rpc_client_api::{
     config::RpcSimulateTransactionConfig,
@@ -44,17 +47,86 @@ impl HttpDispatcher {
                 debug!(error = ?err, "Failed to prepare transaction to simulate")
             })?;
         self.ensure_transaction_accounts(&transaction.txn).await?;
+        let number_of_accounts = transaction.txn.message().account_keys().len();
 
         let replacement_blockhash = config
             .replace_recent_blockhash
             .then(|| RpcBlockhash::from(self.blocks.get_latest()));
+        let inner_instructions_enabled = config.inner_instructions;
+        let accounts_config = config.accounts;
 
         // Submit the transaction to the scheduler for simulation.
         let result = self
             .transactions_scheduler
             .simulate(transaction.txn)
             .await
             .map_err(RpcError::transaction_simulation)?;
+        let TransactionSimulationResult {
+            result,
+            logs,
+            post_simulation_accounts,
+            units_consumed,
+            return_data,
+            inner_instructions: recorded_inner_instructions,
+        } = result;
+        let result_err = result.as_ref().err().cloned();
+        let accounts = if let Some(config_accounts) = accounts_config {
+            let accounts_encoding = config_accounts
+                .encoding
+                .unwrap_or(UiAccountEncoding::Base64);
+
+            if accounts_encoding == UiAccountEncoding::Binary
+                || accounts_encoding == UiAccountEncoding::Base58
+            {
+                return Err(RpcError::invalid_params(
+                    "base58 encoding not supported",
+                ));
+            }
+
+            if config_accounts.addresses.len() > number_of_accounts {
+                return Err(RpcError::invalid_params(format!(
+                    "Too many accounts provided; max {number_of_accounts}"
+                )));
+            }
+
+            if result_err.is_some() {
+                Some(vec![None; config_accounts.addresses.len()])
+            } else {
+                let pubkeys = config_accounts
+                    .addresses
+                    .into_iter()
+                    .map(|address| {
+                        address
+                            .parse::<Pubkey>()
+                            .map_err(RpcError::invalid_params)
+                    })
+                    .collect::<Result<Vec<_>, _>>()?;
+                let current_accounts =
+                    self.read_accounts_with_ensure(&pubkeys).await;
+                let post_simulation_accounts = post_simulation_accounts
+                    .into_iter()
+                    .collect::<HashMap<_, _>>();
+
+                Some(
+                    pubkeys
+                        .into_iter()
+                        .zip(current_accounts)
+                        .map(|(pubkey, account)| {
+                            post_simulation_accounts
+                                .get(&pubkey)
+                                .cloned()
+                                .or(account)
+                                .map(|account| {
+                                    LockedAccount::new(pubkey, account)
+                                        .ui_encode(accounts_encoding, None)
+                                })
+                        })
+                        .collect(),
+                )
+            }
+        } else {
+            None
+        };
 
         // Convert the internal simulation result to the client-facing RPC format.
         let converter = |(index, ixs): (usize, InnerInstructions)| {
@@ -71,9 +143,8 @@ impl HttpDispatcher {
             .into()
         };
 
-        let inner_instructions = config.inner_instructions.then(|| {
-            result
-                .inner_instructions
+        let inner_instructions = inner_instructions_enabled.then(|| {
+            recorded_inner_instructions
                 .into_iter()
                 .flatten()
                 .enumerate()
@@ -82,11 +153,11 @@ impl HttpDispatcher {
         });
 
         let result = RpcSimulateTransactionResult {
-            err: result.result.err(),
-            logs: result.logs,
-            accounts: None,
-            units_consumed: Some(result.units_consumed),
-            return_data: result.return_data.map(Into::into),
+            err: result_err,
+            logs,
+            accounts,
+            units_consumed: Some(units_consumed),
+            return_data: return_data.map(Into::into),
             inner_instructions,
             replacement_blockhash,
         };

magicblock-aperture/tests/transactions.rs

@@ -4,10 +4,12 @@ use magicblock_accounts_db::traits::AccountsBank;
 use magicblock_core::link::blocks::BlockHash;
 use setup::RpcTestEnv;
 use solana_account::ReadableAccount;
+use solana_account_decoder::UiAccountEncoding;
 use solana_pubkey::Pubkey;
 use solana_rpc_client::rpc_client::GetConfirmedSignaturesForAddress2Config;
 use solana_rpc_client_api::config::{
-    RpcSendTransactionConfig, RpcSimulateTransactionConfig,
+    RpcSendTransactionConfig, RpcSimulateTransactionAccountsConfig,
+    RpcSimulateTransactionConfig,
 };
 use solana_signature::Signature;
 use solana_transaction_status::UiTransactionEncoding;
@@ -207,6 +209,50 @@ async fn test_simulate_transaction_success() {
     );
 }
 
+#[tokio::test]
+async fn test_simulate_transaction_returns_requested_accounts() {
+    let env = RpcTestEnv::new().await;
+    let sender = Pubkey::new_unique();
+    let recipient = Pubkey::new_unique();
+    let transfer_tx =
+        env.build_transfer_txn_with_params(sender, recipient, false);
+    let config = RpcSimulateTransactionConfig {
+        accounts: Some(RpcSimulateTransactionAccountsConfig {
+            encoding: Some(UiAccountEncoding::Base64),
+            addresses: vec![
+                sender.to_string(),
+                recipient.to_string(),
+                guinea::ID.to_string(),
+            ],
+        }),
+        ..Default::default()
+    };
+
+    let result = env
+        .rpc
+        .simulate_transaction_with_config(&transfer_tx, config)
+        .await
+        .expect("simulate_transaction request failed")
+        .value;
+    let accounts = result.accounts.expect("accounts should be returned");
+
+    assert_eq!(accounts.len(), 3, "unexpected account count");
+    assert_eq!(
+        accounts[0].as_ref().map(|account| account.lamports),
+        Some(RpcTestEnv::INIT_ACCOUNT_BALANCE - RpcTestEnv::TRANSFER_AMOUNT),
+        "sender should reflect the simulated transfer",
+    );
+    assert_eq!(
+        accounts[1].as_ref().map(|account| account.lamports),
+        Some(RpcTestEnv::INIT_ACCOUNT_BALANCE + RpcTestEnv::TRANSFER_AMOUNT),
+        "recipient should reflect the simulated transfer",
+    );
+    assert!(
+        accounts[2].is_some(),
+        "program account should be returned for requested addresses"
+    );
+}
+
 /// Tests simulation with config options like replacing blockhash and skipping signature verification.
 #[tokio::test]
 async fn test_simulate_transaction_with_config_options() {

magicblock-core/src/link/transactions.rs

@@ -1,15 +1,19 @@
 use flume::{Receiver as MpmcReceiver, Sender as MpmcSender};
 use magicblock_magic_program_api::args::TaskRequest;
 use serde::Serialize;
+use solana_account::AccountSharedData;
 use solana_program::message::{
     inner_instruction::InnerInstructionsList, SimpleAddressLoader,
 };
+use solana_pubkey::Pubkey;
 use solana_transaction::{
     sanitized::SanitizedTransaction, versioned::VersionedTransaction,
     Transaction,
 };
 use solana_transaction_context::TransactionReturnData;
-use solana_transaction_error::TransactionError;
+use solana_transaction_error::{
+    TransactionError, TransactionResult as SolanaTransactionResult,
+};
 use solana_transaction_status_client_types::TransactionStatusMeta;
 use tokio::sync::{
     mpsc::{Receiver, Sender, UnboundedReceiver, UnboundedSender},
@@ -45,7 +49,7 @@ pub type ScheduledTasksTx = UnboundedSender<TaskRequest>;
 pub struct TransactionSchedulerHandle(pub(super) TransactionToProcessTx);
 
 /// The standard result of a transaction execution, indicating success or a `TransactionError`.
-pub type TransactionResult = solana_transaction_error::TransactionResult<()>;
+pub type TransactionResult = SolanaTransactionResult<()>;
 /// The sender half of a one-shot channel used to return the result of a transaction simulation.
 pub type TxnSimulationResultTx = oneshot::Sender<TransactionSimulationResult>;
 /// An optional sender half of a one-shot channel for returning a transaction execution result.
@@ -114,6 +118,7 @@ pub enum TransactionProcessingMode {
 pub struct TransactionSimulationResult {
     pub result: TransactionResult,
     pub logs: Option<Vec<String>>,
+    pub post_simulation_accounts: Vec<(Pubkey, AccountSharedData)>,
     pub units_consumed: u64,
     pub return_data: Option<TransactionReturnData>,
     pub inner_instructions: Option<InnerInstructionsList>,

magicblock-processor/src/executor/processing.rs

@@ -111,23 +111,42 @@ impl super::TransactionExecutor {
         transaction: [SanitizedTransaction; 1],
         tx: TxnSimulationResultTx,
     ) {
+        let number_of_accounts = transaction[0].message().account_keys().len();
         let (result, _) = self.process(&transaction);
         let simulation_result = match result {
             Ok(processed) => {
                 let status = processed.status();
                 let units_consumed = processed.executed_units();
-                let (logs, return_data, inner_instructions) = match processed {
-                    ProcessedTransaction::Executed(ex) => (
-                        ex.execution_details.log_messages,
-                        ex.execution_details.return_data,
-                        ex.execution_details.inner_instructions,
-                    ),
-                    ProcessedTransaction::FeesOnly(_) => Default::default(),
+                let (
+                    logs,
+                    post_simulation_accounts,
+                    return_data,
+                    inner_instructions,
+                ) = match processed {
+                    ProcessedTransaction::Executed(executed) => {
+                        let execution_details = executed.execution_details;
+                        let post_simulation_accounts = executed
+                            .loaded_transaction
+                            .accounts
+                            .into_iter()
+                            .take(number_of_accounts)
+                            .collect();
+                        (
+                            execution_details.log_messages,
+                            post_simulation_accounts,
+                            execution_details.return_data,
+                            execution_details.inner_instructions,
+                        )
+                    }
+                    ProcessedTransaction::FeesOnly(_) => {
+                        (None, vec![], None, None)
+                    }
                 };
                 TransactionSimulationResult {
                     result: status,
                     units_consumed,
                     logs,
+                    post_simulation_accounts,
                     return_data,
                     inner_instructions,
                 }
@@ -136,6 +155,7 @@ impl super::TransactionExecutor {
                 result: Err(error),
                 units_consumed: 0,
                 logs: Default::default(),
+                post_simulation_accounts: vec![],
                 return_data: None,
                 inner_instructions: None,
             },