solana_message/

account_keys.rs

use {
    crate::{
        compiled_instruction::CompiledInstruction,
        v0::{LoadedAddresses, LoadedAddressesView},
        CompileError,
    },
    alloc::{collections::BTreeMap, vec::Vec},
    core::ops::Index,
    solana_address::Address,
    solana_instruction::Instruction,
};

/// Collection of static and dynamically loaded keys used to load accounts
/// during transaction processing.
#[derive(Clone, Default, Debug, Eq)]
pub struct AccountKeys<'a> {
    static_keys: &'a [Address],
    dynamic_keys: Option<LoadedAddressesView<'a>>,
}

impl Index<usize> for AccountKeys<'_> {
    type Output = Address;
    #[inline]
    fn index(&self, index: usize) -> &Self::Output {
        self.get(index).expect("index is invalid")
    }
}

impl<'a> AccountKeys<'a> {
    pub fn new(static_keys: &'a [Address], dynamic_keys: Option<&'a LoadedAddresses>) -> Self {
        Self::new_with_loaded_addresses_view(
            static_keys,
            dynamic_keys.map(LoadedAddressesView::from),
        )
    }

    /// Creates account keys from static keys and an optional borrowed view of
    /// dynamically loaded addresses.
    pub fn new_with_loaded_addresses_view(
        static_keys: &'a [Address],
        dynamic_keys: Option<LoadedAddressesView<'a>>,
    ) -> Self {
        Self {
            static_keys,
            dynamic_keys,
        }
    }

    /// Returns an iterator of account key segments. The ordering of segments
    /// affects how account indexes from compiled instructions are resolved and
    /// so should not be changed.
    #[inline]
    fn key_segment_iter(&self) -> impl Iterator<Item = &'a [Address]> + Clone {
        if let Some(dynamic_keys) = self.dynamic_keys {
            [
                self.static_keys,
                dynamic_keys.writable,
                dynamic_keys.readonly,
            ]
            .into_iter()
        } else {
            // empty segments added for branch type compatibility
            [self.static_keys, &[], &[]].into_iter()
        }
    }

    /// Returns the address of the account at the specified index of the list of
    /// message account keys constructed from static keys, followed by dynamically
    /// loaded writable addresses, and lastly the list of dynamically loaded
    /// readonly addresses.
    #[inline]
    pub fn get(&self, mut index: usize) -> Option<&'a Address> {
        for key_segment in self.key_segment_iter() {
            if index < key_segment.len() {
                return Some(&key_segment[index]);
            }
            index = index.saturating_sub(key_segment.len());
        }

        None
    }

    /// Returns the total length of loaded accounts for a message
    #[inline]
    pub fn len(&self) -> usize {
        let mut len = 0usize;
        for key_segment in self.key_segment_iter() {
            len = len.saturating_add(key_segment.len());
        }
        len
    }

    /// Returns true if this collection of account keys is empty
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Iterator for the addresses of the loaded accounts for a message
    #[inline]
    pub fn iter(&self) -> impl Iterator<Item = &'a Address> + Clone {
        self.key_segment_iter().flatten()
    }

    /// Compile instructions using the order of account keys to determine
    /// compiled instruction account indexes.
    ///
    /// # Panics
    ///
    /// Panics when compiling fails. See [`AccountKeys::try_compile_instructions`]
    /// for a full description of failure scenarios.
    pub fn compile_instructions(&self, instructions: &[Instruction]) -> Vec<CompiledInstruction> {
        self.try_compile_instructions(instructions)
            .expect("compilation failure")
    }

    /// Compile instructions using the order of account keys to determine
    /// compiled instruction account indexes.
    ///
    /// # Errors
    ///
    /// Compilation will fail if any `instructions` use account keys which are not
    /// present in this account key collection.
    ///
    /// Compilation will fail if any `instructions` use account keys which are located
    /// at an index which cannot be cast to a `u8` without overflow.
    pub fn try_compile_instructions(
        &self,
        instructions: &[Instruction],
    ) -> Result<Vec<CompiledInstruction>, CompileError> {
        let mut account_index_map = BTreeMap::<&Address, u8>::new();
        for (index, key) in self.iter().enumerate() {
            let index = u8::try_from(index).map_err(|_| CompileError::AccountIndexOverflow)?;
            account_index_map.insert(key, index);
        }

        let get_account_index = |key: &Address| -> Result<u8, CompileError> {
            account_index_map
                .get(key)
                .cloned()
                .ok_or(CompileError::UnknownInstructionKey(*key))
        };

        instructions
            .iter()
            .map(|ix| {
                let accounts: Vec<u8> = ix
                    .accounts
                    .iter()
                    .map(|account_meta| get_account_index(&account_meta.pubkey))
                    .collect::<Result<Vec<u8>, CompileError>>()?;

                Ok(CompiledInstruction {
                    program_id_index: get_account_index(&ix.program_id)?,
                    data: ix.data.clone(),
                    accounts,
                })
            })
            .collect()
    }
}

impl PartialEq for AccountKeys<'_> {
    fn eq(&self, other: &Self) -> bool {
        self.iter().eq(other.iter())
    }
}

