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

solana_token_toolkit/
state.rs
