use {
crate::{
instruction::{CompiledInstruction, Instruction},
message::v0::LoadedAddresses,
pubkey::Pubkey,
},
std::{collections::BTreeMap, ops::Index},
};
pub struct AccountKeys<'a> {
static_keys: &'a [Pubkey],
dynamic_keys: Option<&'a LoadedAddresses>,
}
impl Index<usize> for AccountKeys<'_> {
type Output = Pubkey;
fn index(&self, index: usize) -> &Self::Output {
self.get(index).expect("index is invalid")
}
}
impl<'a> AccountKeys<'a> {
pub fn new(static_keys: &'a [Pubkey], dynamic_keys: Option<&'a LoadedAddresses>) -> Self {
Self {
static_keys,
dynamic_keys,
}
}
fn key_segment_iter(&self) -> impl Iterator<Item = &'a [Pubkey]> {
if let Some(dynamic_keys) = self.dynamic_keys {
[
self.static_keys,
&dynamic_keys.writable,
&dynamic_keys.readonly,
]
.into_iter()
} else {
[self.static_keys, &[], &[]].into_iter()
}
}
pub fn get(&self, mut index: usize) -> Option<&'a Pubkey> {
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
}
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
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn iter(&self) -> impl Iterator<Item = &'a Pubkey> {
self.key_segment_iter().flatten()
}
pub fn compile_instructions(&self, instructions: &[Instruction]) -> Vec<CompiledInstruction> {
let account_index_map: BTreeMap<&Pubkey, u8> = BTreeMap::from_iter(
self.iter()
.enumerate()
.map(|(index, key)| (key, index as u8)),
);
instructions
.iter()
.map(|ix| {
let accounts: Vec<u8> = ix
.accounts
.iter()
.map(|account_meta| *account_index_map.get(&account_meta.pubkey).unwrap())
.collect();
CompiledInstruction {
program_id_index: *account_index_map.get(&ix.program_id).unwrap(),
data: ix.data.clone(),
accounts,
}
})
.collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn test_account_keys() -> [Pubkey; 6] {
let key0 = Pubkey::new_unique();
let key1 = Pubkey::new_unique();
let key2 = Pubkey::new_unique();
let key3 = Pubkey::new_unique();
let key4 = Pubkey::new_unique();
let key5 = Pubkey::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![
vec![keys[0], keys[1], keys[2]],
vec![keys[3], keys[4]],
vec![keys[5]],
];
assert!(account_keys
.key_segment_iter()
.into_iter()
.eq(expected_segments.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_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]));
}
}