#[cfg(test)]
mod tests {
    use {super::*, alloc::vec, solana_instruction::AccountMeta};

    fn test_account_keys() -> [Address; 6] {
        let key0 = Address::new_unique();
        let key1 = Address::new_unique();
        let key2 = Address::new_unique();
        let key3 = Address::new_unique();
        let key4 = Address::new_unique();
        let key5 = Address::new_unique();

        [key0, key1, key2, key3, key4, key5]
    }

    #[test]
    fn test_key_segment_iter() {
        let keys = test_account_keys();

        let static_keys = vec![keys[0], keys[1], keys[2]];
        let dynamic_keys = LoadedAddresses {
            writable: vec![keys[3], keys[4]],
            readonly: vec![keys[5]],
        };
        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));

        let expected_segments = [
            vec![keys[0], keys[1], keys[2]],
            vec![keys[3], keys[4]],
            vec![keys[5]],
        ];

        assert!(account_keys.key_segment_iter().eq(expected_segments.iter()));
    }

    #[test]
    fn test_key_segment_iter_with_loaded_addresses_view() {
        let keys = test_account_keys();

        let static_keys = [keys[0], keys[1], keys[2]];
        let writable_dynamic_keys = [keys[3], keys[4]];
        let readonly_dynamic_keys = [keys[5]];
        let dynamic_keys = LoadedAddressesView {
            writable: &writable_dynamic_keys,
            readonly: &readonly_dynamic_keys,
        };
        let account_keys =
            AccountKeys::new_with_loaded_addresses_view(&static_keys, Some(dynamic_keys));

        let expected_segments = [
            &static_keys[..],
            &writable_dynamic_keys[..],
            &readonly_dynamic_keys[..],
        ];

        assert!(account_keys
            .key_segment_iter()
            .eq(expected_segments.into_iter()));
    }

    #[test]
    fn test_len() {
        let keys = test_account_keys();

        let static_keys = vec![keys[0], keys[1], keys[2], keys[3], keys[4], keys[5]];
        let account_keys = AccountKeys::new(&static_keys, None);

        assert_eq!(account_keys.len(), keys.len());
    }

    #[test]
    fn test_len_with_dynamic_keys() {
        let keys = test_account_keys();

        let static_keys = vec![keys[0], keys[1], keys[2]];
        let dynamic_keys = LoadedAddresses {
            writable: vec![keys[3], keys[4]],
            readonly: vec![keys[5]],
        };
        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));

        assert_eq!(account_keys.len(), keys.len());
    }

    #[test]
    fn test_iter() {
        let keys = test_account_keys();

        let static_keys = vec![keys[0], keys[1], keys[2], keys[3], keys[4], keys[5]];
        let account_keys = AccountKeys::new(&static_keys, None);

        assert!(account_keys.iter().eq(keys.iter()));
    }

    #[test]
    fn test_iter_with_dynamic_keys() {
        let keys = test_account_keys();

        let static_keys = vec![keys[0], keys[1], keys[2]];
        let dynamic_keys = LoadedAddresses {
            writable: vec![keys[3], keys[4]],
            readonly: vec![keys[5]],
        };
        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));

        assert!(account_keys.iter().eq(keys.iter()));
    }

    #[test]
    fn test_eq() {
        let keys = test_account_keys();

        let static_keys = vec![keys[0], keys[1], keys[2]];
        let split_static_keys = vec![keys[0]];
        let split_dynamic_keys = LoadedAddresses {
            writable: vec![keys[1]],
            readonly: vec![keys[2]],
        };

        assert_eq!(
            AccountKeys::new(&static_keys, None),
            AccountKeys::new(&split_static_keys, Some(&split_dynamic_keys)),
        );
    }

    #[test]
    fn test_ne_with_different_lengths() {
        let keys = test_account_keys();

        let short_static_keys = vec![keys[0]];
        let long_static_keys = vec![keys[0], keys[1]];
        let long_dynamic_keys = LoadedAddresses {
            writable: vec![keys[1]],
            readonly: vec![],
        };

        let short_account_keys = AccountKeys::new(&short_static_keys, None);
        let long_static_account_keys = AccountKeys::new(&long_static_keys, None);
        let long_dynamic_account_keys =
            AccountKeys::new(&short_static_keys, Some(&long_dynamic_keys));

        assert_ne!(short_account_keys, long_static_account_keys);
        assert_ne!(long_static_account_keys, short_account_keys);
        assert_ne!(short_account_keys, long_dynamic_account_keys);
        assert_ne!(long_dynamic_account_keys, short_account_keys);
    }

    #[test]
    fn test_get() {
        let keys = test_account_keys();

        let static_keys = vec![keys[0], keys[1], keys[2], keys[3]];
        let account_keys = AccountKeys::new(&static_keys, None);

        assert_eq!(account_keys.get(0), Some(&keys[0]));
        assert_eq!(account_keys.get(1), Some(&keys[1]));
        assert_eq!(account_keys.get(2), Some(&keys[2]));
        assert_eq!(account_keys.get(3), Some(&keys[3]));
        assert_eq!(account_keys.get(4), None);
        assert_eq!(account_keys.get(5), None);
    }

    #[test]
    fn test_get_with_dynamic_keys() {
        let keys = test_account_keys();

        let static_keys = vec![keys[0], keys[1], keys[2]];
        let dynamic_keys = LoadedAddresses {
            writable: vec![keys[3], keys[4]],
            readonly: vec![keys[5]],
        };
        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));

        assert_eq!(account_keys.get(0), Some(&keys[0]));
        assert_eq!(account_keys.get(1), Some(&keys[1]));
        assert_eq!(account_keys.get(2), Some(&keys[2]));
        assert_eq!(account_keys.get(3), Some(&keys[3]));
        assert_eq!(account_keys.get(4), Some(&keys[4]));
        assert_eq!(account_keys.get(5), Some(&keys[5]));
    }

    #[test]
    fn test_owned_and_view_dynamic_keys_behave_the_same() {
        let keys = test_account_keys();

        let static_keys = [keys[0], keys[1]];
        let dynamic_keys = LoadedAddresses {
            writable: vec![keys[2], keys[3]],
            readonly: vec![keys[4], keys[5]],
        };
        let dynamic_keys_view = LoadedAddressesView::from(&dynamic_keys);
        let owned_backed_account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
        let view_backed_account_keys =
            AccountKeys::new_with_loaded_addresses_view(&static_keys, Some(dynamic_keys_view));

        assert_eq!(
            dynamic_keys_view,
            LoadedAddressesView {
                writable: &dynamic_keys.writable,
                readonly: &dynamic_keys.readonly,
            },
        );
        assert_eq!(
            owned_backed_account_keys.len(),
            view_backed_account_keys.len()
        );
        for index in 0..=keys.len() {
            assert_eq!(
                owned_backed_account_keys.get(index),
                view_backed_account_keys.get(index),
            );
        }
        assert!(owned_backed_account_keys
            .iter()
            .eq(view_backed_account_keys.iter()));
        assert_eq!(owned_backed_account_keys, view_backed_account_keys);
    }

    #[test]
    fn test_loaded_addresses_view_from_borrowed_slices() {
        let keys = test_account_keys();

        let static_keys = [keys[0], keys[1]];
        let writable_dynamic_keys = [keys[2], keys[3]];
        let readonly_dynamic_keys = [keys[4], keys[5]];
        let account_keys = AccountKeys::new_with_loaded_addresses_view(
            &static_keys,
            Some(LoadedAddressesView {
                writable: &writable_dynamic_keys,
                readonly: &readonly_dynamic_keys,
            }),
        );

        assert_eq!(account_keys.len(), keys.len());
        for (index, expected_key) in keys.iter().enumerate() {
            assert_eq!(account_keys.get(index), Some(expected_key));
        }
        assert_eq!(account_keys.get(keys.len()), None);
        assert!(account_keys.iter().eq(keys.iter()));
    }

    #[test]
    fn test_try_compile_instructions() {
        let keys = test_account_keys();

        let static_keys = vec![keys[0]];
        let dynamic_keys = LoadedAddresses {
            writable: vec![keys[1]],
            readonly: vec![keys[2]],
        };
        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));

        let instruction = Instruction {
            program_id: keys[0],
            accounts: vec![
                AccountMeta::new(keys[1], true),
                AccountMeta::new(keys[2], true),
            ],
            data: vec![0],
        };

        assert_eq!(
            account_keys.try_compile_instructions(&[instruction]),
            Ok(vec![CompiledInstruction {
                program_id_index: 0,
                accounts: vec![1, 2],
                data: vec![0],
            }]),
        );
    }

    #[test]
    fn test_try_compile_instructions_with_unknown_key() {
        let static_keys = test_account_keys();
        let account_keys = AccountKeys::new(&static_keys, None);

        let unknown_key = Address::new_unique();
        let test_instructions = [
            Instruction {
                program_id: unknown_key,
                accounts: vec![],
                data: vec![],
            },
            Instruction {
                program_id: static_keys[0],
                accounts: vec![
                    AccountMeta::new(static_keys[1], false),
                    AccountMeta::new(unknown_key, false),
                ],
                data: vec![],
            },
        ];

        for ix in test_instructions {
            assert_eq!(
                account_keys.try_compile_instructions(&[ix]),
                Err(CompileError::UnknownInstructionKey(unknown_key))
            );
        }
    }

    #[test]
    fn test_try_compile_instructions_with_too_many_account_keys() {
        const MAX_LENGTH_WITHOUT_OVERFLOW: usize = u8::MAX as usize + 1;
        let static_keys = vec![Address::default(); MAX_LENGTH_WITHOUT_OVERFLOW];
        let dynamic_keys = LoadedAddresses {
            writable: vec![Address::default()],
            readonly: vec![],
        };
        let account_keys = AccountKeys::new(&static_keys, Some(&dynamic_keys));
        assert_eq!(
            account_keys.try_compile_instructions(&[]),
            Err(CompileError::AccountIndexOverflow)
        );
    }
}