Skip to main content

simulator_client/
builders.rs

1use std::collections::{BTreeMap, BTreeSet};
2
3use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
4use bon::Builder;
5use simulator_api::{
6    AccountData, AccountModifications, BinaryEncoding, ContinueParams, CreateSessionParams,
7};
8use solana_address::Address;
9use solana_client::rpc_client::SerializableTransaction;
10use thiserror::Error;
11
12use crate::BacktestClientError;
13
14/// Error serializing a transaction for injection.
15#[derive(Debug, Error)]
16pub enum SerializeEncodeError {
17    #[error("bincode serialization error: {0}")]
18    Bincode(#[from] bincode::error::EncodeError),
19}
20
21fn serialize_to_base64(value: &impl serde::Serialize) -> Result<String, SerializeEncodeError> {
22    let bytes = bincode::serde::encode_to_vec(
23        value,
24        bincode::config::standard()
25            .with_fixed_int_encoding()
26            .with_little_endian(),
27    )?;
28    Ok(BASE64.encode(&bytes))
29}
30
31/// Builder for `CreateBacktestSession`.
32///
33/// Set either `end_slot` or `slot_count` (not both). Use the helper methods to
34/// add account filters or preload programs.
35#[derive(Debug, Clone, Builder)]
36pub struct CreateSession {
37    pub start_slot: u64,
38
39    pub end_slot: Option<u64>,
40
41    pub slot_count: Option<u64>,
42
43    #[builder(default)]
44    pub signer_filter: BTreeSet<Address>,
45
46    #[builder(default)]
47    pub preload_programs: BTreeSet<Address>,
48
49    #[builder(default)]
50    pub preload_account_bundles: Vec<String>,
51
52    /// When true, include a session summary with transaction statistics in the Completed response.
53    #[builder(default)]
54    pub send_summary: bool,
55
56    pub disconnect_timeout_secs: Option<u16>,
57}
58
59impl CreateSession {
60    /// Add an account to the signer filter.
61    pub fn add_signer_filter(mut self, address: Address) -> Self {
62        self.signer_filter.insert(address);
63        self
64    }
65
66    /// Add a program to preload before the first continue.
67    pub fn add_preload_program(mut self, address: Address) -> Self {
68        self.preload_programs.insert(address);
69        self
70    }
71
72    /// Convert the builder into API parameters, validating slot options.
73    pub fn into_params(self) -> Result<CreateSessionParams, BacktestClientError> {
74        let end_slot = match (self.end_slot, self.slot_count) {
75            (Some(_), Some(_)) => {
76                return Err(BacktestClientError::InvalidParams {
77                    message: "CreateSession: set only one of end_slot or slot_count".to_string(),
78                });
79            }
80            (Some(end_slot), None) => end_slot,
81            (None, Some(slot_count)) => {
82                self.start_slot.checked_add(slot_count).ok_or_else(|| {
83                    BacktestClientError::InvalidParams {
84                        message: "CreateSession: start_slot + slot_count overflow".to_string(),
85                    }
86                })?
87            }
88            (None, None) => {
89                return Err(BacktestClientError::InvalidParams {
90                    message: "CreateSession: must set end_slot or slot_count".to_string(),
91                });
92            }
93        };
94
95        if end_slot < self.start_slot {
96            return Err(BacktestClientError::InvalidParams {
97                message: format!(
98                    "CreateSession: end_slot ({end_slot}) must be >= start_slot ({})",
99                    self.start_slot
100                ),
101            });
102        }
103
104        Ok(CreateSessionParams {
105            start_slot: self.start_slot,
106            end_slot,
107            signer_filter: self.signer_filter,
108            preload_programs: self.preload_programs,
109            preload_account_bundles: self.preload_account_bundles,
110            send_summary: self.send_summary,
111            disconnect_timeout_secs: self.disconnect_timeout_secs,
112        })
113    }
114}
115
116/// Builder for `Continue` requests.
117///
118/// Use this to advance the simulation, inject transactions, or patch accounts.
119#[derive(Debug, Builder)]
120pub struct Continue {
121    #[builder(default = ContinueParams::default_advance_count())]
122    pub advance_count: u64,
123
124    #[builder(default)]
125    pub transactions: Vec<String>,
126
127    #[builder(default)]
128    pub modify_accounts: BTreeMap<Address, AccountData>,
129}
130
131impl Continue {
132    /// Append a base64-encoded transaction payload.
133    pub fn push_transaction_base64(mut self, data: impl Into<String>) -> Self {
134        self.transactions.push(data.into());
135        self
136    }
137
138    /// Append a raw transaction payload encoded as base64.
139    pub fn push_transaction_bytes(mut self, bytes: &[u8]) -> Self {
140        self.transactions.push(BinaryEncoding::Base64.encode(bytes));
141        self
142    }
143
144    /// Append a serializable transaction encoded as base64.
145    pub fn push_transaction(
146        mut self,
147        transaction: &impl SerializableTransaction,
148    ) -> Result<Self, SerializeEncodeError> {
149        self.transactions.push(serialize_to_base64(&transaction)?);
150        Ok(self)
151    }
152
153    /// Modify an account state prior to execution.
154    pub fn modify_account(mut self, address: Address, account: AccountData) -> Self {
155        self.modify_accounts.insert(address, account);
156        self
157    }
158
159    /// Convert the builder into API parameters.
160    pub fn into_params(self) -> ContinueParams {
161        ContinueParams {
162            advance_count: self.advance_count,
163            transactions: self.transactions,
164            modify_account_states: AccountModifications(self.modify_accounts),
165        }
166    }
167}