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

miraland_program/message/
account_keys.rs
