From 299676f321052bf5debf138a81940c2726680deb Mon Sep 17 00:00:00 2001
From: Callum <callum.mcintyre@solana.com>
Date: Tue, 5 Mar 2024 19:00:59 +0000
Subject: [PATCH] Return JSON parsed token accounts from simulateBundle

---
 rpc/src/rpc.rs               | 50 ++++++++++++++++++++++++++++++++----
 svm/src/account_overrides.rs |  8 ++++++
 2 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/rpc/src/rpc.rs b/rpc/src/rpc.rs
index db12852d53..a03c2960df 100644
--- a/rpc/src/rpc.rs
+++ b/rpc/src/rpc.rs
@@ -3281,7 +3281,7 @@ pub mod rpc_accounts_scan {
 
 pub mod utils {
     use {
-        crate::rpc::encode_account,
+        crate::rpc::{encode_account, get_encoded_account},
         jsonrpc_core::Error,
         solana_account_decoder::{UiAccount, UiAccountEncoding},
         solana_bundle::{
@@ -3295,7 +3295,10 @@ pub mod utils {
             },
             config::RpcSimulateTransactionAccountsConfig,
         },
+        solana_runtime::bank::Bank,
         solana_sdk::{account::AccountSharedData, pubkey::Pubkey},
+        solana_svm::account_overrides::AccountOverrides,
+        std::collections::HashMap,
         std::str::FromStr,
     };
 
@@ -3304,12 +3307,27 @@ pub mod utils {
     fn try_encode_accounts(
         accounts: &Option<Vec<(Pubkey, AccountSharedData)>>,
         encoding: UiAccountEncoding,
+        bank: &Bank,
+        account_overrides: Option<&AccountOverrides>,
     ) -> Result<Option<Vec<UiAccount>>, Error> {
         if let Some(accounts) = accounts {
+            let overwrite_accounts = account_overrides.map(|a| AccountOverrides::get_accounts(a));
             Ok(Some(
                 accounts
                     .iter()
-                    .map(|(pubkey, account)| encode_account(account, pubkey, encoding, None))
+                    .map(|(pubkey, account)| {
+                        get_encoded_account(
+                            bank,
+                            pubkey,
+                            encoding,
+                            None,
+                            overwrite_accounts.as_ref(),
+                        )
+                        .and_then(|ui_account| match ui_account {
+                            Some(acc) => Ok(acc),
+                            None => encode_account(account, pubkey, encoding, None),
+                        })
+                    })
                     .collect::<Result<Vec<UiAccount>, Error>>()?,
             ))
         } else {
@@ -3320,6 +3338,7 @@ pub mod utils {
     pub fn rpc_bundle_result_from_bank_result(
         bundle_execution_result: LoadAndExecuteBundleOutput,
         rpc_config: RpcSimulateBundleConfig,
+        bank: &Bank,
     ) -> Result<RpcSimulateBundleResult, Error> {
         let summary = match bundle_execution_result.result() {
             Ok(_) => RpcBundleSimulationSummary::Succeeded,
@@ -3341,6 +3360,7 @@ pub mod utils {
         };
 
         let mut transaction_results = Vec::new();
+        let account_overrides: &mut AccountOverrides = &mut AccountOverrides::new(HashMap::new());
         for bundle_output in bundle_execution_result.bundle_transaction_results() {
             for (index, execution_result) in bundle_output
                 .execution_results()
@@ -3364,15 +3384,35 @@ pub mod utils {
                 let pre_execution_accounts = if let Some(pre_tx_accounts) =
                     bundle_output.pre_tx_execution_accounts().get(index)
                 {
-                    try_encode_accounts(pre_tx_accounts, account_encoding)?
+                    try_encode_accounts(
+                        pre_tx_accounts,
+                        account_encoding,
+                        bank,
+                        Some(&account_overrides),
+                    )?
                 } else {
                     None
                 };
 
+                // update account_overrides for this transaction
+                let post_simulation_accounts_map: HashMap<_, _> =
+                    details.post_accounts.clone().into_iter().collect();
+                let account_overrides_transaction =
+                    AccountOverrides::new(post_simulation_accounts_map);
+                AccountOverrides::upsert_account_overrides(
+                    account_overrides,
+                    account_overrides_transaction,
+                );
+
                 let post_execution_accounts = if let Some(post_tx_accounts) =
                     bundle_output.post_tx_execution_accounts().get(index)
                 {
-                    try_encode_accounts(post_tx_accounts, account_encoding)?
+                    try_encode_accounts(
+                        post_tx_accounts,
+                        account_encoding,
+                        bank,
+                        Some(&account_overrides),
+                    )?
                 } else {
                     None
                 };
@@ -4152,7 +4192,7 @@ pub mod rpc_full {
             }
 
             let rpc_bundle_result =
-                rpc_bundle_result_from_bank_result(bundle_execution_result, config)?;
+                rpc_bundle_result_from_bank_result(bundle_execution_result, config, &bank)?;
 
             Ok(new_response(&bank, rpc_bundle_result))
         }
diff --git a/svm/src/account_overrides.rs b/svm/src/account_overrides.rs
index d5d3286426..8ae94d1a06 100644
--- a/svm/src/account_overrides.rs
+++ b/svm/src/account_overrides.rs
@@ -10,6 +10,10 @@ pub struct AccountOverrides {
 }
 
 impl AccountOverrides {
+    pub fn new(accounts: HashMap<Pubkey, AccountSharedData>) -> Self {
+        AccountOverrides { accounts }
+    }
+
     pub fn upsert_account_overrides(&mut self, other: AccountOverrides) {
         self.accounts.extend(other.accounts);
     }
@@ -32,4 +36,8 @@ impl AccountOverrides {
     pub fn get(&self, pubkey: &Pubkey) -> Option<&AccountSharedData> {
         self.accounts.get(pubkey)
     }
+
+    pub fn get_accounts(&self) -> HashMap<Pubkey, AccountSharedData> {
+        self.accounts.clone()
+    }
 }