solana_token_toolkit/
state.rs

1//! Tier 1 — on-chain state queries for an owner's token accounts.
2
3use std::collections::HashMap;
4
5use solana_account::Account;
6use solana_client::nonblocking::rpc_client::RpcClient;
7use solana_pubkey::Pubkey;
8use spl_associated_token_account_interface::address::get_associated_token_address_with_program_id;
9use spl_token_interface::native_mint;
10
11use crate::TokenError;
12
13/// On-chain state for an owner's token accounts across a set of mints.
14#[derive(Debug, Clone)]
15pub struct TokenAccountState {
16    /// The owner whose ATAs were queried.
17    pub owner: Pubkey,
18    /// Map of mint pubkey to per-mint state.
19    pub mints: HashMap<Pubkey, MintAndAta>,
20}
21
22impl TokenAccountState {
23    /// Build an empty state (no mints) for the given owner.
24    #[must_use]
25    pub fn empty(owner: Pubkey) -> Self {
26        Self {
27            owner,
28            mints: HashMap::new(),
29        }
30    }
31}
32
33/// Per-mint state: the mint account itself, the ATA address, and the
34/// ATA account if it currently exists on chain.
35#[derive(Debug, Clone)]
36pub struct MintAndAta {
37    /// The raw mint account (use `spl-token-2022-interface` to unpack).
38    pub mint_account: Account,
39    /// Derived ATA address for `(owner, mint, mint_account.owner)`.
40    pub ata_address: Pubkey,
41    /// `Some` if the ATA exists on chain, `None` if it needs creation.
42    pub ata_account: Option<Account>,
43}
44
45/// Fetch the on-chain state for an owner across a set of mints. Performs
46/// two RPC `get_multiple_accounts` calls (mints, then ATAs).
47///
48/// Returns `Err(TokenError::MintNotFound(mint))` if any requested mint
49/// account does not exist on chain.
50///
51/// # Example
52///
53/// ```no_run
54/// # use solana_token_toolkit::{fetch_token_account_state, TokenError};
55/// # use solana_client::nonblocking::rpc_client::RpcClient;
56/// # use solana_pubkey::Pubkey;
57/// # async fn run() -> Result<(), TokenError> {
58/// let rpc = RpcClient::new("http://localhost:8899".to_string());
59/// let owner = Pubkey::new_unique();
60/// let mints = vec![Pubkey::new_unique()];
61/// let state = fetch_token_account_state(&rpc, owner, &mints).await?;
62/// # let _ = state;
63/// # Ok(())
64/// # }
65/// ```
66pub async fn fetch_token_account_state(
67    rpc: &RpcClient,
68    owner: Pubkey,
69    mints: &[Pubkey],
70) -> Result<TokenAccountState, TokenError> {
71    if mints.is_empty() {
72        return Ok(TokenAccountState::empty(owner));
73    }
74
75    let mint_account_opts = rpc.get_multiple_accounts(mints).await?;
76    let mut mint_accounts = Vec::with_capacity(mints.len());
77    for (i, opt) in mint_account_opts.into_iter().enumerate() {
78        match (mints[i] == native_mint::ID, opt) {
79            (_, Some(account)) => mint_accounts.push(account),
80            (true, None) => mint_accounts.push(native_mint_account()),
81            (false, None) => return Err(TokenError::MintNotFound(mints[i])),
82        }
83    }
84
85    assemble_token_account_state(rpc, owner, mints, &mint_accounts).await
86}
87
88fn native_mint_account() -> Account {
89    use solana_program_pack::Pack;
90    use spl_token::state::Mint;
91
92    let mint = Mint {
93        mint_authority: spl_token::solana_program::program_option::COption::None,
94        supply: 0,
95        decimals: native_mint::DECIMALS,
96        is_initialized: true,
97        freeze_authority: spl_token::solana_program::program_option::COption::None,
98    };
99    let mut data = vec![0u8; Mint::LEN];
100    Mint::pack(mint, &mut data).expect("native mint account serialization is infallible");
101
102    Account {
103        lamports: 0,
104        data,
105        owner: spl_token::ID,
106        executable: false,
107        rent_epoch: 0,
108    }
109}
110
111/// Assemble token account state from caller-provided mint accounts. Performs
112/// one RPC `get_multiple_accounts` call (ATAs only).
113///
114/// `mints` and `mint_accounts` must be aligned by index.
115///
116/// # Example
117///
118/// ```no_run
119/// # use solana_token_toolkit::{assemble_token_account_state, TokenError};
120/// # use solana_client::nonblocking::rpc_client::RpcClient;
121/// # use solana_pubkey::Pubkey;
122/// # async fn run() -> Result<(), TokenError> {
123/// let rpc = RpcClient::new("http://localhost:8899".to_string());
124/// let owner = Pubkey::new_unique();
125/// let mints = vec![Pubkey::new_unique()];
126/// let mint_accounts = rpc.get_multiple_accounts(&mints).await?
127///     .into_iter()
128///     .collect::<Option<Vec<_>>>()
129///     .ok_or(TokenError::MintNotFound(mints[0]))?;
130/// let state = assemble_token_account_state(&rpc, owner, &mints, &mint_accounts).await?;
131/// # let _ = state;
132/// # Ok(())
133/// # }
134/// ```
135pub async fn assemble_token_account_state(
136    rpc: &RpcClient,
137    owner: Pubkey,
138    mints: &[Pubkey],
139    mint_accounts: &[Account],
140) -> Result<TokenAccountState, TokenError> {
141    if mints.len() != mint_accounts.len() {
142        return Err(TokenError::LengthMismatch {
143            mints: mints.len(),
144            mint_accounts: mint_accounts.len(),
145        });
146    }
147    if mints.is_empty() {
148        return Ok(TokenAccountState::empty(owner));
149    }
150
151    let ata_addresses: Vec<Pubkey> = mints
152        .iter()
153        .zip(mint_accounts.iter())
154        .map(|(mint, acc)| get_associated_token_address_with_program_id(&owner, mint, &acc.owner))
155        .collect();
156
157    let ata_account_opts = rpc.get_multiple_accounts(&ata_addresses).await?;
158
159    let mints_map = mints
160        .iter()
161        .zip(mint_accounts.iter())
162        .zip(ata_addresses.iter())
163        .zip(ata_account_opts)
164        .map(|(((mint, mint_account), ata_address), ata_account)| {
165            (
166                *mint,
167                MintAndAta {
168                    mint_account: mint_account.clone(),
169                    ata_address: *ata_address,
170                    ata_account,
171                },
172            )
173        })
174        .collect();
175
176    Ok(TokenAccountState {
177        owner,
178        mints: mints_map,
179    })
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185
186    #[tokio::test]
187    async fn empty_mints_returns_empty_state_no_rpc() {
188        let rpc = RpcClient::new_mock("succeeds".to_string());
189        let owner = Pubkey::new_unique();
190        let state = assemble_token_account_state(&rpc, owner, &[], &[])
191            .await
192            .unwrap();
193        assert_eq!(state.owner, owner);
194        assert!(state.mints.is_empty());
195    }
196
197    #[tokio::test]
198    async fn length_mismatch_returns_error_no_rpc() {
199        let rpc = RpcClient::new_mock("succeeds".to_string());
200        let owner = Pubkey::new_unique();
201        let mints = vec![Pubkey::new_unique(), Pubkey::new_unique()];
202        let mint_accounts = vec![Account {
203            lamports: 0,
204            data: vec![],
205            owner: spl_token::ID,
206            executable: false,
207            rent_epoch: 0,
208        }];
209        let err = assemble_token_account_state(&rpc, owner, &mints, &mint_accounts)
210            .await
211            .unwrap_err();
212        match err {
213            TokenError::LengthMismatch {
214                mints: m,
215                mint_accounts: ma,
216            } => {
217                assert_eq!(m, 2);
218                assert_eq!(ma, 1);
219            }
220            _ => panic!("expected LengthMismatch, got {err:?}"),
221        }
222    }
223
224    #[test]
225    fn empty_state_constructor() {
226        let owner = Pubkey::new_unique();
227        let state = TokenAccountState::empty(owner);
228        assert_eq!(state.owner, owner);
229        assert!(state.mints.is_empty());
230    }
231
232    #[tokio::test]
233    async fn fetch_empty_mints_returns_empty_state_no_rpc() {
234        let rpc = RpcClient::new_mock("succeeds".to_string());
235        let owner = Pubkey::new_unique();
236        let state = fetch_token_account_state(&rpc, owner, &[]).await.unwrap();
237        assert_eq!(state.owner, owner);
238        assert!(state.mints.is_empty());
239    }
240
241    #[test]
242    fn native_mint_account_is_valid_classic_mint_owned_by_spl_token() {
243        use solana_program_pack::Pack;
244        use spl_token::state::Mint;
245
246        let account = native_mint_account();
247        assert_eq!(account.owner, spl_token::ID);
248        let mint = Mint::unpack(&account.data).unwrap();
249        assert!(mint.is_initialized);
250        assert_eq!(mint.decimals, native_mint::DECIMALS);
251    }
252}
solana_token_toolkit/state.rs

solana_token_toolkit/
state.rs
