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

solana_program/message/
account_keys.rs
