light_sdk/instruction/
pack_accounts.rs

1use std::collections::HashMap;
2
3use crate::{
4    instruction::system_accounts::{get_light_system_account_metas, SystemAccountMetaConfig},
5    AccountMeta, Pubkey,
6};
7
8#[derive(Default, Debug)]
9pub struct PackedAccounts {
10    pre_accounts: Vec<AccountMeta>,
11    system_accounts: Vec<AccountMeta>,
12    next_index: u8,
13    map: HashMap<Pubkey, (u8, AccountMeta)>,
14}
15
16impl PackedAccounts {
17    pub fn new_with_system_accounts(config: SystemAccountMetaConfig) -> Self {
18        let mut remaining_accounts = PackedAccounts::default();
19        remaining_accounts.add_system_accounts(config);
20        remaining_accounts
21    }
22
23    pub fn add_pre_accounts_signer(&mut self, pubkey: Pubkey) {
24        self.pre_accounts.push(AccountMeta {
25            pubkey,
26            is_signer: true,
27            is_writable: false,
28        });
29    }
30
31    pub fn add_pre_accounts_signer_mut(&mut self, pubkey: Pubkey) {
32        self.pre_accounts.push(AccountMeta {
33            pubkey,
34            is_signer: true,
35            is_writable: true,
36        });
37    }
38
39    pub fn add_pre_accounts_meta(&mut self, account_meta: AccountMeta) {
40        self.pre_accounts.push(account_meta);
41    }
42
43    pub fn add_system_accounts(&mut self, config: SystemAccountMetaConfig) {
44        self.system_accounts
45            .extend(get_light_system_account_metas(config));
46    }
47
48    /// Returns the index of the provided `pubkey` in the collection.
49    ///
50    /// If the provided `pubkey` is not a part of the collection, it gets
51    /// inserted with a `next_index`.
52    ///
53    /// If the privided `pubkey` already exists in the collection, its already
54    /// existing index is returned.
55    pub fn insert_or_get(&mut self, pubkey: Pubkey) -> u8 {
56        self.insert_or_get_config(pubkey, false, true)
57    }
58
59    pub fn insert_or_get_read_only(&mut self, pubkey: Pubkey) -> u8 {
60        self.insert_or_get_config(pubkey, false, false)
61    }
62
63    pub fn insert_or_get_config(
64        &mut self,
65        pubkey: Pubkey,
66        is_signer: bool,
67        is_writable: bool,
68    ) -> u8 {
69        self.map
70            .entry(pubkey)
71            .or_insert_with(|| {
72                let index = self.next_index;
73                self.next_index += 1;
74                (
75                    index,
76                    AccountMeta {
77                        pubkey,
78                        is_signer,
79                        is_writable,
80                    },
81                )
82            })
83            .0
84    }
85
86    fn hash_set_accounts_to_metas(&self) -> Vec<AccountMeta> {
87        let mut packed_accounts = self.map.iter().collect::<Vec<_>>();
88        // hash maps are not sorted so we need to sort manually and collect into a vector again
89        packed_accounts.sort_by(|a, b| a.1 .0.cmp(&b.1 .0));
90        let packed_accounts = packed_accounts
91            .iter()
92            .map(|(_, (_, k))| k.clone())
93            .collect::<Vec<AccountMeta>>();
94        packed_accounts
95    }
96
97    fn get_offsets(&self) -> (usize, usize) {
98        let system_accounts_start_offset = self.pre_accounts.len();
99        let packed_accounts_start_offset =
100            system_accounts_start_offset + self.system_accounts.len();
101        (system_accounts_start_offset, packed_accounts_start_offset)
102    }
103
104    /// Converts the collection of accounts to a vector of
105    /// [`AccountMeta`](solana_instruction::AccountMeta), which can be used
106    /// as remaining accounts in instructions or CPI calls.
107    pub fn to_account_metas(&self) -> (Vec<AccountMeta>, usize, usize) {
108        let packed_accounts = self.hash_set_accounts_to_metas();
109        let (system_accounts_start_offset, packed_accounts_start_offset) = self.get_offsets();
110        (
111            [
112                self.pre_accounts.clone(),
113                self.system_accounts.clone(),
114                packed_accounts,
115            ]
116            .concat(),
117            system_accounts_start_offset,
118            packed_accounts_start_offset,
119        )
120    }
121}
122
123#[cfg(test)]
124mod test {
125    use super::*;
126
127    #[test]
128    fn test_remaining_accounts() {
129        let mut remaining_accounts = PackedAccounts::default();
130
131        let pubkey_1 = Pubkey::new_unique();
132        let pubkey_2 = Pubkey::new_unique();
133        let pubkey_3 = Pubkey::new_unique();
134        let pubkey_4 = Pubkey::new_unique();
135
136        // Initial insertion.
137        assert_eq!(remaining_accounts.insert_or_get(pubkey_1), 0);
138        assert_eq!(remaining_accounts.insert_or_get(pubkey_2), 1);
139        assert_eq!(remaining_accounts.insert_or_get(pubkey_3), 2);
140
141        assert_eq!(
142            remaining_accounts.to_account_metas().0.as_slice(),
143            &[
144                AccountMeta {
145                    pubkey: pubkey_1,
146                    is_signer: false,
147                    is_writable: true,
148                },
149                AccountMeta {
150                    pubkey: pubkey_2,
151                    is_signer: false,
152                    is_writable: true,
153                },
154                AccountMeta {
155                    pubkey: pubkey_3,
156                    is_signer: false,
157                    is_writable: true,
158                }
159            ]
160        );
161
162        // Insertion of already existing pubkeys.
163        assert_eq!(remaining_accounts.insert_or_get(pubkey_1), 0);
164        assert_eq!(remaining_accounts.insert_or_get(pubkey_2), 1);
165        assert_eq!(remaining_accounts.insert_or_get(pubkey_3), 2);
166
167        assert_eq!(
168            remaining_accounts.to_account_metas().0.as_slice(),
169            &[
170                AccountMeta {
171                    pubkey: pubkey_1,
172                    is_signer: false,
173                    is_writable: true,
174                },
175                AccountMeta {
176                    pubkey: pubkey_2,
177                    is_signer: false,
178                    is_writable: true,
179                },
180                AccountMeta {
181                    pubkey: pubkey_3,
182                    is_signer: false,
183                    is_writable: true,
184                }
185            ]
186        );
187
188        // Again, initial insertion.
189        assert_eq!(remaining_accounts.insert_or_get(pubkey_4), 3);
190
191        assert_eq!(
192            remaining_accounts.to_account_metas().0.as_slice(),
193            &[
194                AccountMeta {
195                    pubkey: pubkey_1,
196                    is_signer: false,
197                    is_writable: true,
198                },
199                AccountMeta {
200                    pubkey: pubkey_2,
201                    is_signer: false,
202                    is_writable: true,
203                },
204                AccountMeta {
205                    pubkey: pubkey_3,
206                    is_signer: false,
207                    is_writable: true,
208                },
209                AccountMeta {
210                    pubkey: pubkey_4,
211                    is_signer: false,
212                    is_writable: true,
213                }
214            ]
215        );
216    }
217}