waffles_solana_program/message/
account_keys.rs

1use {
2    crate::{
3        instruction::{CompiledInstruction, Instruction},
4        message::{v0::LoadedAddresses, CompileError},
5        pubkey::Pubkey,
6    },
7    std::{collections::BTreeMap, ops::Index},
8};
9
10/// Collection of static and dynamically loaded keys used to load accounts
11/// during transaction processing.
12pub struct AccountKeys<'a> {
13    static_keys: &'a [Pubkey],
14    dynamic_keys: Option<&'a LoadedAddresses>,
15}
16
17impl Index<usize> for AccountKeys<'_> {
18    type Output = Pubkey;
19    fn index(&self, index: usize) -> &Self::Output {
20        self.get(index).expect("index is invalid")
21    }
22}
23
24impl<'a> AccountKeys<'a> {
25    pub fn new(static_keys: &'a [Pubkey], dynamic_keys: Option<&'a LoadedAddresses>) -> Self {
26        Self {
27            static_keys,
28            dynamic_keys,
29        }
30    }
31
32    /// Returns an iterator of account key segments. The ordering of segments
33    /// affects how account indexes from compiled instructions are resolved and
34    /// so should not be changed.
35    fn key_segment_iter(&self) -> impl Iterator<Item = &'a [Pubkey]> {
36        if let Some(dynamic_keys) = self.dynamic_keys {
37            [
38                self.static_keys,
39                &dynamic_keys.writable,
40                &dynamic_keys.readonly,
41            ]
42            .into_iter()
43        } else {
44            // empty segments added for branch type compatibility
45            [self.static_keys, &[], &[]].into_iter()
46        }
47    }
48
49    /// Returns the address of the account at the specified index of the list of
50    /// message account keys constructed from static keys, followed by dynamically
51    /// loaded writable addresses, and lastly the list of dynamically loaded
52    /// readonly addresses.
53    pub fn get(&self, mut index: usize) -> Option<&'a Pubkey> {
54        for key_segment in self.key_segment_iter() {
55            if index < key_segment.len() {
56                return Some(&key_segment[index]);
57            }
58            index = index.saturating_sub(key_segment.len());
59        }
60
61        None
62    }
63
64    /// Returns the total length of loaded accounts for a message
65    pub fn len(&self) -> usize {
66        let mut len = 0usize;
67        for key_segment in self.key_segment_iter() {
68            len = len.saturating_add(key_segment.len());
69        }
70        len
71    }
72
73    /// Returns true if this collection of account keys is empty
74    pub fn is_empty(&self) -> bool {
75        self.len() == 0
76    }
77
78    /// Iterator for the addresses of the loaded accounts for a message
79    pub fn iter(&self) -> impl Iterator<Item = &'a Pubkey> {
80        self.key_segment_iter().flatten()
81    }
82
83    /// Compile instructions using the order of account keys to determine
84    /// compiled instruction account indexes.
85    ///
86    /// # Panics
87    ///
88    /// Panics when compiling fails. See [`AccountKeys::try_compile_instructions`]
89    /// for a full description of failure scenarios.
90    pub fn compile_instructions(&self, instructions: &[Instruction]) -> Vec<CompiledInstruction> {
91        self.try_compile_instructions(instructions)
92            .expect("compilation failure")
93    }
94
95    /// Compile instructions using the order of account keys to determine
96    /// compiled instruction account indexes.
97    ///
98    /// # Errors
99    ///
100    /// Compilation will fail if any `instructions` use account keys which are not
101    /// present in this account key collection.
102    ///
103    /// Compilation will fail if any `instructions` use account keys which are located
104    /// at an index which cannot be cast to a `u8` without overflow.
105    pub fn try_compile_instructions(
106        &self,
107        instructions: &[Instruction],
108    ) -> Result<Vec<CompiledInstruction>, CompileError> {
109        let mut account_index_map = BTreeMap::<&Pubkey, u8>::new();
110        for (index, key) in self.iter().enumerate() {
111            let index = u8::try_from(index).map_err(|_| CompileError::AccountIndexOverflow)?;
112            account_index_map.insert(key, index);
113        }
114
115        let get_account_index = |key: &Pubkey| -> Result<u8, CompileError> {
116            account_index_map
117                .get(key)
118                .cloned()
119                .ok_or(CompileError::UnknownInstructionKey(*key))
120        };
121
122        instructions
123            .iter()
124            .map(|ix| {
125                let accounts: Vec<u8> = ix
126                    .accounts
127                    .iter()
128                    .map(|account_meta| get_account_index(&account_meta.pubkey))
129                    .collect::<Result<Vec<u8>, CompileError>>()?;
130
131                Ok(CompiledInstruction {
132                    program_id_index: get_account_index(&ix.program_id)?,
133                    data: ix.data.clone(),
134                    accounts,
135                })
136            })
137            .collect()
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use {super::*, crate::instruction::AccountMeta};
144
145    fn test_account_keys() -> [Pubkey; 6] {
146        let key0 = Pubkey::new_unique();
147        let key1 = Pubkey::new_unique();
148        let key2 = Pubkey::new_unique();
149        let key3 = Pubkey::new_unique();
150        let key4 = Pubkey::new_unique();
151        let key5 = Pubkey::new_unique();
152
153        [key0, key1, key2, key3, key4, key5]
154    }
155
156    #[test]
157    fn test_key_segment_iter() {
158        let keys = test_account_keys();
159
160        let static_keys = vec![keys[0], keys[1], keys[2]];
161        let dynamic_keys = LoadedAddresses {
162            writable: vec![keys[3], keys[4]],
163            readonly: vec![keys[5]],
164        };
165        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
166
167        let expected_segments = vec![
168            vec![keys[0], keys[1], keys[2]],
169            vec![keys[3], keys[4]],
170            vec![keys[5]],
171        ];
172
173        assert!(account_keys.key_segment_iter().eq(expected_segments.iter()));
174    }
175
176    #[test]
177    fn test_len() {
178        let keys = test_account_keys();
179
180        let static_keys = vec![keys[0], keys[1], keys[2], keys[3], keys[4], keys[5]];
181        let account_keys = AccountKeys::new(&static_keys, None);
182
183        assert_eq!(account_keys.len(), keys.len());
184    }
185
186    #[test]
187    fn test_len_with_dynamic_keys() {
188        let keys = test_account_keys();
189
190        let static_keys = vec![keys[0], keys[1], keys[2]];
191        let dynamic_keys = LoadedAddresses {
192            writable: vec![keys[3], keys[4]],
193            readonly: vec![keys[5]],
194        };
195        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
196
197        assert_eq!(account_keys.len(), keys.len());
198    }
199
200    #[test]
201    fn test_iter() {
202        let keys = test_account_keys();
203
204        let static_keys = vec![keys[0], keys[1], keys[2], keys[3], keys[4], keys[5]];
205        let account_keys = AccountKeys::new(&static_keys, None);
206
207        assert!(account_keys.iter().eq(keys.iter()));
208    }
209
210    #[test]
211    fn test_iter_with_dynamic_keys() {
212        let keys = test_account_keys();
213
214        let static_keys = vec![keys[0], keys[1], keys[2]];
215        let dynamic_keys = LoadedAddresses {
216            writable: vec![keys[3], keys[4]],
217            readonly: vec![keys[5]],
218        };
219        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
220
221        assert!(account_keys.iter().eq(keys.iter()));
222    }
223
224    #[test]
225    fn test_get() {
226        let keys = test_account_keys();
227
228        let static_keys = vec![keys[0], keys[1], keys[2], keys[3]];
229        let account_keys = AccountKeys::new(&static_keys, None);
230
231        assert_eq!(account_keys.get(0), Some(&keys[0]));
232        assert_eq!(account_keys.get(1), Some(&keys[1]));
233        assert_eq!(account_keys.get(2), Some(&keys[2]));
234        assert_eq!(account_keys.get(3), Some(&keys[3]));
235        assert_eq!(account_keys.get(4), None);
236        assert_eq!(account_keys.get(5), None);
237    }
238
239    #[test]
240    fn test_get_with_dynamic_keys() {
241        let keys = test_account_keys();
242
243        let static_keys = vec![keys[0], keys[1], keys[2]];
244        let dynamic_keys = LoadedAddresses {
245            writable: vec![keys[3], keys[4]],
246            readonly: vec![keys[5]],
247        };
248        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
249
250        assert_eq!(account_keys.get(0), Some(&keys[0]));
251        assert_eq!(account_keys.get(1), Some(&keys[1]));
252        assert_eq!(account_keys.get(2), Some(&keys[2]));
253        assert_eq!(account_keys.get(3), Some(&keys[3]));
254        assert_eq!(account_keys.get(4), Some(&keys[4]));
255        assert_eq!(account_keys.get(5), Some(&keys[5]));
256    }
257
258    #[test]
259    fn test_try_compile_instructions() {
260        let keys = test_account_keys();
261
262        let static_keys = vec![keys[0]];
263        let dynamic_keys = LoadedAddresses {
264            writable: vec![keys[1]],
265            readonly: vec![keys[2]],
266        };
267        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
268
269        let instruction = Instruction {
270            program_id: keys[0],
271            accounts: vec![
272                AccountMeta::new(keys[1], true),
273                AccountMeta::new(keys[2], true),
274            ],
275            data: vec![0],
276        };
277
278        assert_eq!(
279            account_keys.try_compile_instructions(&[instruction]),
280            Ok(vec![CompiledInstruction {
281                program_id_index: 0,
282                accounts: vec![1, 2],
283                data: vec![0],
284            }]),
285        );
286    }
287
288    #[test]
289    fn test_try_compile_instructions_with_unknown_key() {
290        let static_keys = test_account_keys();
291        let account_keys = AccountKeys::new(&static_keys, None);
292
293        let unknown_key = Pubkey::new_unique();
294        let test_instructions = [
295            Instruction {
296                program_id: unknown_key,
297                accounts: vec![],
298                data: vec![],
299            },
300            Instruction {
301                program_id: static_keys[0],
302                accounts: vec![
303                    AccountMeta::new(static_keys[1], false),
304                    AccountMeta::new(unknown_key, false),
305                ],
306                data: vec![],
307            },
308        ];
309
310        for ix in test_instructions {
311            assert_eq!(
312                account_keys.try_compile_instructions(&[ix]),
313                Err(CompileError::UnknownInstructionKey(unknown_key))
314            );
315        }
316    }
317
318    #[test]
319    fn test_try_compile_instructions_with_too_many_account_keys() {
320        const MAX_LENGTH_WITHOUT_OVERFLOW: usize = u8::MAX as usize + 1;
321        let static_keys = vec![Pubkey::default(); MAX_LENGTH_WITHOUT_OVERFLOW];
322        let dynamic_keys = LoadedAddresses {
323            writable: vec![Pubkey::default()],
324            readonly: vec![],
325        };
326        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
327        assert_eq!(
328            account_keys.try_compile_instructions(&[]),
329            Err(CompileError::AccountIndexOverflow)
330        );
331    }
332}
waffles_solana_program/message/account_keys.rs

waffles_solana_program/message/
account_keys.rs
