gmsol_treasury/states/
treasury.rs

1use anchor_lang::prelude::*;
2use bytemuck::Zeroable;
3use gmsol_store::{states::Seed, utils::pubkey::to_bytes, CoreError};
4
5pub use gmsol_utils::token_config::{TokenFlag, MAX_TREASURY_TOKEN_FLAGS};
6
7pub(crate) const MAX_TOKENS: usize = 16;
8
9/// Treasury vault config account.
10#[account(zero_copy)]
11#[cfg_attr(feature = "debug", derive(derive_more::Debug))]
12pub struct TreasuryVaultConfig {
13    version: u8,
14    pub(crate) bump: u8,
15    index: u16,
16    #[cfg_attr(feature = "debug", debug(skip))]
17    padding: [u8; 12],
18    pub(crate) config: Pubkey,
19    #[cfg_attr(feature = "debug", debug(skip))]
20    reserved: [u8; 256],
21    tokens: TokenMap,
22}
23
24impl Seed for TreasuryVaultConfig {
25    const SEED: &'static [u8] = b"treasury_vault_config";
26}
27
28impl gmsol_utils::InitSpace for TreasuryVaultConfig {
29    const INIT_SPACE: usize = std::mem::size_of::<Self>();
30}
31
32impl TreasuryVaultConfig {
33    pub(crate) fn init(&mut self, bump: u8, index: u16, config: &Pubkey) {
34        self.bump = bump;
35        self.index = index;
36        self.config = *config;
37    }
38
39    pub(crate) fn insert_token(&mut self, token: &Pubkey) -> Result<()> {
40        self.tokens
41            .insert_with_options(token, TokenConfig::default(), true)?;
42        Ok(())
43    }
44
45    pub(crate) fn remove_token(&mut self, token: &Pubkey) -> Result<()> {
46        self.tokens
47            .remove(token)
48            .ok_or_else(|| error!(CoreError::NotFound))?;
49        Ok(())
50    }
51
52    pub(crate) fn toggle_token_flag(
53        &mut self,
54        token: &Pubkey,
55        flag: TokenFlag,
56        value: bool,
57    ) -> Result<bool> {
58        let config = self
59            .tokens
60            .get_mut(token)
61            .ok_or_else(|| error!(CoreError::NotFound))?;
62
63        require_neq!(
64            config.flags.get_flag(flag),
65            value,
66            CoreError::PreconditionsAreNotMet
67        );
68
69        Ok(config.flags.set_flag(flag, value))
70    }
71
72    fn get_token_config(&self, token: &Pubkey) -> Result<&TokenConfig> {
73        self.tokens
74            .get(token)
75            .ok_or_else(|| error!(CoreError::NotFound))
76    }
77
78    pub(crate) fn is_deposit_allowed(&self, token: &Pubkey) -> Result<bool> {
79        Ok(self
80            .get_token_config(token)?
81            .flags
82            .get_flag(TokenFlag::AllowDeposit))
83    }
84
85    pub(crate) fn is_withdrawal_allowed(&self, token: &Pubkey) -> Result<bool> {
86        Ok(self
87            .get_token_config(token)?
88            .flags
89            .get_flag(TokenFlag::AllowWithdrawal))
90    }
91
92    pub(crate) fn signer(&self) -> TreasuryVaultSigner {
93        TreasuryVaultSigner {
94            config: self.config,
95            index_bytes: self.index.to_le_bytes(),
96            bump_bytes: [self.bump],
97        }
98    }
99
100    /// Get the number of tokens.
101    pub fn num_tokens(&self) -> usize {
102        self.tokens.len()
103    }
104
105    /// Get all tokens.
106    pub fn tokens(&self) -> impl Iterator<Item = Pubkey> + '_ {
107        self.tokens
108            .entries()
109            .map(|(key, _)| Pubkey::new_from_array(*key))
110    }
111}
112
113/// Treasury Vault Signer.
114pub struct TreasuryVaultSigner {
115    config: Pubkey,
116    index_bytes: [u8; 2],
117    bump_bytes: [u8; 1],
118}
119
120impl TreasuryVaultSigner {
121    /// As signer seeds.
122    pub fn as_seeds(&self) -> [&[u8]; 4] {
123        [
124            TreasuryVaultConfig::SEED,
125            self.config.as_ref(),
126            &self.index_bytes,
127            &self.bump_bytes,
128        ]
129    }
130}
131
132/// Token config for treasury.
133#[zero_copy]
134#[cfg_attr(feature = "debug", derive(Debug))]
135pub struct TokenConfig {
136    flags: TokenFlagContainer,
137    reserved: [u8; 64],
138}
139
140impl Default for TokenConfig {
141    fn default() -> Self {
142        Self::zeroed()
143    }
144}
145
146gmsol_utils::fixed_map!(TokenMap, Pubkey, to_bytes, TokenConfig, MAX_TOKENS, 0);
147
148gmsol_utils::flags!(TokenFlag, MAX_TREASURY_TOKEN_FLAGS, u8);