hodor_program/swap/
state.rs

1use arrayref::{array_mut_ref, array_ref, array_refs, mut_array_refs};
2use solana_program::program_error::ProgramError;
3use solana_program::pubkey::Pubkey;
4
5#[derive(Debug, PartialEq)]
6pub struct CreatorFee {
7    pub rate: u32,
8    pub balance_a: u64,
9    pub balance_b: u64,
10    pub withdraw_authority: Pubkey,
11}
12
13#[derive(Debug, PartialEq)]
14pub struct SwapPool {
15    pub seed: [u8; 32],
16    pub token_account_a: Pubkey,
17    pub token_account_b: Pubkey,
18    pub balance_a: u64,
19    pub balance_b: u64,
20    pub lp_mint: Pubkey,
21    pub lp_fee_rate: u32,
22    pub creator_fee: Option<CreatorFee>,
23}
24
25
26impl SwapPool {
27    pub const BASE_SIZE: usize = 1 + 32 + 32 + 32 + 8 + 8 + 32 + 4;
28    pub const CREATOR_FEE_SIZE: usize = 4 + 8 + 8 + 32;
29    pub const WITH_CREATOR_FEE_SIZE: usize = SwapPool::BASE_SIZE + SwapPool::CREATOR_FEE_SIZE;
30    pub const TYPE_MARKER: u8 = 1;
31
32    pub fn pack(&self, dst: &mut [u8]) -> Result<(), ProgramError> {
33        if (self.creator_fee.is_none() && dst.len() != SwapPool::BASE_SIZE) ||
34            (self.creator_fee.is_some() && dst.len() != SwapPool::WITH_CREATOR_FEE_SIZE) {
35            return Err(ProgramError::InvalidAccountData);
36        }
37
38        let dst_ref = array_mut_ref![dst, 0, SwapPool::BASE_SIZE];
39        let (type_marker_dst, seed_dst, token_acc_a_dst, token_acc_b_dst, balance_a_dst, balance_b_dst, lp_mint_dst, lp_fee_rate_dst)
40            = mut_array_refs![dst_ref, 1, 32, 32, 32, 8, 8, 32, 4];
41
42        *type_marker_dst = [SwapPool::TYPE_MARKER];
43        seed_dst.copy_from_slice(self.seed.as_ref());
44        token_acc_a_dst.copy_from_slice(self.token_account_a.as_ref());
45        token_acc_b_dst.copy_from_slice(self.token_account_b.as_ref());
46        *balance_a_dst = self.balance_a.to_le_bytes();
47        *balance_b_dst = self.balance_b.to_le_bytes();
48        lp_mint_dst.copy_from_slice(self.lp_mint.as_ref());
49        *lp_fee_rate_dst = self.lp_fee_rate.to_le_bytes();
50
51        if let Some(creator_fee) = &self.creator_fee {
52            let dst_ref = array_mut_ref![dst, SwapPool::BASE_SIZE, SwapPool::CREATOR_FEE_SIZE];
53
54            let (rate_dst, balance_a_dst, balance_b_dst, withdraw_authority_dst)
55                = mut_array_refs![dst_ref, 4, 8, 8, 32];
56
57            *rate_dst = creator_fee.rate.to_le_bytes();
58            *balance_a_dst = creator_fee.balance_a.to_le_bytes();
59            *balance_b_dst = creator_fee.balance_b.to_le_bytes();
60            withdraw_authority_dst.copy_from_slice(creator_fee.withdraw_authority.as_ref());
61        }
62
63        Ok(())
64    }
65
66    pub fn unpack(src: &[u8]) -> Result<Self, ProgramError> {
67        // todo: size check
68
69        let src_array_ref = array_ref![src, 0, SwapPool::BASE_SIZE];
70        let (type_marker, seed, token_acc_a, token_acc_b,
71            balance_a, balance_b, lp_mint, lp_fee_rate)
72            = array_refs![src_array_ref, 1, 32, 32, 32, 8, 8, 32, 4];
73
74        if *type_marker != [SwapPool::TYPE_MARKER] {
75            return Err(ProgramError::InvalidAccountData);
76        }
77
78        let creator_fee = if src.len() == SwapPool::WITH_CREATOR_FEE_SIZE {
79            let src_array_ref = array_ref![src, SwapPool::BASE_SIZE, SwapPool::CREATOR_FEE_SIZE];
80            let (rate, balance_a, balance_b, withdraw_authority)
81                = array_refs![src_array_ref, 4, 8, 8, 32];
82
83            Some(CreatorFee {
84                rate: u32::from_le_bytes(*rate),
85                balance_a: u64::from_le_bytes(*balance_a),
86                balance_b: u64::from_le_bytes(*balance_b),
87                withdraw_authority: Pubkey::new_from_array(*withdraw_authority),
88            })
89        } else {
90            None
91        };
92
93        Ok(SwapPool {
94            seed: *seed, // todo: should we clone ?
95            token_account_a: Pubkey::new_from_array(*token_acc_a),
96            token_account_b: Pubkey::new_from_array(*token_acc_b),
97            balance_a: u64::from_le_bytes(*balance_a),
98            balance_b: u64::from_le_bytes(*balance_b),
99            lp_mint: Pubkey::new_from_array(*lp_mint),
100            lp_fee_rate: u32::from_le_bytes(*lp_fee_rate),
101            creator_fee,
102        })
103    }
104}
105
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[test]
112    fn test_swap_pool_pack_unpack() {
113        let pool_without_creator_fee = SwapPool {
114            seed: Pubkey::new_unique().to_bytes(),
115            token_account_a: Pubkey::new_unique(),
116            token_account_b: Pubkey::new_unique(),
117            balance_a: 100,
118            balance_b: 150,
119            lp_mint: Pubkey::new_unique(),
120            lp_fee_rate: 5_000,
121            creator_fee: None,
122        };
123        let mut state_array = [0u8; SwapPool::BASE_SIZE];
124        pool_without_creator_fee.pack(&mut state_array).unwrap();
125        assert_eq!(pool_without_creator_fee, SwapPool::unpack(&state_array).unwrap());
126        assert!(pool_without_creator_fee.pack(&mut [0u8; SwapPool::WITH_CREATOR_FEE_SIZE]).is_err());
127
128
129        let pool_with_creator_fee = SwapPool {
130            seed: Pubkey::new_unique().to_bytes(),
131            token_account_a: Pubkey::new_unique(),
132            token_account_b: Pubkey::new_unique(),
133            balance_a: 0,
134            balance_b: 120,
135            lp_mint: Pubkey::new_unique(),
136            lp_fee_rate: 5_000,
137            creator_fee: Some(CreatorFee {
138                rate: 10_000,
139                balance_a: 5_000,
140                balance_b: 6_000,
141                withdraw_authority: Default::default()
142            }),
143        };
144        let mut state_array = [0u8; SwapPool::WITH_CREATOR_FEE_SIZE];
145        pool_with_creator_fee.pack(&mut state_array).unwrap();
146        assert_eq!(pool_with_creator_fee, SwapPool::unpack(&state_array).unwrap());
147        assert!(pool_with_creator_fee.pack(&mut [0u8; SwapPool::BASE_SIZE]).is_err());
148    }
149}