solana_message/
account_keys.rs

1use {
2    crate::{compiled_instruction::CompiledInstruction, v0::LoadedAddresses, CompileError},
3    alloc::{collections::BTreeMap, vec::Vec},
4    core::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        self.iter().eq(other.iter())
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_eq() {
237        let keys = test_account_keys();
238
239        let static_keys = vec![keys[0], keys[1], keys[2]];
240        let split_static_keys = vec![keys[0]];
241        let split_dynamic_keys = LoadedAddresses {
242            writable: vec![keys[1]],
243            readonly: vec![keys[2]],
244        };
245
246        assert_eq!(
247            AccountKeys::new(&static_keys, None),
248            AccountKeys::new(&split_static_keys, Some(&split_dynamic_keys)),
249        );
250    }
251
252    #[test]
253    fn test_ne_with_different_lengths() {
254        let keys = test_account_keys();
255
256        let short_static_keys = vec![keys[0]];
257        let long_static_keys = vec![keys[0], keys[1]];
258        let long_dynamic_keys = LoadedAddresses {
259            writable: vec![keys[1]],
260            readonly: vec![],
261        };
262
263        let short_account_keys = AccountKeys::new(&short_static_keys, None);
264        let long_static_account_keys = AccountKeys::new(&long_static_keys, None);
265        let long_dynamic_account_keys =
266            AccountKeys::new(&short_static_keys, Some(&long_dynamic_keys));
267
268        assert_ne!(short_account_keys, long_static_account_keys);
269        assert_ne!(long_static_account_keys, short_account_keys);
270        assert_ne!(short_account_keys, long_dynamic_account_keys);
271        assert_ne!(long_dynamic_account_keys, short_account_keys);
272    }
273
274    #[test]
275    fn test_get() {
276        let keys = test_account_keys();
277
278        let static_keys = vec![keys[0], keys[1], keys[2], keys[3]];
279        let account_keys = AccountKeys::new(&static_keys, None);
280
281        assert_eq!(account_keys.get(0), Some(&keys[0]));
282        assert_eq!(account_keys.get(1), Some(&keys[1]));
283        assert_eq!(account_keys.get(2), Some(&keys[2]));
284        assert_eq!(account_keys.get(3), Some(&keys[3]));
285        assert_eq!(account_keys.get(4), None);
286        assert_eq!(account_keys.get(5), None);
287    }
288
289    #[test]
290    fn test_get_with_dynamic_keys() {
291        let keys = test_account_keys();
292
293        let static_keys = vec![keys[0], keys[1], keys[2]];
294        let dynamic_keys = LoadedAddresses {
295            writable: vec![keys[3], keys[4]],
296            readonly: vec![keys[5]],
297        };
298        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
299
300        assert_eq!(account_keys.get(0), Some(&keys[0]));
301        assert_eq!(account_keys.get(1), Some(&keys[1]));
302        assert_eq!(account_keys.get(2), Some(&keys[2]));
303        assert_eq!(account_keys.get(3), Some(&keys[3]));
304        assert_eq!(account_keys.get(4), Some(&keys[4]));
305        assert_eq!(account_keys.get(5), Some(&keys[5]));
306    }
307
308    #[test]
309    fn test_try_compile_instructions() {
310        let keys = test_account_keys();
311
312        let static_keys = vec![keys[0]];
313        let dynamic_keys = LoadedAddresses {
314            writable: vec![keys[1]],
315            readonly: vec![keys[2]],
316        };
317        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
318
319        let instruction = Instruction {
320            program_id: keys[0],
321            accounts: vec![
322                AccountMeta::new(keys[1], true),
323                AccountMeta::new(keys[2], true),
324            ],
325            data: vec![0],
326        };
327
328        assert_eq!(
329            account_keys.try_compile_instructions(&[instruction]),
330            Ok(vec![CompiledInstruction {
331                program_id_index: 0,
332                accounts: vec![1, 2],
333                data: vec![0],
334            }]),
335        );
336    }
337
338    #[test]
339    fn test_try_compile_instructions_with_unknown_key() {
340        let static_keys = test_account_keys();
341        let account_keys = AccountKeys::new(&static_keys, None);
342
343        let unknown_key = Address::new_unique();
344        let test_instructions = [
345            Instruction {
346                program_id: unknown_key,
347                accounts: vec![],
348                data: vec![],
349            },
350            Instruction {
351                program_id: static_keys[0],
352                accounts: vec![
353                    AccountMeta::new(static_keys[1], false),
354                    AccountMeta::new(unknown_key, false),
355                ],
356                data: vec![],
357            },
358        ];
359
360        for ix in test_instructions {
361            assert_eq!(
362                account_keys.try_compile_instructions(&[ix]),
363                Err(CompileError::UnknownInstructionKey(unknown_key))
364            );
365        }
366    }
367
368    #[test]
369    fn test_try_compile_instructions_with_too_many_account_keys() {
370        const MAX_LENGTH_WITHOUT_OVERFLOW: usize = u8::MAX as usize + 1;
371        let static_keys = vec![Address::default(); MAX_LENGTH_WITHOUT_OVERFLOW];
372        let dynamic_keys = LoadedAddresses {
373            writable: vec![Address::default()],
374            readonly: vec![],
375        };
376        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
377        assert_eq!(
378            account_keys.try_compile_instructions(&[]),
379            Err(CompileError::AccountIndexOverflow)
380        );
381    }
382}
solana_message/account_keys.rs

solana_message/
account_keys.rs
