surfpool_core/surfnet/
locker.rs

1use std::{collections::BTreeMap, sync::Arc};
2
3use bincode::serialized_size;
4use crossbeam_channel::{Receiver, Sender};
5use itertools::Itertools;
6use jsonrpc_core::futures::future::join_all;
7use litesvm::types::{FailedTransactionMetadata, SimulatedTransactionInfo, TransactionResult};
8use solana_account::Account;
9use solana_account_decoder::{
10    encode_ui_account,
11    parse_account_data::AccountAdditionalDataV3,
12    parse_bpf_loader::{parse_bpf_upgradeable_loader, BpfUpgradeableLoaderAccountType, UiProgram},
13    parse_token::UiTokenAmount,
14    UiAccount, UiAccountEncoding,
15};
16use solana_address_lookup_table_interface::state::AddressLookupTable;
17use solana_client::{
18    rpc_config::{
19        RpcAccountInfoConfig, RpcLargestAccountsConfig, RpcLargestAccountsFilter,
20        RpcSignaturesForAddressConfig,
21    },
22    rpc_filter::RpcFilterType,
23    rpc_request::TokenAccountsFilter,
24    rpc_response::{
25        RpcAccountBalance, RpcConfirmedTransactionStatusWithSignature, RpcKeyedAccount,
26        RpcTokenAccountBalance,
27    },
28};
29use solana_clock::Slot;
30use solana_commitment_config::{CommitmentConfig, CommitmentLevel};
31use solana_epoch_info::EpochInfo;
32use solana_hash::Hash;
33use solana_message::{
34    v0::{LoadedAddresses, MessageAddressTableLookup},
35    VersionedMessage,
36};
37use solana_pubkey::Pubkey;
38use solana_rpc_client_api::response::SlotInfo;
39use solana_sdk::{
40    bpf_loader_upgradeable::{get_program_data_address, UpgradeableLoaderState},
41    program_pack::Pack,
42    transaction::VersionedTransaction,
43};
44use solana_signature::Signature;
45use solana_transaction_error::TransactionError;
46use solana_transaction_status::{
47    EncodedConfirmedTransactionWithStatusMeta,
48    TransactionConfirmationStatus as SolanaTransactionConfirmationStatus, UiTransactionEncoding,
49};
50use spl_token::state::Mint;
51use spl_token_2022::extension::StateWithExtensions;
52use surfpool_types::{
53    ComputeUnitsEstimationResult, ProfileResult, ProfileState, SimnetEvent,
54    TransactionConfirmationStatus, TransactionStatusEvent,
55};
56use tokio::sync::RwLock;
57
58use super::{
59    remote::{SomeRemoteCtx, SurfnetRemoteClient},
60    AccountFactory, GetAccountResult, GetTransactionResult, GeyserEvent, SignatureSubscriptionType,
61    SurfnetSvm,
62};
63use crate::{
64    error::{SurfpoolError, SurfpoolResult},
65    rpc::utils::{convert_transaction_metadata_from_canonical, verify_pubkey},
66    surfnet::FINALIZATION_SLOT_THRESHOLD,
67    types::TransactionWithStatusMeta,
68};
69
70pub struct SvmAccessContext<T> {
71    pub slot: Slot,
72    pub latest_epoch_info: EpochInfo,
73    pub latest_blockhash: Hash,
74    pub inner: T,
75}
76
77impl<T> SvmAccessContext<T> {
78    pub fn new(slot: Slot, latest_epoch_info: EpochInfo, latest_blockhash: Hash, inner: T) -> Self {
79        Self {
80            slot,
81            latest_blockhash,
82            latest_epoch_info,
83            inner,
84        }
85    }
86
87    pub fn inner(&self) -> &T {
88        &self.inner
89    }
90
91    pub fn with_new_value<N>(&self, inner: N) -> SvmAccessContext<N> {
92        SvmAccessContext {
93            slot: self.slot,
94            latest_blockhash: self.latest_blockhash,
95            latest_epoch_info: self.latest_epoch_info.clone(),
96            inner,
97        }
98    }
99}
100
101pub type SurfpoolContextualizedResult<T> = SurfpoolResult<SvmAccessContext<T>>;
102
103pub struct SurfnetSvmLocker(pub Arc<RwLock<SurfnetSvm>>);
104
105impl Clone for SurfnetSvmLocker {
106    fn clone(&self) -> Self {
107        Self(self.0.clone())
108    }
109}
110
111/// Functions for reading and writing to the underlying SurfnetSvm instance
112impl SurfnetSvmLocker {
113    /// Executes a read-only operation on the underlying `SurfnetSvm` by acquiring a blocking read lock.
114    /// Accepts a closure that receives a shared reference to `SurfnetSvm` and returns a value.
115    ///
116    /// # Returns
117    /// The result produced by the closure.
118    pub fn with_svm_reader<T, F>(&self, reader: F) -> T
119    where
120        F: Fn(&SurfnetSvm) -> T + Send + Sync,
121        T: Send + 'static,
122    {
123        let read_lock = self.0.clone();
124        tokio::task::block_in_place(move || {
125            let read_guard = read_lock.blocking_read();
126            reader(&read_guard)
127        })
128    }
129
130    /// Executes a read-only operation and wraps the result in `SvmAccessContext`, capturing
131    /// slot, epoch info, and blockhash along with the closure's result.
132    fn with_contextualized_svm_reader<T, F>(&self, reader: F) -> SvmAccessContext<T>
133    where
134        F: Fn(&SurfnetSvm) -> T + Send + Sync,
135        T: Send + 'static,
136    {
137        let read_lock = self.0.clone();
138        tokio::task::block_in_place(move || {
139            let read_guard = read_lock.blocking_read();
140            let res = reader(&read_guard);
141
142            SvmAccessContext::new(
143                read_guard.get_latest_absolute_slot(),
144                read_guard.latest_epoch_info(),
145                read_guard.latest_blockhash(),
146                res,
147            )
148        })
149    }
150
151    /// Executes a write operation on the underlying `SurfnetSvm` by acquiring a blocking write lock.
152    /// Accepts a closure that receives a mutable reference to `SurfnetSvm` and returns a value.
153    ///
154    /// # Returns
155    /// The result produced by the closure.
156    pub fn with_svm_writer<T, F>(&self, writer: F) -> T
157    where
158        F: Fn(&mut SurfnetSvm) -> T + Send + Sync,
159        T: Send + 'static,
160    {
161        let write_lock = self.0.clone();
162        tokio::task::block_in_place(move || {
163            let mut write_guard = write_lock.blocking_write();
164            writer(&mut write_guard)
165        })
166    }
167}
168
169/// Functions for creating and initializing the underlying SurfnetSvm instance
170impl SurfnetSvmLocker {
171    /// Constructs a new `SurfnetSvmLocker` wrapping the given `SurfnetSvm` instance.
172    pub fn new(svm: SurfnetSvm) -> Self {
173        Self(Arc::new(RwLock::new(svm)))
174    }
175
176    /// Initializes the locked `SurfnetSvm` by fetching or defaulting epoch info,
177    /// then calling its `initialize` method. Returns the epoch info on success.
178    pub async fn initialize(
179        &self,
180        remote_ctx: &Option<SurfnetRemoteClient>,
181    ) -> SurfpoolResult<EpochInfo> {
182        let epoch_info = if let Some(remote_client) = remote_ctx {
183            remote_client.get_epoch_info().await?
184        } else {
185            EpochInfo {
186                epoch: 0,
187                slot_index: 0,
188                slots_in_epoch: 0,
189                absolute_slot: 0,
190                block_height: 0,
191                transaction_count: None,
192            }
193        };
194
195        self.with_svm_writer(|svm_writer| {
196            svm_writer.initialize(epoch_info.clone(), remote_ctx);
197        });
198        Ok(epoch_info)
199    }
200}
201
202/// Functions for getting accounts from the underlying SurfnetSvm instance or remote client
203impl SurfnetSvmLocker {
204    /// Retrieves a local account from the SVM cache, returning a contextualized result.
205    pub fn get_account_local(&self, pubkey: &Pubkey) -> SvmAccessContext<GetAccountResult> {
206        self.with_contextualized_svm_reader(|svm_reader| {
207            match svm_reader.inner.get_account(pubkey) {
208                Some(account) => GetAccountResult::FoundAccount(
209                    *pubkey, account,
210                    // mark as not an account that should be updated in the SVM, since this is a local read and it already exists
211                    false,
212                ),
213                None => GetAccountResult::None(*pubkey),
214            }
215        })
216    }
217
218    /// Attempts local retrieval, then fetches from remote if missing, returning a contextualized result.
219    pub async fn get_account_local_then_remote(
220        &self,
221        client: &SurfnetRemoteClient,
222        pubkey: &Pubkey,
223        commitment_config: CommitmentConfig,
224    ) -> SurfpoolContextualizedResult<GetAccountResult> {
225        let result = self.get_account_local(pubkey);
226
227        if result.inner.is_none() {
228            let remote_account = client.get_account(pubkey, commitment_config).await?;
229            Ok(result.with_new_value(remote_account))
230        } else {
231            Ok(result)
232        }
233    }
234
235    /// Retrieves an account, using local or remote based on context, applying a default factory if provided.
236    pub async fn get_account(
237        &self,
238        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
239        pubkey: &Pubkey,
240        factory: Option<AccountFactory>,
241    ) -> SurfpoolContextualizedResult<GetAccountResult> {
242        let result = if let Some((remote_client, commitment_config)) = remote_ctx {
243            self.get_account_local_then_remote(remote_client, pubkey, *commitment_config)
244                .await?
245        } else {
246            self.get_account_local(pubkey)
247        };
248
249        match (&result.inner, factory) {
250            (&GetAccountResult::None(_), Some(factory)) => {
251                let default = factory(self.clone());
252                Ok(result.with_new_value(default))
253            }
254            _ => Ok(result),
255        }
256    }
257
258    /// Retrieves an account, using local or remote based on context, applying a default factory if provided.
259    pub fn get_local_account_associated_data(
260        &self,
261        account: &GetAccountResult,
262    ) -> SvmAccessContext<Option<AccountAdditionalDataV3>> {
263        self.with_contextualized_svm_reader(|svm_reader| {
264            let associated_data = match account {
265                GetAccountResult::FoundAccount(_, account, _) => {
266                    if !account.owner.eq(&spl_token_2022::id()) {
267                        return None;
268                    }
269
270                    let Ok(token_data) =
271                        StateWithExtensions::<spl_token_2022::state::Account>::unpack(
272                            &account.data,
273                        )
274                    else {
275                        return None;
276                    };
277                    svm_reader
278                        .account_associated_data
279                        .get(&token_data.base.mint)
280                        .copied()
281                }
282                _ => None,
283            };
284            associated_data
285        })
286    }
287
288    /// Retrieves multiple accounts from local cache, returning a contextualized result.
289    pub fn get_multiple_accounts_local(
290        &self,
291        pubkeys: &[Pubkey],
292    ) -> SvmAccessContext<Vec<GetAccountResult>> {
293        self.with_contextualized_svm_reader(|svm_reader| {
294            let mut accounts = vec![];
295
296            for pubkey in pubkeys.iter() {
297                let res = match svm_reader.inner.get_account(pubkey) {
298                    Some(account) => GetAccountResult::FoundAccount(
299                        *pubkey, account,
300                        // mark as not an account that should be updated in the SVM, since this is a local read and it already exists
301                        false,
302                    ),
303                    None => GetAccountResult::None(*pubkey),
304                };
305                accounts.push(res);
306            }
307            accounts
308        })
309    }
310
311    /// Retrieves multiple accounts, fetching missing ones from remote, returning a contextualized result.
312    pub async fn get_multiple_accounts_local_then_remote(
313        &self,
314        client: &SurfnetRemoteClient,
315        pubkeys: &[Pubkey],
316        commitment_config: CommitmentConfig,
317    ) -> SurfpoolContextualizedResult<Vec<GetAccountResult>> {
318        let SvmAccessContext {
319            slot,
320            latest_epoch_info,
321            latest_blockhash,
322            inner: local_results,
323        } = self.get_multiple_accounts_local(pubkeys);
324
325        let mut missing_accounts = vec![];
326        let mut found_accounts = vec![];
327        for result in local_results.into_iter() {
328            if let GetAccountResult::None(pubkey) = result {
329                missing_accounts.push(pubkey)
330            } else {
331                found_accounts.push(result.clone());
332            }
333        }
334
335        if missing_accounts.is_empty() {
336            return Ok(SvmAccessContext::new(
337                slot,
338                latest_epoch_info,
339                latest_blockhash,
340                found_accounts,
341            ));
342        }
343
344        let mut remote_results = client
345            .get_multiple_accounts(&missing_accounts, commitment_config)
346            .await?;
347        let mut combined_results = found_accounts.clone();
348        combined_results.append(&mut remote_results);
349
350        Ok(SvmAccessContext::new(
351            slot,
352            latest_epoch_info,
353            latest_blockhash,
354            combined_results,
355        ))
356    }
357
358    /// Retrieves multiple accounts, using local or remote context and applying factory defaults if provided.
359    pub async fn get_multiple_accounts(
360        &self,
361        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
362        pubkeys: &[Pubkey],
363        factory: Option<AccountFactory>,
364    ) -> SurfpoolContextualizedResult<Vec<GetAccountResult>> {
365        let results = if let Some((remote_client, commitment_config)) = remote_ctx {
366            self.get_multiple_accounts_local_then_remote(remote_client, pubkeys, *commitment_config)
367                .await?
368        } else {
369            self.get_multiple_accounts_local(pubkeys)
370        };
371
372        let mut combined = Vec::with_capacity(results.inner.len());
373        for result in results.inner.clone() {
374            match (&result, &factory) {
375                (&GetAccountResult::None(_), Some(factory)) => {
376                    let default = factory(self.clone());
377                    combined.push(default);
378                }
379                _ => combined.push(result),
380            }
381        }
382        Ok(results.with_new_value(combined))
383    }
384
385    /// Retrieves largest accounts from local cache, returning a contextualized result.
386    pub fn get_largest_accounts_local(
387        &self,
388        config: RpcLargestAccountsConfig,
389    ) -> SvmAccessContext<Vec<RpcAccountBalance>> {
390        self.with_contextualized_svm_reader(|svm_reader| {
391            let non_circulating_accounts: Vec<_> = svm_reader
392                .non_circulating_accounts
393                .iter()
394                .flat_map(|acct| verify_pubkey(acct))
395                .collect();
396
397            let ordered_accounts = svm_reader
398                .accounts_registry
399                .iter()
400                .sorted_by(|a, b| b.1.lamports.cmp(&a.1.lamports))
401                .collect::<Vec<_>>();
402
403            let ordered_filtered_accounts = match config.filter {
404                Some(RpcLargestAccountsFilter::NonCirculating) => ordered_accounts
405                    .into_iter()
406                    .filter(|(pubkey, _)| non_circulating_accounts.contains(pubkey))
407                    .collect::<Vec<_>>(),
408                Some(RpcLargestAccountsFilter::Circulating) => ordered_accounts
409                    .into_iter()
410                    .filter(|(pubkey, _)| !non_circulating_accounts.contains(pubkey))
411                    .collect::<Vec<_>>(),
412                None => ordered_accounts,
413            };
414
415            ordered_filtered_accounts
416                .iter()
417                .take(20)
418                .map(|(pubkey, account)| RpcAccountBalance {
419                    address: pubkey.to_string(),
420                    lamports: account.lamports,
421                })
422                .collect()
423        })
424    }
425
426    pub async fn get_largest_accounts_local_then_remote(
427        &self,
428        client: &SurfnetRemoteClient,
429        config: RpcLargestAccountsConfig,
430        commitment_config: CommitmentConfig,
431    ) -> SurfpoolContextualizedResult<Vec<RpcAccountBalance>> {
432        // get all non-circulating and circulating pubkeys from the remote client first,
433        // and insert them locally
434        {
435            let mut remote_non_circulating_pubkeys = client
436                .get_largest_accounts(Some(RpcLargestAccountsConfig {
437                    filter: Some(RpcLargestAccountsFilter::NonCirculating),
438                    ..config.clone()
439                }))
440                .await?
441                .iter()
442                .map(|account_balance| verify_pubkey(&account_balance.address))
443                .collect::<SurfpoolResult<Vec<_>>>()?;
444
445            let mut remote_circulating_pubkeys = client
446                .get_largest_accounts(Some(RpcLargestAccountsConfig {
447                    filter: Some(RpcLargestAccountsFilter::Circulating),
448                    ..config.clone()
449                }))
450                .await?
451                .iter()
452                .map(|account_balance| verify_pubkey(&account_balance.address))
453                .collect::<SurfpoolResult<Vec<_>>>()?;
454
455            let mut combined = Vec::with_capacity(
456                remote_non_circulating_pubkeys.len() + remote_circulating_pubkeys.len(),
457            );
458            combined.append(&mut remote_non_circulating_pubkeys);
459            combined.append(&mut remote_circulating_pubkeys);
460
461            let get_account_results = self
462                .get_multiple_accounts_local_then_remote(client, &combined, commitment_config)
463                .await?
464                .inner;
465
466            self.write_multiple_account_updates(&get_account_results);
467        }
468
469        // now that our local cache is aware of all large remote accounts, we can get the largest accounts locally
470        // and filter according to the config
471        Ok(self.get_largest_accounts_local(config))
472    }
473
474    pub async fn get_largest_accounts(
475        &self,
476        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
477        config: RpcLargestAccountsConfig,
478    ) -> SurfpoolContextualizedResult<Vec<RpcAccountBalance>> {
479        let results = if let Some((remote_client, commitment_config)) = remote_ctx {
480            self.get_largest_accounts_local_then_remote(remote_client, config, *commitment_config)
481                .await?
482        } else {
483            self.get_largest_accounts_local(config)
484        };
485
486        Ok(results)
487    }
488}
489
490/// Get signatures for Addresses
491impl SurfnetSvmLocker {
492    pub fn get_signatures_for_address_local(
493        &self,
494        pubkey: &Pubkey,
495        config: Option<RpcSignaturesForAddressConfig>,
496    ) -> SvmAccessContext<Vec<RpcConfirmedTransactionStatusWithSignature>> {
497        self.with_contextualized_svm_reader(|svm_reader| {
498            let current_slot = svm_reader.get_latest_absolute_slot();
499
500            let config = config.clone().unwrap_or_default();
501            let limit = config.limit.unwrap_or(1000);
502
503            let mut before_slot = None;
504            let mut until_slot = None;
505
506            let sigs: Vec<_> = svm_reader
507                .transactions
508                .iter()
509                .filter_map(|(sig, status)| {
510                    let TransactionWithStatusMeta(slot, tx, _, err) = status.expect_processed();
511
512                    if *slot < config.clone().min_context_slot.unwrap_or_default() {
513                        return None;
514                    }
515
516                    if Some(sig.to_string()) == config.clone().before {
517                        before_slot = Some(*slot)
518                    }
519
520                    if Some(sig.to_string()) == config.clone().until {
521                        until_slot = Some(*slot)
522                    }
523
524                    // Check if the pubkey is a signer
525                    let is_signer = tx
526                        .message
527                        .static_account_keys()
528                        .iter()
529                        .position(|pk| pk == pubkey)
530                        .map(|i| tx.message.is_signer(i))
531                        .unwrap_or(false);
532
533                    if !is_signer {
534                        return None;
535                    }
536
537                    // Determine confirmation status
538                    let confirmation_status = match current_slot {
539                        cs if cs == *slot => SolanaTransactionConfirmationStatus::Processed,
540                        cs if cs < slot + FINALIZATION_SLOT_THRESHOLD => {
541                            SolanaTransactionConfirmationStatus::Confirmed
542                        }
543                        _ => SolanaTransactionConfirmationStatus::Finalized,
544                    };
545
546                    Some(RpcConfirmedTransactionStatusWithSignature {
547                        err: err.clone(),
548                        slot: *slot,
549                        memo: None,
550                        block_time: None,
551                        confirmation_status: Some(confirmation_status),
552                        signature: sig.to_string(),
553                    })
554                })
555                .collect();
556
557            sigs.into_iter()
558                .filter(|sig| {
559                    if config.before.is_none() && config.until.is_none() {
560                        return true;
561                    }
562
563                    if config.before.is_some() && before_slot >= Some(sig.slot) {
564                        return true;
565                    }
566
567                    if config.until.is_some() && until_slot <= Some(sig.slot) {
568                        return true;
569                    }
570
571                    false
572                })
573                .take(limit)
574                .collect()
575        })
576    }
577
578    pub async fn get_signatures_for_address_local_then_remote(
579        &self,
580        client: &SurfnetRemoteClient,
581        pubkey: &Pubkey,
582        config: Option<RpcSignaturesForAddressConfig>,
583    ) -> SurfpoolContextualizedResult<Vec<RpcConfirmedTransactionStatusWithSignature>> {
584        let results = self.get_signatures_for_address_local(pubkey, config.clone());
585        let limit = config.clone().and_then(|c| c.limit).unwrap_or(1000);
586
587        let mut combined_results = results.inner.clone();
588        if combined_results.len() < limit {
589            let mut remote_results = client.get_signatures_for_address(pubkey, config).await?;
590            combined_results.append(&mut remote_results);
591        }
592
593        Ok(results.with_new_value(combined_results))
594    }
595
596    pub async fn get_signatures_for_address(
597        &self,
598        remote_ctx: &Option<(SurfnetRemoteClient, ())>,
599        pubkey: &Pubkey,
600        config: Option<RpcSignaturesForAddressConfig>,
601    ) -> SurfpoolContextualizedResult<Vec<RpcConfirmedTransactionStatusWithSignature>> {
602        let results = if let Some((remote_client, _)) = remote_ctx {
603            self.get_signatures_for_address_local_then_remote(remote_client, pubkey, config.clone())
604                .await?
605        } else {
606            self.get_signatures_for_address_local(pubkey, config)
607        };
608
609        Ok(results)
610    }
611}
612
613/// Functions for getting transactions from the underlying SurfnetSvm instance or remote client
614impl SurfnetSvmLocker {
615    /// Retrieves a transaction by signature, using local or remote based on context.
616    pub async fn get_transaction(
617        &self,
618        remote_ctx: &Option<(SurfnetRemoteClient, Option<UiTransactionEncoding>)>,
619        signature: &Signature,
620    ) -> SvmAccessContext<GetTransactionResult> {
621        if let Some((remote_client, encoding)) = remote_ctx {
622            self.get_transaction_local_then_remote(remote_client, signature, *encoding)
623                .await
624        } else {
625            self.get_transaction_local(signature)
626        }
627    }
628
629    /// Retrieves a transaction from local cache, returning a contextualized result.
630    pub fn get_transaction_local(
631        &self,
632        signature: &Signature,
633    ) -> SvmAccessContext<GetTransactionResult> {
634        self.with_contextualized_svm_reader(|svm_reader| {
635            let latest_absolute_slot = svm_reader.get_latest_absolute_slot();
636
637            match svm_reader.transactions.get(signature).map(|entry| {
638                Into::<EncodedConfirmedTransactionWithStatusMeta>::into(
639                    entry.expect_processed().clone(),
640                )
641            }) {
642                Some(tx) => {
643                    GetTransactionResult::found_transaction(*signature, tx, latest_absolute_slot)
644                }
645                None => GetTransactionResult::None(*signature),
646            }
647        })
648    }
649
650    /// Retrieves a transaction locally then from remote if missing, returning a contextualized result.
651    pub async fn get_transaction_local_then_remote(
652        &self,
653        client: &SurfnetRemoteClient,
654        signature: &Signature,
655        encoding: Option<UiTransactionEncoding>,
656    ) -> SvmAccessContext<GetTransactionResult> {
657        let local_result = self.get_transaction_local(signature);
658        if local_result.inner.is_none() {
659            let remote_result = client
660                .get_transaction(*signature, encoding, local_result.slot)
661                .await;
662            local_result.with_new_value(remote_result)
663        } else {
664            local_result
665        }
666    }
667}
668
669/// Functions for simulating and processing transactions in the underlying SurfnetSvm instance
670impl SurfnetSvmLocker {
671    /// Simulates a transaction on the SVM, returning detailed info or failure metadata.
672    pub fn simulate_transaction(
673        &self,
674        transaction: VersionedTransaction,
675        sigverify: bool,
676    ) -> Result<SimulatedTransactionInfo, FailedTransactionMetadata> {
677        self.with_svm_reader(|svm_reader| {
678            svm_reader.simulate_transaction(transaction.clone(), sigverify)
679        })
680    }
681
682    /// Processes a transaction: verifies signatures, preflight sim, sends to SVM, and enqueues status events.
683    pub async fn process_transaction(
684        &self,
685        remote_ctx: &Option<SurfnetRemoteClient>,
686        transaction: VersionedTransaction,
687        status_tx: Sender<TransactionStatusEvent>,
688        skip_preflight: bool,
689    ) -> SurfpoolContextualizedResult<()> {
690        let remote_ctx = &remote_ctx.get_remote_ctx(CommitmentConfig::confirmed());
691        let (latest_absolute_slot, latest_epoch_info, latest_blockhash) =
692            self.with_svm_writer(|svm_writer| {
693                let latest_absolute_slot = svm_writer.get_latest_absolute_slot();
694                svm_writer.notify_signature_subscribers(
695                    SignatureSubscriptionType::received(),
696                    &transaction.signatures[0],
697                    latest_absolute_slot,
698                    None,
699                );
700                (
701                    latest_absolute_slot,
702                    svm_writer.latest_epoch_info(),
703                    svm_writer.latest_blockhash(),
704                )
705            });
706
707        let signature = transaction.signatures[0];
708
709        // find accounts that are needed for this transaction but are missing from the local
710        // svm cache, fetch them from the RPC, and insert them locally
711        let pubkeys_from_message = self
712            .get_pubkeys_from_message(remote_ctx, &transaction.message)
713            .await?;
714
715        let account_updates = self
716            .get_multiple_accounts(remote_ctx, &pubkeys_from_message, None)
717            .await?
718            .inner;
719
720        self.with_svm_writer(|svm_writer| {
721            for update in &account_updates {
722                svm_writer.write_account_update(update.clone());
723            }
724
725            let accounts_before = pubkeys_from_message
726                .iter()
727                .map(|p| svm_writer.inner.get_account(p))
728                .collect::<Vec<Option<Account>>>();
729
730            // if not skipping preflight, simulate the transaction
731            if !skip_preflight {
732                match svm_writer.simulate_transaction(transaction.clone(), true) {
733                    Ok(_) => {}
734                    Err(res) => {
735                        let _ = svm_writer
736                            .simnet_events_tx
737                            .try_send(SimnetEvent::error(format!(
738                                "Transaction simulation failed: {}",
739                                res.err
740                            )));
741                        let meta = convert_transaction_metadata_from_canonical(&res.meta);
742                        let _ = status_tx.try_send(TransactionStatusEvent::SimulationFailure((
743                            res.err.clone(),
744                            meta,
745                        )));
746                        svm_writer.notify_signature_subscribers(
747                            SignatureSubscriptionType::processed(),
748                            &signature,
749                            latest_absolute_slot,
750                            Some(res.err),
751                        );
752                        return;
753                    }
754                }
755            }
756            // send the transaction to the SVM
757            let err = match svm_writer
758                .send_transaction(transaction.clone(), false /* cu_analysis_enabled */)
759            {
760                Ok(res) => {
761                    let accounts_after = pubkeys_from_message
762                        .iter()
763                        .map(|p| svm_writer.inner.get_account(p))
764                        .collect::<Vec<Option<Account>>>();
765
766                    for (pubkey, (before, after)) in pubkeys_from_message
767                        .iter()
768                        .zip(accounts_before.iter().zip(accounts_after))
769                    {
770                        if before.ne(&after) {
771                            if let Some(after) = &after {
772                                svm_writer.update_account_registries(pubkey, after);
773                            }
774                            svm_writer
775                                .notify_account_subscribers(pubkey, &after.unwrap_or_default());
776                        }
777                    }
778
779                    let transaction_meta = convert_transaction_metadata_from_canonical(&res);
780                    let _ = svm_writer
781                        .geyser_events_tx
782                        .send(GeyserEvent::NewTransaction(
783                            transaction.clone(),
784                            transaction_meta.clone(),
785                            latest_absolute_slot,
786                        ));
787                    let _ = status_tx.try_send(TransactionStatusEvent::Success(
788                        TransactionConfirmationStatus::Processed,
789                    ));
790                    svm_writer
791                        .transactions_queued_for_confirmation
792                        .push_back((transaction.clone(), status_tx.clone()));
793                    None
794                }
795                Err(res) => {
796                    let transaction_meta = convert_transaction_metadata_from_canonical(&res.meta);
797                    let _ = svm_writer
798                        .simnet_events_tx
799                        .try_send(SimnetEvent::error(format!(
800                            "Transaction execution failed: {}",
801                            res.err
802                        )));
803                    let _ = status_tx.try_send(TransactionStatusEvent::ExecutionFailure((
804                        res.err.clone(),
805                        transaction_meta,
806                    )));
807                    Some(res.err)
808                }
809            };
810
811            svm_writer.notify_signature_subscribers(
812                SignatureSubscriptionType::processed(),
813                &signature,
814                latest_absolute_slot,
815                err,
816            );
817        });
818
819        Ok(SvmAccessContext::new(
820            latest_absolute_slot,
821            latest_epoch_info,
822            latest_blockhash,
823            (),
824        ))
825    }
826}
827
828/// Functions for writing account updates to the underlying SurfnetSvm instance
829impl SurfnetSvmLocker {
830    /// Writes a single account update into the SVM state if present.
831    pub fn write_account_update(&self, account_update: GetAccountResult) {
832        if !account_update.requires_update() {
833            return;
834        }
835
836        self.with_svm_writer(move |svm_writer| {
837            svm_writer.write_account_update(account_update.clone())
838        })
839    }
840
841    /// Writes multiple account updates into the SVM state when any are present.
842    pub fn write_multiple_account_updates(&self, account_updates: &[GetAccountResult]) {
843        if account_updates
844            .iter()
845            .all(|update| !update.requires_update())
846        {
847            return;
848        }
849
850        self.with_svm_writer(move |svm_writer| {
851            for update in account_updates {
852                svm_writer.write_account_update(update.clone());
853            }
854        });
855    }
856}
857
858/// Token account related functions
859impl SurfnetSvmLocker {
860    /// Fetches all token accounts for an owner, returning remote results and missing pubkeys contexts.
861    pub async fn get_token_accounts_by_owner(
862        &self,
863        remote_ctx: &Option<SurfnetRemoteClient>,
864        owner: Pubkey,
865        filter: &TokenAccountsFilter,
866        config: &RpcAccountInfoConfig,
867    ) -> SurfpoolContextualizedResult<Vec<RpcKeyedAccount>> {
868        if let Some(remote_client) = remote_ctx {
869            self.get_token_accounts_by_owner_local_then_remote(owner, filter, remote_client, config)
870                .await
871        } else {
872            Ok(self.get_token_accounts_by_owner_local(owner, filter, config))
873        }
874    }
875
876    pub fn get_token_accounts_by_owner_local(
877        &self,
878        owner: Pubkey,
879        filter: &TokenAccountsFilter,
880        config: &RpcAccountInfoConfig,
881    ) -> SvmAccessContext<Vec<RpcKeyedAccount>> {
882        self.with_contextualized_svm_reader(|svm_reader| {
883            svm_reader
884                .get_parsed_token_accounts_by_owner(&owner)
885                .iter()
886                .filter_map(|(pubkey, token_account)| {
887                    let account = svm_reader.accounts_registry.get(pubkey)?;
888                    if match filter {
889                        TokenAccountsFilter::Mint(mint) => token_account.mint.eq(mint),
890                        TokenAccountsFilter::ProgramId(program_id) => account.owner.eq(program_id),
891                    } {
892                        let account_data = encode_ui_account(
893                            pubkey,
894                            account,
895                            config.encoding.unwrap_or(UiAccountEncoding::Base64),
896                            None,
897                            config.data_slice,
898                        );
899                        Some(RpcKeyedAccount {
900                            pubkey: pubkey.to_string(),
901                            account: account_data,
902                        })
903                    } else {
904                        None
905                    }
906                })
907                .collect::<Vec<_>>()
908        })
909    }
910
911    pub async fn get_token_accounts_by_owner_local_then_remote(
912        &self,
913        owner: Pubkey,
914        filter: &TokenAccountsFilter,
915        remote_client: &SurfnetRemoteClient,
916        config: &RpcAccountInfoConfig,
917    ) -> SurfpoolContextualizedResult<Vec<RpcKeyedAccount>> {
918        let SvmAccessContext {
919            slot,
920            latest_epoch_info,
921            latest_blockhash,
922            inner: local_accounts,
923        } = self.get_token_accounts_by_owner_local(owner, filter, config);
924
925        let remote_accounts = remote_client
926            .get_token_accounts_by_owner(owner, filter, config)
927            .await?;
928
929        let mut combined_accounts = remote_accounts;
930
931        for local_account in local_accounts {
932            if let Some((pos, _)) = combined_accounts
933                .iter()
934                .find_position(|RpcKeyedAccount { pubkey, .. }| pubkey.eq(&local_account.pubkey))
935            {
936                combined_accounts[pos] = local_account;
937            } else {
938                combined_accounts.push(local_account);
939            }
940        }
941
942        Ok(SvmAccessContext::new(
943            slot,
944            latest_epoch_info,
945            latest_blockhash,
946            combined_accounts,
947        ))
948    }
949
950    pub async fn get_token_accounts_by_delegate(
951        &self,
952        remote_ctx: &Option<SurfnetRemoteClient>,
953        delegate: Pubkey,
954        filter: &TokenAccountsFilter,
955        config: &RpcAccountInfoConfig,
956    ) -> SurfpoolContextualizedResult<Vec<RpcKeyedAccount>> {
957        // Validate that the program is supported if using ProgramId filter
958        if let TokenAccountsFilter::ProgramId(program_id) = filter {
959            if !is_supported_token_program(program_id) {
960                return Err(SurfpoolError::unsupported_token_program(*program_id));
961            }
962        }
963
964        if let Some(remote_client) = remote_ctx {
965            self.get_token_accounts_by_delegate_local_then_remote(
966                delegate,
967                filter,
968                remote_client,
969                config,
970            )
971            .await
972        } else {
973            Ok(self.get_token_accounts_by_delegate_local(delegate, filter, config))
974        }
975    }
976}
977
978/// Token account by delegate related functions
979impl SurfnetSvmLocker {
980    pub fn get_token_accounts_by_delegate_local(
981        &self,
982        delegate: Pubkey,
983        filter: &TokenAccountsFilter,
984        config: &RpcAccountInfoConfig,
985    ) -> SvmAccessContext<Vec<RpcKeyedAccount>> {
986        self.with_contextualized_svm_reader(|svm_reader| {
987            svm_reader
988                .get_token_accounts_by_delegate(&delegate)
989                .iter()
990                .filter_map(|(pubkey, token_account)| {
991                    let include = match filter {
992                        TokenAccountsFilter::Mint(mint) => token_account.mint == *mint,
993                        TokenAccountsFilter::ProgramId(program_id) => {
994                            if let Some(account) = svm_reader.accounts_registry.get(pubkey) {
995                                account.owner == *program_id
996                                    && is_supported_token_program(program_id)
997                            } else {
998                                false
999                            }
1000                        }
1001                    };
1002
1003                    if include {
1004                        svm_reader
1005                            .accounts_registry
1006                            .get(pubkey)
1007                            .map(|account| RpcKeyedAccount {
1008                                pubkey: pubkey.to_string(),
1009                                account: encode_ui_account(
1010                                    pubkey,
1011                                    account,
1012                                    config.encoding.unwrap_or(UiAccountEncoding::Base64),
1013                                    None,
1014                                    config.data_slice,
1015                                ),
1016                            })
1017                    } else {
1018                        None
1019                    }
1020                })
1021                .collect::<Vec<_>>()
1022        })
1023    }
1024
1025    pub async fn get_token_accounts_by_delegate_local_then_remote(
1026        &self,
1027        delegate: Pubkey,
1028        filter: &TokenAccountsFilter,
1029        remote_client: &SurfnetRemoteClient,
1030        config: &RpcAccountInfoConfig,
1031    ) -> SurfpoolContextualizedResult<Vec<RpcKeyedAccount>> {
1032        let SvmAccessContext {
1033            slot,
1034            latest_epoch_info,
1035            latest_blockhash,
1036            inner: local_accounts,
1037        } = self.get_token_accounts_by_delegate_local(delegate, filter, config);
1038
1039        let remote_accounts = remote_client
1040            .get_token_accounts_by_delegate(delegate, filter, config)
1041            .await?;
1042
1043        let mut combined_accounts = remote_accounts;
1044
1045        for local_account in local_accounts {
1046            if let Some((pos, _)) = combined_accounts
1047                .iter()
1048                .find_position(|RpcKeyedAccount { pubkey, .. }| pubkey.eq(&local_account.pubkey))
1049            {
1050                // Replace remote account with local one (local takes precedence)
1051                combined_accounts[pos] = local_account;
1052            } else {
1053                // Add local account that wasn't found in remote results
1054                combined_accounts.push(local_account);
1055            }
1056        }
1057
1058        Ok(SvmAccessContext::new(
1059            slot,
1060            latest_epoch_info,
1061            latest_blockhash,
1062            combined_accounts,
1063        ))
1064    }
1065}
1066
1067/// Get largest account related account
1068impl SurfnetSvmLocker {
1069    pub fn get_token_largest_accounts_local(
1070        &self,
1071        mint: &Pubkey,
1072    ) -> SvmAccessContext<Vec<RpcTokenAccountBalance>> {
1073        self.with_contextualized_svm_reader(|svm_reader| {
1074            let token_accounts = svm_reader.get_token_accounts_by_mint(mint);
1075
1076            // get mint information to determine decimals
1077            let mint_decimals = if let Some(mint_account) = svm_reader.accounts_registry.get(mint) {
1078                if let Ok(mint_data) = Mint::unpack(&mint_account.data) {
1079                    mint_data.decimals
1080                } else {
1081                    0
1082                }
1083            } else {
1084                0
1085            };
1086
1087            // convert to RpcTokenAccountBalance and sort by balance
1088            let mut balances: Vec<RpcTokenAccountBalance> = token_accounts
1089                .into_iter()
1090                .map(|(pubkey, token_account)| {
1091                    let ui_amount = if mint_decimals > 0 {
1092                        Some(
1093                            token_account.amount as f64 / (10_u64.pow(mint_decimals as u32) as f64),
1094                        )
1095                    } else {
1096                        Some(token_account.amount as f64)
1097                    };
1098
1099                    let ui_amount_string = if mint_decimals > 0 {
1100                        format!(
1101                            "{:.precision$}",
1102                            token_account.amount as f64 / (10_u64.pow(mint_decimals as u32) as f64),
1103                            precision = mint_decimals as usize
1104                        )
1105                    } else {
1106                        token_account.amount.to_string()
1107                    };
1108
1109                    RpcTokenAccountBalance {
1110                        address: pubkey.to_string(),
1111                        amount: UiTokenAmount {
1112                            amount: token_account.amount.to_string(),
1113                            decimals: mint_decimals,
1114                            ui_amount,
1115                            ui_amount_string,
1116                        },
1117                    }
1118                })
1119                .collect();
1120
1121            // sort by amount in descending order
1122            balances.sort_by(|a, b| {
1123                let amount_a: u64 = a.amount.amount.parse().unwrap_or(0);
1124                let amount_b: u64 = b.amount.amount.parse().unwrap_or(0);
1125                amount_b.cmp(&amount_a)
1126            });
1127
1128            // limit to top 20 accounts
1129            balances.truncate(20);
1130
1131            balances
1132        })
1133    }
1134
1135    pub async fn get_token_largest_accounts_local_then_remote(
1136        &self,
1137        client: &SurfnetRemoteClient,
1138        mint: &Pubkey,
1139        commitment_config: CommitmentConfig,
1140    ) -> SurfpoolContextualizedResult<Vec<RpcTokenAccountBalance>> {
1141        let SvmAccessContext {
1142            slot,
1143            latest_epoch_info,
1144            latest_blockhash,
1145            inner: local_accounts,
1146        } = self.get_token_largest_accounts_local(mint);
1147
1148        let remote_accounts = client
1149            .get_token_largest_accounts(mint, commitment_config)
1150            .await?;
1151
1152        let mut combined_accounts = remote_accounts;
1153
1154        // if the account is in both the local and remote list, add the local one and not the remote
1155        for local_account in local_accounts {
1156            if let Some((pos, _)) = combined_accounts
1157                .iter()
1158                .find_position(|remote_account| remote_account.address == local_account.address)
1159            {
1160                combined_accounts[pos] = local_account;
1161            } else {
1162                combined_accounts.push(local_account);
1163            }
1164        }
1165
1166        // re-sort and limit after combining
1167        combined_accounts.sort_by(|a, b| {
1168            let amount_a: u64 = a.amount.amount.parse().unwrap_or(0);
1169            let amount_b: u64 = b.amount.amount.parse().unwrap_or(0);
1170            amount_b.cmp(&amount_a)
1171        });
1172        combined_accounts.truncate(20);
1173
1174        Ok(SvmAccessContext::new(
1175            slot,
1176            latest_epoch_info,
1177            latest_blockhash,
1178            combined_accounts,
1179        ))
1180    }
1181
1182    /// Fetches the largest token accounts for a specific mint, returning contextualized results.
1183    pub async fn get_token_largest_accounts(
1184        &self,
1185        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
1186        mint: &Pubkey,
1187    ) -> SurfpoolContextualizedResult<Vec<RpcTokenAccountBalance>> {
1188        if let Some((remote_client, commitment_config)) = remote_ctx {
1189            self.get_token_largest_accounts_local_then_remote(
1190                remote_client,
1191                mint,
1192                *commitment_config,
1193            )
1194            .await
1195        } else {
1196            Ok(self.get_token_largest_accounts_local(mint))
1197        }
1198    }
1199}
1200
1201/// Address lookup table related functions
1202impl SurfnetSvmLocker {
1203    /// Extracts pubkeys from a VersionedMessage, resolving address lookup tables as needed.
1204    pub async fn get_pubkeys_from_message(
1205        &self,
1206        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
1207        message: &VersionedMessage,
1208    ) -> SurfpoolResult<Vec<Pubkey>> {
1209        match message {
1210            VersionedMessage::Legacy(message) => Ok(message.account_keys.clone()),
1211            VersionedMessage::V0(message) => {
1212                let alts = message.address_table_lookups.clone();
1213                let mut acc_keys = message.account_keys.clone();
1214                let mut alt_pubkeys = alts.iter().map(|msg| msg.account_key).collect::<Vec<_>>();
1215
1216                let mut table_entries = join_all(alts.iter().map(|msg| async {
1217                    let loaded_addresses = self
1218                        .get_lookup_table_addresses(remote_ctx, msg)
1219                        .await?
1220                        .inner;
1221                    let mut combined = loaded_addresses.writable;
1222                    combined.extend(loaded_addresses.readonly);
1223                    Ok::<_, SurfpoolError>(combined)
1224                }))
1225                .await
1226                .into_iter()
1227                .collect::<Result<Vec<Vec<Pubkey>>, SurfpoolError>>()?
1228                .into_iter()
1229                .flatten()
1230                .collect();
1231
1232                acc_keys.append(&mut alt_pubkeys);
1233                acc_keys.append(&mut table_entries);
1234                Ok(acc_keys)
1235            }
1236        }
1237    }
1238
1239    /// Retrieves loaded addresses from a lookup table account, validating owner and indices.
1240    pub async fn get_lookup_table_addresses(
1241        &self,
1242        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
1243        address_table_lookup: &MessageAddressTableLookup,
1244    ) -> SurfpoolContextualizedResult<LoadedAddresses> {
1245        let result = self
1246            .get_account(remote_ctx, &address_table_lookup.account_key, None)
1247            .await?;
1248        let table_account = result.inner.clone().map_account()?;
1249
1250        if table_account.owner == solana_sdk_ids::address_lookup_table::id() {
1251            let SvmAccessContext {
1252                slot: current_slot,
1253                inner: slot_hashes,
1254                ..
1255            } = self.with_contextualized_svm_reader(|svm_reader| {
1256                svm_reader
1257                    .inner
1258                    .get_sysvar::<solana_sdk::sysvar::slot_hashes::SlotHashes>()
1259            });
1260
1261            //let current_slot = self.get_latest_absolute_slot(); // or should i use this?
1262            let data = &table_account.data.clone();
1263            let lookup_table = AddressLookupTable::deserialize(data).map_err(|_ix_err| {
1264                SurfpoolError::invalid_account_data(
1265                    address_table_lookup.account_key,
1266                    table_account.data,
1267                    Some("Attempted to lookup addresses from an invalid account"),
1268                )
1269            })?;
1270
1271            let loaded_addresses = LoadedAddresses {
1272                writable: lookup_table
1273                    .lookup(
1274                        current_slot,
1275                        &address_table_lookup.writable_indexes,
1276                        &slot_hashes,
1277                    )
1278                    .map_err(|_ix_err| {
1279                        SurfpoolError::invalid_lookup_index(address_table_lookup.account_key)
1280                    })?,
1281                readonly: lookup_table
1282                    .lookup(
1283                        current_slot,
1284                        &address_table_lookup.readonly_indexes,
1285                        &slot_hashes,
1286                    )
1287                    .map_err(|_ix_err| {
1288                        SurfpoolError::invalid_lookup_index(address_table_lookup.account_key)
1289                    })?,
1290            };
1291            Ok(result.with_new_value(loaded_addresses))
1292        } else {
1293            Err(SurfpoolError::invalid_account_owner(
1294                table_account.owner,
1295                Some("Attempted to lookup addresses from an account owned by the wrong program"),
1296            ))
1297        }
1298    }
1299}
1300
1301/// Profiling helper functions
1302impl SurfnetSvmLocker {
1303    /// Estimates compute units for a transaction via contextualized simulation.
1304    pub fn estimate_compute_units(
1305        &self,
1306        transaction: &VersionedTransaction,
1307    ) -> SvmAccessContext<ComputeUnitsEstimationResult> {
1308        self.with_contextualized_svm_reader(|svm_reader| {
1309            svm_reader.estimate_compute_units(transaction)
1310        })
1311    }
1312    pub async fn profile_transaction(
1313        &self,
1314        remote_ctx: &Option<SurfnetRemoteClient>,
1315        transaction: &VersionedTransaction,
1316        encoding: Option<UiAccountEncoding>,
1317    ) -> SurfpoolContextualizedResult<ProfileResult> {
1318        let SvmAccessContext {
1319            slot,
1320            latest_epoch_info,
1321            latest_blockhash,
1322            inner: mut svm_clone,
1323        } = self.with_contextualized_svm_reader(|svm_reader| svm_reader.clone());
1324
1325        let (dummy_simnet_tx, _) = crossbeam_channel::bounded(1);
1326        let (dummy_geyser_tx, _) = crossbeam_channel::bounded(1);
1327        svm_clone.simnet_events_tx = dummy_simnet_tx;
1328        svm_clone.geyser_events_tx = dummy_geyser_tx;
1329
1330        let svm_locker = SurfnetSvmLocker::new(svm_clone);
1331
1332        let remote_ctx_with_config = remote_ctx
1333            .clone()
1334            .map(|client| (client, CommitmentConfig::confirmed()));
1335
1336        let account_keys = svm_locker
1337            .get_pubkeys_from_message(&remote_ctx_with_config, &transaction.message)
1338            .await?;
1339
1340        let pre_execution_capture = {
1341            let mut capture = BTreeMap::new();
1342            for pubkey in &account_keys {
1343                let account = svm_locker
1344                    .get_account(&remote_ctx_with_config, pubkey, None)
1345                    .await?
1346                    .inner;
1347
1348                snapshot_get_account_result(
1349                    &mut capture,
1350                    account,
1351                    encoding.unwrap_or(UiAccountEncoding::Base64),
1352                );
1353            }
1354            capture
1355        };
1356
1357        let compute_units_estimation_result = svm_locker.estimate_compute_units(transaction).inner;
1358
1359        let (status_tx, status_rx) = crossbeam_channel::unbounded();
1360        let _ = svm_locker
1361            .process_transaction(remote_ctx, transaction.clone(), status_tx, true)
1362            .await?;
1363
1364        let simnet_events_tx = self.simnet_events_tx();
1365        loop {
1366            if let Ok(status) = status_rx.try_recv() {
1367                match status {
1368                    TransactionStatusEvent::Success(_) => break,
1369                    TransactionStatusEvent::ExecutionFailure((err, _)) => {
1370                        let _ = simnet_events_tx.try_send(SimnetEvent::WarnLog(
1371                            chrono::Local::now(),
1372                            format!(
1373                                "Transaction {} failed during snapshot simulation: {}",
1374                                transaction.signatures[0], err
1375                            ),
1376                        ));
1377                        return Err(SurfpoolError::internal(format!(
1378                            "Transaction {} failed during snapshot simulation: {}",
1379                            transaction.signatures[0], err
1380                        )));
1381                    }
1382                    TransactionStatusEvent::SimulationFailure(_) => unreachable!(),
1383                    TransactionStatusEvent::VerificationFailure(_) => {
1384                        let _ = simnet_events_tx.try_send(SimnetEvent::WarnLog(
1385                            chrono::Local::now(),
1386                            format!(
1387                                "Transaction {} verification failed during snapshot simulation",
1388                                transaction.signatures[0]
1389                            ),
1390                        ));
1391                        return Err(SurfpoolError::internal(format!(
1392                            "Transaction {} verification failed during snapshot simulation",
1393                            transaction.signatures[0]
1394                        )));
1395                    }
1396                }
1397            }
1398        }
1399
1400        let post_execution_capture = {
1401            let mut capture = BTreeMap::new();
1402            for pubkey in &account_keys {
1403                // get the account locally after execution, because execution should have inserted from
1404                // the remote if not found locally
1405                let account = svm_locker.get_account_local(pubkey).inner;
1406
1407                snapshot_get_account_result(
1408                    &mut capture,
1409                    account,
1410                    encoding.unwrap_or(UiAccountEncoding::Base64),
1411                );
1412            }
1413            capture
1414        };
1415
1416        Ok(SvmAccessContext::new(
1417            slot,
1418            latest_epoch_info,
1419            latest_blockhash,
1420            ProfileResult {
1421                compute_units: compute_units_estimation_result,
1422                state: ProfileState::new(pre_execution_capture, post_execution_capture),
1423            },
1424        ))
1425    }
1426
1427    /// Records profiling results under a tag and emits a tagged profile event.
1428    pub fn write_profiling_results(&self, tag: String, profile_result: ProfileResult) {
1429        self.with_svm_writer(|svm_writer| {
1430            svm_writer
1431                .tagged_profiling_results
1432                .entry(tag.clone())
1433                .or_default()
1434                .push(profile_result.clone());
1435            let _ = svm_writer
1436                .simnet_events_tx
1437                .try_send(SimnetEvent::tagged_profile(
1438                    profile_result.clone(),
1439                    tag.clone(),
1440                ));
1441        });
1442    }
1443}
1444
1445/// Program account related functions
1446impl SurfnetSvmLocker {
1447    /// Clones a program account from source to destination, handling upgradeable loader state.
1448    pub async fn clone_program_account(
1449        &self,
1450        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
1451        source_program_id: &Pubkey,
1452        destination_program_id: &Pubkey,
1453    ) -> SurfpoolContextualizedResult<()> {
1454        let expected_source_program_data_address = get_program_data_address(source_program_id);
1455
1456        let result = self
1457            .get_multiple_accounts(
1458                remote_ctx,
1459                &[*source_program_id, expected_source_program_data_address],
1460                None,
1461            )
1462            .await?;
1463
1464        let mut accounts = result
1465            .inner
1466            .clone()
1467            .into_iter()
1468            .map(|a| a.map_account())
1469            .collect::<SurfpoolResult<Vec<Account>>>()?;
1470
1471        let source_program_data_account = accounts.remove(1);
1472        let source_program_account = accounts.remove(0);
1473
1474        let BpfUpgradeableLoaderAccountType::Program(UiProgram {
1475            program_data: source_program_data_address,
1476        }) = parse_bpf_upgradeable_loader(&source_program_account.data).map_err(|e| {
1477            SurfpoolError::invalid_program_account(source_program_id, e.to_string())
1478        })?
1479        else {
1480            return Err(SurfpoolError::expected_program_account(source_program_id));
1481        };
1482
1483        if source_program_data_address.ne(&expected_source_program_data_address.to_string()) {
1484            return Err(SurfpoolError::invalid_program_account(
1485                source_program_id,
1486                format!(
1487                    "Program data address mismatch: expected {}, found {}",
1488                    expected_source_program_data_address, source_program_data_address
1489                ),
1490            ));
1491        }
1492
1493        let destination_program_data_address = get_program_data_address(destination_program_id);
1494
1495        // create a new program account that has the `program_data` field set to the
1496        // destination program data address
1497        let mut new_program_account = source_program_account;
1498        new_program_account.data = bincode::serialize(&UpgradeableLoaderState::Program {
1499            programdata_address: destination_program_data_address,
1500        })
1501        .map_err(|e| SurfpoolError::internal(format!("Failed to serialize program data: {}", e)))?;
1502
1503        self.with_svm_writer(|svm_writer| {
1504            svm_writer.set_account(
1505                &destination_program_data_address,
1506                source_program_data_account.clone(),
1507            )?;
1508
1509            svm_writer.set_account(destination_program_id, new_program_account.clone())?;
1510            Ok::<(), SurfpoolError>(())
1511        })?;
1512
1513        Ok(result.with_new_value(()))
1514    }
1515
1516    pub async fn set_program_authority(
1517        &self,
1518        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
1519        program_id: Pubkey,
1520        new_authority: Option<Pubkey>,
1521    ) -> SurfpoolContextualizedResult<()> {
1522        let SvmAccessContext {
1523            slot,
1524            latest_epoch_info,
1525            latest_blockhash,
1526            inner: mut get_account_result,
1527        } = self.get_account(remote_ctx, &program_id, None).await?;
1528
1529        let original_authority = match &mut get_account_result {
1530            GetAccountResult::None(pubkey) => {
1531                return Err(SurfpoolError::invalid_program_account(
1532                    pubkey,
1533                    "Account not found",
1534                ))
1535            }
1536            GetAccountResult::FoundAccount(pubkey, program_account, _) => {
1537                let programdata_address = get_program_data_address(pubkey);
1538                let mut programdata_account_result = self
1539                    .get_account(remote_ctx, &programdata_address, None)
1540                    .await?
1541                    .inner;
1542                match &mut programdata_account_result {
1543                    GetAccountResult::None(pubkey) => {
1544                        return Err(SurfpoolError::invalid_program_account(
1545                            pubkey,
1546                            "Program data account does not exist",
1547                        ));
1548                    }
1549                    GetAccountResult::FoundAccount(_, programdata_account, _) => {
1550                        let original_authority = update_programdata_account(
1551                            &program_id,
1552                            programdata_account,
1553                            new_authority,
1554                        )?;
1555
1556                        get_account_result = GetAccountResult::FoundProgramAccount(
1557                            (pubkey.clone(), program_account.clone()),
1558                            (programdata_address, Some(programdata_account.clone())),
1559                        );
1560
1561                        original_authority
1562                    }
1563                    GetAccountResult::FoundProgramAccount(_, _) => {
1564                        return Err(SurfpoolError::invalid_program_account(
1565                            pubkey,
1566                            "Not a program account",
1567                        ));
1568                    }
1569                }
1570            }
1571            GetAccountResult::FoundProgramAccount(_, (_, None)) => {
1572                return Err(SurfpoolError::invalid_program_account(
1573                    &program_id,
1574                    "Program data account does not exist",
1575                ))
1576            }
1577            GetAccountResult::FoundProgramAccount(_, (_, Some(programdata_account))) => {
1578                update_programdata_account(&program_id, programdata_account, new_authority)?
1579            }
1580        };
1581
1582        let simnet_events_tx = self.simnet_events_tx();
1583        match (original_authority, new_authority) {
1584            (Some(original), Some(new)) => {
1585                if original != new {
1586                    let _ = simnet_events_tx.send(SimnetEvent::info(format!(
1587                        "Setting new authority for program {}",
1588                        program_id
1589                    )));
1590                    let _ = simnet_events_tx
1591                        .send(SimnetEvent::info(format!("Old Authority: {}", original)));
1592                    let _ =
1593                        simnet_events_tx.send(SimnetEvent::info(format!("New Authority: {}", new)));
1594                } else {
1595                    let _ = simnet_events_tx.send(SimnetEvent::info(format!(
1596                        "No authority change for program {}",
1597                        program_id
1598                    )));
1599                }
1600            }
1601            (Some(original), None) => {
1602                let _ = simnet_events_tx.send(SimnetEvent::info(format!(
1603                    "Removing authority for program {}",
1604                    program_id
1605                )));
1606                let _ = simnet_events_tx
1607                    .send(SimnetEvent::info(format!("Old Authority: {}", original)));
1608            }
1609            (None, Some(new)) => {
1610                let _ = simnet_events_tx.send(SimnetEvent::info(format!(
1611                    "Setting new authority for program {}",
1612                    program_id
1613                )));
1614                let _ = simnet_events_tx.send(SimnetEvent::info(format!("Old Authority: None")));
1615                let _ = simnet_events_tx.send(SimnetEvent::info(format!("New Authority: {}", new)));
1616            }
1617            (None, None) => {
1618                let _ = simnet_events_tx.send(SimnetEvent::info(format!(
1619                    "No authority change for program {}",
1620                    program_id
1621                )));
1622            }
1623        };
1624
1625        self.write_account_update(get_account_result);
1626
1627        Ok(SvmAccessContext::new(
1628            slot,
1629            latest_epoch_info,
1630            latest_blockhash,
1631            (),
1632        ))
1633    }
1634
1635    pub async fn get_program_accounts(
1636        &self,
1637        remote_ctx: &Option<SurfnetRemoteClient>,
1638        program_id: &Pubkey,
1639        account_config: RpcAccountInfoConfig,
1640        filters: Option<Vec<RpcFilterType>>,
1641    ) -> SurfpoolContextualizedResult<Vec<RpcKeyedAccount>> {
1642        if let Some(remote_client) = remote_ctx {
1643            self.get_program_accounts_local_then_remote(
1644                remote_client,
1645                program_id,
1646                account_config,
1647                filters,
1648            )
1649            .await
1650        } else {
1651            self.get_program_accounts_local(program_id, account_config, filters)
1652        }
1653    }
1654
1655    /// Retrieves program accounts from the local SVM cache, returning a contextualized result.
1656    pub fn get_program_accounts_local(
1657        &self,
1658        program_id: &Pubkey,
1659        account_config: RpcAccountInfoConfig,
1660        filters: Option<Vec<RpcFilterType>>,
1661    ) -> SurfpoolContextualizedResult<Vec<RpcKeyedAccount>> {
1662        let res = self.with_contextualized_svm_reader(|svm_reader| {
1663            svm_reader.get_account_owned_by(*program_id)
1664        });
1665
1666        let mut filtered = vec![];
1667        for (pubkey, account) in res.inner.iter() {
1668            if let Some(ref active_filters) = filters {
1669                match apply_rpc_filters(&account.data, active_filters) {
1670                    Ok(true) => {}           // Account matches all filters
1671                    Ok(false) => continue,   // Filtered out
1672                    Err(e) => return Err(e), // Error applying filter, already JsonRpcError
1673                }
1674            }
1675            let data_slice = account_config.data_slice;
1676
1677            filtered.push(RpcKeyedAccount {
1678                pubkey: pubkey.to_string(),
1679                account: encode_ui_account(
1680                    pubkey,
1681                    account,
1682                    account_config.encoding.unwrap_or(UiAccountEncoding::Base64),
1683                    None, // No additional data for now
1684                    data_slice,
1685                ),
1686            });
1687        }
1688        Ok(res.with_new_value(filtered))
1689    }
1690
1691    /// Retrieves program accounts from the local cache and remote client, combining results.
1692    pub async fn get_program_accounts_local_then_remote(
1693        &self,
1694        client: &SurfnetRemoteClient,
1695        program_id: &Pubkey,
1696        account_config: RpcAccountInfoConfig,
1697        filters: Option<Vec<RpcFilterType>>,
1698    ) -> SurfpoolContextualizedResult<Vec<RpcKeyedAccount>> {
1699        let SvmAccessContext {
1700            slot,
1701            latest_epoch_info,
1702            latest_blockhash,
1703            inner: local_accounts,
1704        } = self.get_program_accounts_local(program_id, account_config.clone(), filters.clone())?;
1705        let remote_accounts = client
1706            .get_program_accounts(program_id, account_config, filters)
1707            .await?;
1708
1709        let mut combined_accounts = remote_accounts;
1710
1711        for local_account in local_accounts {
1712            // if the local account is in the remote set, replace it with the local one
1713            if let Some((pos, _)) = combined_accounts.iter().find_position(
1714                |RpcKeyedAccount {
1715                     pubkey: remote_pubkey,
1716                     ..
1717                 }| remote_pubkey.eq(&local_account.pubkey),
1718            ) {
1719                combined_accounts[pos] = local_account;
1720            } else {
1721                // otherwise, add the local account to the combined list
1722                combined_accounts.push(local_account);
1723            };
1724        }
1725
1726        Ok(SvmAccessContext {
1727            slot,
1728            latest_epoch_info,
1729            latest_blockhash,
1730            inner: combined_accounts,
1731        })
1732    }
1733}
1734
1735impl SurfnetSvmLocker {
1736    pub fn get_genesis_hash_local(&self) -> SvmAccessContext<Hash> {
1737        self.with_contextualized_svm_reader(|svm_reader| svm_reader.genesis_config.hash())
1738    }
1739
1740    pub async fn get_genesis_hash(
1741        &self,
1742        remote_ctx: &Option<SurfnetRemoteClient>,
1743    ) -> SurfpoolContextualizedResult<Hash> {
1744        if let Some(client) = remote_ctx {
1745            let remote_hash = client.get_genesis_hash().await?;
1746            Ok(self.with_contextualized_svm_reader(|_| remote_hash))
1747        } else {
1748            Ok(self.get_genesis_hash_local())
1749        }
1750    }
1751}
1752
1753/// Pass through functions for accessing the underlying SurfnetSvm instance
1754impl SurfnetSvmLocker {
1755    /// Returns a sender for simulation events from the underlying SVM.
1756    pub fn simnet_events_tx(&self) -> Sender<SimnetEvent> {
1757        self.with_svm_reader(|svm_reader| svm_reader.simnet_events_tx.clone())
1758    }
1759
1760    /// Retrieves the latest epoch info from the underlying SVM.
1761    pub fn get_epoch_info(&self) -> EpochInfo {
1762        self.with_svm_reader(|svm_reader| svm_reader.latest_epoch_info.clone())
1763    }
1764
1765    /// Retrieves the latest absolute slot from the underlying SVM.
1766    pub fn get_latest_absolute_slot(&self) -> Slot {
1767        self.with_svm_reader(|svm_reader| svm_reader.get_latest_absolute_slot())
1768    }
1769
1770    pub fn get_slot_for_commitment(&self, commitment: &CommitmentConfig) -> Slot {
1771        self.with_svm_reader(|svm_reader| {
1772            let slot = svm_reader.get_latest_absolute_slot();
1773            match commitment.commitment {
1774                CommitmentLevel::Processed => slot,
1775                CommitmentLevel::Confirmed => slot.saturating_sub(1),
1776                CommitmentLevel::Finalized => slot.saturating_sub(FINALIZATION_SLOT_THRESHOLD),
1777            }
1778        })
1779    }
1780
1781    /// Executes an airdrop via the underlying SVM.
1782    pub fn airdrop(&self, pubkey: &Pubkey, lamports: u64) -> TransactionResult {
1783        self.with_svm_writer(|svm_writer| svm_writer.airdrop(pubkey, lamports))
1784    }
1785
1786    /// Executes a batch airdrop via the underlying SVM.
1787    pub fn airdrop_pubkeys(&self, lamports: u64, addresses: &[Pubkey]) {
1788        self.with_svm_writer(|svm_writer| svm_writer.airdrop_pubkeys(lamports, addresses))
1789    }
1790
1791    /// Confirms the current block on the underlying SVM, returning `Ok(())` or an error.
1792    pub fn confirm_current_block(&self) -> SurfpoolResult<()> {
1793        self.with_svm_writer(|svm_writer| svm_writer.confirm_current_block())
1794    }
1795
1796    /// Subscribes for signature updates (confirmed/finalized) and returns a receiver of events.
1797    pub fn subscribe_for_signature_updates(
1798        &self,
1799        signature: &Signature,
1800        subscription_type: SignatureSubscriptionType,
1801    ) -> Receiver<(Slot, Option<TransactionError>)> {
1802        self.with_svm_writer(|svm_writer| {
1803            svm_writer.subscribe_for_signature_updates(signature, subscription_type.clone())
1804        })
1805    }
1806
1807    /// Subscribes for account updates and returns a receiver of account updates.
1808    pub fn subscribe_for_account_updates(
1809        &self,
1810        account_pubkey: &Pubkey,
1811        encoding: Option<UiAccountEncoding>,
1812    ) -> Receiver<UiAccount> {
1813        // Handles the locking/unlocking safely
1814        self.with_svm_writer(|svm_writer| {
1815            svm_writer.subscribe_for_account_updates(account_pubkey, encoding)
1816        })
1817    }
1818
1819    /// Subscribes for slot updates and returns a receiver of slot updates.
1820    pub fn subscribe_for_slot_updates(&self) -> Receiver<SlotInfo> {
1821        self.with_svm_writer(|svm_writer| svm_writer.subscribe_for_slot_updates())
1822    }
1823}
1824
1825fn snapshot_get_account_result(
1826    capture: &mut BTreeMap<Pubkey, Option<UiAccount>>,
1827    result: GetAccountResult,
1828    encoding: UiAccountEncoding,
1829) {
1830    match result {
1831        GetAccountResult::None(pubkey) => {
1832            capture.insert(pubkey, None);
1833        }
1834        GetAccountResult::FoundAccount(pubkey, account, _) => {
1835            capture.insert(
1836                pubkey,
1837                Some(encode_ui_account(&pubkey, &account, encoding, None, None)),
1838            );
1839        }
1840        GetAccountResult::FoundProgramAccount((pubkey, account), (data_pubkey, data_account)) => {
1841            capture.insert(
1842                pubkey,
1843                Some(encode_ui_account(&pubkey, &account, encoding, None, None)),
1844            );
1845            if let Some(data_account) = data_account {
1846                capture.insert(
1847                    data_pubkey,
1848                    Some(encode_ui_account(
1849                        &data_pubkey,
1850                        &data_account,
1851                        encoding,
1852                        None,
1853                        None,
1854                    )),
1855                );
1856            }
1857        }
1858    }
1859}
1860
1861// Helper function to apply filters
1862fn apply_rpc_filters(account_data: &[u8], filters: &[RpcFilterType]) -> SurfpoolResult<bool> {
1863    for filter in filters {
1864        match filter {
1865            RpcFilterType::DataSize(size) => {
1866                if account_data.len() as u64 != *size {
1867                    return Ok(false);
1868                }
1869            }
1870            RpcFilterType::Memcmp(memcmp_filter) => {
1871                // Use the public bytes_match method from solana_client::rpc_filter::Memcmp
1872                if !memcmp_filter.bytes_match(account_data) {
1873                    return Ok(false); // Content mismatch or out of bounds handled by bytes_match
1874                }
1875            }
1876            RpcFilterType::TokenAccountState => {
1877                return Err(SurfpoolError::internal(
1878                    "TokenAccountState filter is not supported",
1879                ));
1880            }
1881        }
1882    }
1883    Ok(true)
1884}
1885
1886// used in the remote.rs
1887pub fn is_supported_token_program(program_id: &Pubkey) -> bool {
1888    *program_id == spl_token::ID || *program_id == spl_token_2022::ID
1889}
1890
1891fn update_programdata_account(
1892    program_id: &Pubkey,
1893    programdata_account: &mut Account,
1894    new_authority: Option<Pubkey>,
1895) -> SurfpoolResult<Option<Pubkey>> {
1896    let upgradeable_loader_state =
1897        bincode::deserialize::<UpgradeableLoaderState>(&programdata_account.data).map_err(|e| {
1898            SurfpoolError::invalid_program_account(
1899                &program_id,
1900                format!("Failed to serialize program data: {}", e),
1901            )
1902        })?;
1903    if let UpgradeableLoaderState::ProgramData {
1904        upgrade_authority_address,
1905        slot,
1906    } = upgradeable_loader_state
1907    {
1908        let offset = if upgrade_authority_address.is_some() {
1909            UpgradeableLoaderState::size_of_programdata_metadata()
1910        } else {
1911            UpgradeableLoaderState::size_of_programdata_metadata()
1912                - serialized_size(&Pubkey::default()).unwrap() as usize
1913        };
1914
1915        let mut data = bincode::serialize(&UpgradeableLoaderState::ProgramData {
1916            upgrade_authority_address: new_authority,
1917            slot,
1918        })
1919        .map_err(|e| {
1920            SurfpoolError::invalid_program_account(
1921                &program_id,
1922                format!("Failed to serialize program data: {}", e),
1923            )
1924        })?;
1925
1926        data.append(&mut programdata_account.data[offset..].to_vec());
1927
1928        programdata_account.data = data;
1929
1930        Ok(upgrade_authority_address)
1931    } else {
1932        return Err(SurfpoolError::invalid_program_account(
1933            &program_id,
1934            "Invalid program data account",
1935        ));
1936    }
1937}