solomka_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
174            .key_segment_iter()
175            .into_iter()
176            .eq(expected_segments.iter()));
177    }
178
179    #[test]
180    fn test_len() {
181        let keys = test_account_keys();
182
183        let static_keys = vec![keys[0], keys[1], keys[2], keys[3], keys[4], keys[5]];
184        let account_keys = AccountKeys::new(&static_keys, None);
185
186        assert_eq!(account_keys.len(), keys.len());
187    }
188
189    #[test]
190    fn test_len_with_dynamic_keys() {
191        let keys = test_account_keys();
192
193        let static_keys = vec![keys[0], keys[1], keys[2]];
194        let dynamic_keys = LoadedAddresses {
195            writable: vec![keys[3], keys[4]],
196            readonly: vec![keys[5]],
197        };
198        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
199
200        assert_eq!(account_keys.len(), keys.len());
201    }
202
203    #[test]
204    fn test_iter() {
205        let keys = test_account_keys();
206
207        let static_keys = vec![keys[0], keys[1], keys[2], keys[3], keys[4], keys[5]];
208        let account_keys = AccountKeys::new(&static_keys, None);
209
210        assert!(account_keys.iter().eq(keys.iter()));
211    }
212
213    #[test]
214    fn test_iter_with_dynamic_keys() {
215        let keys = test_account_keys();
216
217        let static_keys = vec![keys[0], keys[1], keys[2]];
218        let dynamic_keys = LoadedAddresses {
219            writable: vec![keys[3], keys[4]],
220            readonly: vec![keys[5]],
221        };
222        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
223
224        assert!(account_keys.iter().eq(keys.iter()));
225    }
226
227    #[test]
228    fn test_get() {
229        let keys = test_account_keys();
230
231        let static_keys = vec![keys[0], keys[1], keys[2], keys[3]];
232        let account_keys = AccountKeys::new(&static_keys, None);
233
234        assert_eq!(account_keys.get(0), Some(&keys[0]));
235        assert_eq!(account_keys.get(1), Some(&keys[1]));
236        assert_eq!(account_keys.get(2), Some(&keys[2]));
237        assert_eq!(account_keys.get(3), Some(&keys[3]));
238        assert_eq!(account_keys.get(4), None);
239        assert_eq!(account_keys.get(5), None);
240    }
241
242    #[test]
243    fn test_get_with_dynamic_keys() {
244        let keys = test_account_keys();
245
246        let static_keys = vec![keys[0], keys[1], keys[2]];
247        let dynamic_keys = LoadedAddresses {
248            writable: vec![keys[3], keys[4]],
249            readonly: vec![keys[5]],
250        };
251        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
252
253        assert_eq!(account_keys.get(0), Some(&keys[0]));
254        assert_eq!(account_keys.get(1), Some(&keys[1]));
255        assert_eq!(account_keys.get(2), Some(&keys[2]));
256        assert_eq!(account_keys.get(3), Some(&keys[3]));
257        assert_eq!(account_keys.get(4), Some(&keys[4]));
258        assert_eq!(account_keys.get(5), Some(&keys[5]));
259    }
260
261    #[test]
262    fn test_try_compile_instructions() {
263        let keys = test_account_keys();
264
265        let static_keys = vec![keys[0]];
266        let dynamic_keys = LoadedAddresses {
267            writable: vec![keys[1]],
268            readonly: vec![keys[2]],
269        };
270        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
271
272        let instruction = Instruction {
273            program_id: keys[0],
274            accounts: vec![
275                AccountMeta::new(keys[1], true),
276                AccountMeta::new(keys[2], true),
277            ],
278            data: vec![0],
279        };
280
281        assert_eq!(
282            account_keys.try_compile_instructions(&[instruction]),
283            Ok(vec![CompiledInstruction {
284                program_id_index: 0,
285                accounts: vec![1, 2],
286                data: vec![0],
287            }]),
288        );
289    }
290
291    #[test]
292    fn test_try_compile_instructions_with_unknown_key() {
293        let static_keys = test_account_keys();
294        let account_keys = AccountKeys::new(&static_keys, None);
295
296        let unknown_key = Pubkey::new_unique();
297        let test_instructions = [
298            Instruction {
299                program_id: unknown_key,
300                accounts: vec![],
301                data: vec![],
302            },
303            Instruction {
304                program_id: static_keys[0],
305                accounts: vec![
306                    AccountMeta::new(static_keys[1], false),
307                    AccountMeta::new(unknown_key, false),
308                ],
309                data: vec![],
310            },
311        ];
312
313        for ix in test_instructions {
314            assert_eq!(
315                account_keys.try_compile_instructions(&[ix]),
316                Err(CompileError::UnknownInstructionKey(unknown_key))
317            );
318        }
319    }
320
321    #[test]
322    fn test_try_compile_instructions_with_too_many_account_keys() {
323        const MAX_LENGTH_WITHOUT_OVERFLOW: usize = u8::MAX as usize + 1;
324        let static_keys = vec![Pubkey::default(); MAX_LENGTH_WITHOUT_OVERFLOW];
325        let dynamic_keys = LoadedAddresses {
326            writable: vec![Pubkey::default()],
327            readonly: vec![],
328        };
329        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
330        assert_eq!(
331            account_keys.try_compile_instructions(&[]),
332            Err(CompileError::AccountIndexOverflow)
333        );
334    }
335}
solomka_program/message/account_keys.rs

solomka_program/message/
account_keys.rs
