From 44831c18d29f5c43feb0ffdb0abe258be34d4a99 Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)"
 <75863576+jeffwashington@users.noreply.github.com>
Date: Fri, 21 May 2021 10:07:04 -0500
Subject: [PATCH] reuse work in is_non_loader_key (#16521)

---
 runtime/src/accounts.rs    |  6 +++---
 sdk/program/src/message.rs | 35 ++++++++++++++++++++++++++++++++++-
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/runtime/src/accounts.rs b/runtime/src/accounts.rs
index c07c060dd7..1e15a31feb 100644
--- a/runtime/src/accounts.rs
+++ b/runtime/src/accounts.rs
@@ -26,7 +26,7 @@ use solana_sdk::{
     fee_calculator::{FeeCalculator, FeeConfig},
     genesis_config::ClusterType,
     hash::Hash,
-    message::Message,
+    message::{Message, MessageProgramIdsCache},
     native_loader, nonce,
     pubkey::Pubkey,
     transaction::Result,
@@ -205,9 +205,9 @@ impl Accounts {
             let mut account_deps = Vec::with_capacity(message.account_keys.len());
             let demote_sysvar_write_locks =
                 feature_set.is_active(&feature_set::demote_sysvar_write_locks::id());
-
+            let mut key_check = MessageProgramIdsCache::new(&message);
             for (i, key) in message.account_keys.iter().enumerate() {
-                let account = if message.is_non_loader_key(key, i) {
+                let account = if key_check.is_non_loader_key(key, i) {
                     if payer_index.is_none() {
                         payer_index = Some(i);
                     }
diff --git a/sdk/program/src/message.rs b/sdk/program/src/message.rs
index 98728dc7d9..a8b0aa5fe8 100644
--- a/sdk/program/src/message.rs
+++ b/sdk/program/src/message.rs
@@ -240,6 +240,26 @@ impl Sanitize for Message {
         Ok(())
     }
 }
+pub struct MessageProgramIdsCache<'a> {
+    program_ids: Option<Vec<&'a Pubkey>>,
+    message: &'a Message,
+}
+
+impl<'a> MessageProgramIdsCache<'a> {
+    pub fn new(message: &'a Message) -> Self {
+        Self {
+            program_ids: None,
+            message,
+        }
+    }
+    pub fn is_non_loader_key(&mut self, key: &Pubkey, key_index: usize) -> bool {
+        if self.program_ids.is_none() {
+            self.program_ids = Some(self.message.program_ids());
+        }
+        self.message
+            .is_non_loader_key_internal(self.program_ids.as_ref().unwrap(), key, key_index)
+    }
+}
 
 impl Message {
     pub fn new_with_compiled_instructions(
@@ -349,8 +369,17 @@ impl Message {
         false
     }
 
+    pub(crate) fn is_non_loader_key_internal(
+        &self,
+        program_ids: &[&Pubkey],
+        key: &Pubkey,
+        key_index: usize,
+    ) -> bool {
+        !program_ids.contains(&key) || self.is_key_passed_to_program(key_index)
+    }
+
     pub fn is_non_loader_key(&self, key: &Pubkey, key_index: usize) -> bool {
-        !self.program_ids().contains(&key) || self.is_key_passed_to_program(key_index)
+        self.is_non_loader_key_internal(&self.program_ids(), key, key_index)
     }
 
     pub fn program_position(&self, index: usize) -> Option<usize> {
@@ -1013,10 +1042,14 @@ mod tests {
             Hash::default(),
             instructions,
         );
+        let mut helper = MessageProgramIdsCache::new(&message);
 
         assert!(message.is_non_loader_key(&key0, 0));
         assert!(message.is_non_loader_key(&key1, 1));
         assert!(!message.is_non_loader_key(&loader2, 2));
+        assert!(helper.is_non_loader_key(&key0, 0));
+        assert!(helper.is_non_loader_key(&key1, 1));
+        assert!(!helper.is_non_loader_key(&loader2, 2));
     }
 
     #[test]