solana_message/
account_keys.rs

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

solana_message/
account_keys.rs
