Refactor: Add AccountKeys struct for static and dynamic message keys (#22960)

sdk/program/src/message/account_keys.rs

@@ -0,0 +1,199 @@
+use {
+    crate::{message::v0::LoadedAddresses, pubkey::Pubkey},
+    std::ops::Index,
+};
+
+/// Collection of static and dynamically loaded keys used to load accounts
+/// during transaction processing.
+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,
+        }
+    }
+
+    /// 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.
+    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 {
+            // 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.
+    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
+    }
+
+    /// Returns the total length of loaded accounts for a message
+    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
+    pub fn iter(&self) -> impl Iterator<Item = &'a Pubkey> {
+        self.key_segment_iter().flatten()
+    }
+}
+
+#[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]));
+    }
+}

sdk/program/src/message/mod.rs

@@ -5,10 +5,11 @@ pub mod legacy;
 #[cfg(not(target_arch = "bpf"))]
 #[path = ""]
 mod non_bpf_modules {
+    mod account_keys;
     mod sanitized;
     mod versions;
 
-    pub use {sanitized::*, versions::*};
+    pub use {account_keys::*, sanitized::*, versions::*};
 }
 
 pub use legacy::Message;

sdk/program/src/message/sanitized.rs

@@ -5,7 +5,7 @@ use {
         message::{
             legacy::Message as LegacyMessage,
             v0::{self, LoadedAddresses},
-            MessageHeader,
+            AccountKeys, MessageHeader,
         },
         nonce::NONCED_TX_MARKER_IX_INDEX,
         program_utils::limited_deserialize,
@@ -85,7 +85,8 @@ impl SanitizedMessage {
 
     /// Returns the fee payer for the transaction
     pub fn fee_payer(&self) -> &Pubkey {
-        self.get_account_key(0)
+        self.account_keys()
+            .get(0)
             .expect("sanitized message always has non-program fee payer at index 0")
     }
 
@@ -117,34 +118,19 @@ impl SanitizedMessage {
         }
         .map(move |ix| {
             (
-                self.get_account_key(usize::from(ix.program_id_index))
+                self.account_keys()
+                    .get(usize::from(ix.program_id_index))
                     .expect("program id index is sanitized"),
                 ix,
             )
         })
     }
 
-    /// Iterator of all account keys referenced in this message, including dynamically loaded keys.
-    pub fn account_keys_iter(&self) -> Box<dyn Iterator<Item = &Pubkey> + '_> {
+    /// Returns the list of account keys that are loaded for this message.
+    pub fn account_keys(&self) -> AccountKeys {
         match self {
-            Self::Legacy(message) => Box::new(message.account_keys.iter()),
-            Self::V0(message) => Box::new(message.account_keys_iter()),
-        }
-    }
-
-    /// Length of all account keys referenced in this message, including dynamically loaded keys.
-    pub fn account_keys_len(&self) -> usize {
-        match self {
-            Self::Legacy(message) => message.account_keys.len(),
-            Self::V0(message) => message.account_keys_len(),
-        }
-    }
-
-    /// Returns the address of the account at the specified index.
-    pub fn get_account_key(&self, index: usize) -> Option<&Pubkey> {
-        match self {
-            Self::Legacy(message) => message.account_keys.get(index),
-            Self::V0(message) => message.get_account_key(index),
+            Self::Legacy(message) => AccountKeys::new(&message.account_keys, None),
+            Self::V0(message) => message.account_keys(),
         }
     }
 
@@ -210,7 +196,7 @@ impl SanitizedMessage {
     }
 
     fn try_position(&self, key: &Pubkey) -> Option<u8> {
-        u8::try_from(self.account_keys_iter().position(|k| k == key)?).ok()
+        u8::try_from(self.account_keys().iter().position(|k| k == key)?).ok()
     }
 
     /// Try to compile an instruction using the account keys in this message.
@@ -230,6 +216,7 @@ impl SanitizedMessage {
 
     /// Decompile message instructions without cloning account keys
     pub fn decompile_instructions(&self) -> Vec<BorrowedInstruction> {
+        let account_keys = self.account_keys();
         self.program_instructions_iter()
             .map(|(program_id, instruction)| {
                 let accounts = instruction
@@ -240,7 +227,7 @@ impl SanitizedMessage {
                         BorrowedAccountMeta {
                             is_signer: self.is_signer(account_index),
                             is_writable: self.is_writable(account_index),
-                            pubkey: self.get_account_key(account_index).unwrap(),
+                            pubkey: account_keys.get(account_index).unwrap(),
                         }
                     })
                     .collect();
@@ -267,7 +254,7 @@ impl SanitizedMessage {
         self.instructions()
             .get(NONCED_TX_MARKER_IX_INDEX as usize)
             .filter(
-                |ix| match self.get_account_key(ix.program_id_index as usize) {
+                |ix| match self.account_keys().get(ix.program_id_index as usize) {
                     Some(program_id) => system_program::check_id(program_id),
                     _ => false,
                 },
@@ -284,7 +271,7 @@ impl SanitizedMessage {
                     if nonce_must_be_writable && !self.is_writable(idx) {
                         None
                     } else {
-                        self.get_account_key(idx)
+                        self.account_keys().get(idx)
                     }
                 })
             })

sdk/program/src/message/versions/mod.rs

@@ -50,45 +50,6 @@ impl VersionedMessage {
         }
     }
 
-    pub fn into_static_account_keys(self) -> Vec<Pubkey> {
-        match self {
-            Self::Legacy(message) => message.account_keys,
-            Self::V0(message) => message.account_keys,
-        }
-    }
-
-    pub fn static_account_keys_iter(&self) -> impl Iterator<Item = &Pubkey> {
-        match self {
-            Self::Legacy(message) => message.account_keys.iter(),
-            Self::V0(message) => message.account_keys.iter(),
-        }
-    }
-
-    pub fn static_account_keys_len(&self) -> usize {
-        match self {
-            Self::Legacy(message) => message.account_keys.len(),
-            Self::V0(message) => message.account_keys.len(),
-        }
-    }
-
-    pub fn total_account_keys_len(&self) -> usize {
-        match self {
-            Self::Legacy(message) => message.account_keys.len(),
-            Self::V0(message) => message.account_keys.len().saturating_add(
-                message
-                    .address_table_lookups
-                    .iter()
-                    .map(|lookup| {
-                        lookup
-                            .writable_indexes
-                            .len()
-                            .saturating_add(lookup.readonly_indexes.len())
-                    })
-                    .sum(),
-            ),
-        }
-    }
-
     pub fn recent_blockhash(&self) -> &Hash {
         match self {
             Self::Legacy(message) => &message.recent_blockhash,

sdk/program/src/message/versions/v0/loaded.rs

@@ -1,7 +1,7 @@
 use {
     crate::{
         bpf_loader_upgradeable,
-        message::{legacy::BUILTIN_PROGRAMS_KEYS, v0},
+        message::{legacy::BUILTIN_PROGRAMS_KEYS, v0, AccountKeys},
         pubkey::Pubkey,
         sysvar,
     },
@@ -57,66 +57,18 @@ impl LoadedAddresses {
     pub fn len(&self) -> usize {
         self.writable.len().saturating_add(self.readonly.len())
     }
-
-    /// Iterate over loaded addresses in the order that they are used
-    /// as account indexes
-    pub fn ordered_iter(&self) -> impl Iterator<Item = &Pubkey> {
-        self.writable.iter().chain(self.readonly.iter())
-    }
-
-    /// Iterate over loaded addresses in the order that they are used
-    /// as account indexes
-    pub fn into_ordered_iter(self) -> impl Iterator<Item = Pubkey> {
-        self.writable.into_iter().chain(self.readonly.into_iter())
-    }
 }
 
 impl LoadedMessage {
-    /// 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.
-    fn account_keys_segment_iter(&self) -> impl Iterator<Item = &Vec<Pubkey>> {
-        vec![
-            &self.message.account_keys,
-            &self.loaded_addresses.writable,
-            &self.loaded_addresses.readonly,
-        ]
-        .into_iter()
-    }
-
-    /// Returns the total length of loaded accounts for this message
-    pub fn account_keys_len(&self) -> usize {
-        let mut len = 0usize;
-        for key_segment in self.account_keys_segment_iter() {
-            len = len.saturating_add(key_segment.len());
-        }
-        len
-    }
-
-    /// Iterator for the addresses of the loaded accounts for this message
-    pub fn account_keys_iter(&self) -> impl Iterator<Item = &Pubkey> {
-        self.account_keys_segment_iter().flatten()
+    /// Returns the list of account keys that are loaded for this message.
+    pub fn account_keys(&self) -> AccountKeys {
+        AccountKeys::new(&self.account_keys, Some(&self.loaded_addresses))
     }
 
     /// Returns true if any account keys are duplicates
     pub fn has_duplicates(&self) -> bool {
         let mut uniq = HashSet::new();
-        self.account_keys_iter().any(|x| !uniq.insert(x))
-    }
-
-    /// 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.
-    pub fn get_account_key(&self, mut index: usize) -> Option<&Pubkey> {
-        for key_segment in self.account_keys_segment_iter() {
-            if index < key_segment.len() {
-                return Some(&key_segment[index]);
-            }
-            index = index.saturating_sub(key_segment.len());
-        }
-
-        None
+        self.account_keys().iter().any(|x| !uniq.insert(x))
     }
 
     /// Returns true if the account at the specified index was requested to be
@@ -144,7 +96,7 @@ impl LoadedMessage {
     /// Returns true if the account at the specified index was loaded as writable
     pub fn is_writable(&self, key_index: usize) -> bool {
         if self.is_writable_index(key_index) {
-            if let Some(key) = self.get_account_key(key_index) {
+            if let Some(key) = self.account_keys().get(key_index) {
                 let demote_program_id = self.is_key_called_as_program(key_index)
                     && !self.is_upgradeable_loader_present();
                 return !(sysvar::is_sysvar_id(key)
@@ -169,7 +121,8 @@ impl LoadedMessage {
 
     /// Returns true if any account is the bpf upgradeable loader
     pub fn is_upgradeable_loader_present(&self) -> bool {
-        self.account_keys_iter()
+        self.account_keys()
+            .iter()
             .any(|&key| key == bpf_loader_upgradeable::id())
     }
 }
@@ -209,39 +162,6 @@ mod tests {
         (message, [key0, key1, key2, key3, key4, key5])
     }
 
-    #[test]
-    fn test_account_keys_segment_iter() {
-        let (message, keys) = check_test_loaded_message();
-
-        let expected_segments = vec![
-            vec![keys[0], keys[1], keys[2], keys[3]],
-            vec![keys[4]],
-            vec![keys[5]],
-        ];
-
-        let mut iter = message.account_keys_segment_iter();
-        for expected_segment in expected_segments {
-            assert_eq!(iter.next(), Some(&expected_segment));
-        }
-    }
-
-    #[test]
-    fn test_account_keys_len() {
-        let (message, keys) = check_test_loaded_message();
-
-        assert_eq!(message.account_keys_len(), keys.len());
-    }
-
-    #[test]
-    fn test_account_keys_iter() {
-        let (message, keys) = check_test_loaded_message();
-
-        let mut iter = message.account_keys_iter();
-        for expected_key in keys {
-            assert_eq!(iter.next(), Some(&expected_key));
-        }
-    }
-
     #[test]
     fn test_has_duplicates() {
         let message = check_test_loaded_message().0;
@@ -276,18 +196,6 @@ mod tests {
         }
     }
 
-    #[test]
-    fn test_get_account_key() {
-        let (message, keys) = check_test_loaded_message();
-
-        assert_eq!(message.get_account_key(0), Some(&keys[0]));
-        assert_eq!(message.get_account_key(1), Some(&keys[1]));
-        assert_eq!(message.get_account_key(2), Some(&keys[2]));
-        assert_eq!(message.get_account_key(3), Some(&keys[3]));
-        assert_eq!(message.get_account_key(4), Some(&keys[4]));
-        assert_eq!(message.get_account_key(5), Some(&keys[5]));
-    }
-
     #[test]
     fn test_is_writable_index() {
         let message = check_test_loaded_message().0;