From b053bc2790b81531f3561a1d903b76e0fc3be9e7 Mon Sep 17 00:00:00 2001
From: Tyera Eulberg <tyera@solana.com>
Date: Thu, 7 Mar 2019 09:51:56 -0700
Subject: [PATCH] Load accounts by program owner for program subscriptions

---
 runtime/src/accounts.rs | 64 +++++++++++++++++++++++++++++++++++++++++
 runtime/src/bank.rs     |  8 ++++++
 2 files changed, 72 insertions(+)

diff --git a/runtime/src/accounts.rs b/runtime/src/accounts.rs
index 2bf88a2d8b..6547aa3061 100644
--- a/runtime/src/accounts.rs
+++ b/runtime/src/accounts.rs
@@ -381,6 +381,25 @@ impl AccountsDB {
         None
     }
 
+    fn load_by_program(
+        &self,
+        fork: Fork,
+        program_id: &Pubkey,
+        walk_back: bool,
+    ) -> Vec<(Pubkey, Account)> {
+        self.account_index
+            .account_maps
+            .read()
+            .unwrap()
+            .iter()
+            .filter_map(|(pubkey, _)| {
+                self.load(fork, pubkey, walk_back)
+                    .filter(|account| account.owner == *program_id)
+                    .map(|account| (*pubkey, account))
+            })
+            .collect()
+    }
+
     fn get_storage_id(&self, start: usize, current: usize) -> usize {
         let mut id = current;
         let len: usize;
@@ -837,6 +856,19 @@ impl Accounts {
             .filter(|acc| acc.lamports != 0)
     }
 
+    /// Slow because lock is held for 1 operation insted of many
+    pub fn load_by_program_slow_no_parent(
+        &self,
+        fork: Fork,
+        program_id: &Pubkey,
+    ) -> Vec<(Pubkey, Account)> {
+        self.accounts_db
+            .load_by_program(fork, program_id, false)
+            .into_iter()
+            .filter(|(_, acc)| acc.lamports != 0)
+            .collect()
+    }
+
     /// Slow because lock is held for 1 operation insted of many
     pub fn store_slow(&self, fork: Fork, pubkey: &Pubkey, account: &Account) {
         self.accounts_db.store(fork, pubkey, account);
@@ -1854,4 +1886,36 @@ mod tests {
         accounts.squash(1);
         assert_eq!(accounts.transaction_count(1), 2);
     }
+
+    #[test]
+    fn test_load_by_program() {
+        let paths = get_tmp_accounts_path!();
+        let accounts_db = AccountsDB::new(0, &paths.paths);
+
+        // Load accounts owned by various programs into AccountsDB
+        let pubkey0 = Keypair::new().pubkey();
+        let account0 = Account::new(1, 0, Pubkey::new(&[2; 32]));
+        accounts_db.store(0, &pubkey0, &account0);
+        let pubkey1 = Keypair::new().pubkey();
+        let account1 = Account::new(1, 0, Pubkey::new(&[2; 32]));
+        accounts_db.store(0, &pubkey1, &account1);
+        let pubkey2 = Keypair::new().pubkey();
+        let account2 = Account::new(1, 0, Pubkey::new(&[3; 32]));
+        accounts_db.store(0, &pubkey2, &account2);
+
+        let accounts = accounts_db.load_by_program(0, &Pubkey::new(&[2; 32]), false);
+        assert_eq!(accounts.len(), 2);
+        let accounts = accounts_db.load_by_program(0, &Pubkey::new(&[3; 32]), false);
+        assert_eq!(accounts, vec![(pubkey2, account2)]);
+        let accounts = accounts_db.load_by_program(0, &Pubkey::new(&[4; 32]), false);
+        assert_eq!(accounts, vec![]);
+
+        // Accounts method
+        let mut accounts_proper = Accounts::new(0, None);
+        accounts_proper.accounts_db = accounts_db;
+        let accounts = accounts_proper.load_by_program_slow_no_parent(0, &Pubkey::new(&[2; 32]));
+        assert_eq!(accounts.len(), 2);
+        let accounts = accounts_proper.load_by_program_slow_no_parent(0, &Pubkey::new(&[4; 32]));
+        assert_eq!(accounts, vec![]);
+    }
 }
diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index cf378f9fcf..1e7223c001 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -796,6 +796,14 @@ impl Bank {
         self.accounts().load_slow(self.accounts_id, pubkey)
     }
 
+    pub fn get_program_accounts_modified_since_parent(
+        &self,
+        program_id: &Pubkey,
+    ) -> Vec<(Pubkey, Account)> {
+        self.accounts()
+            .load_by_program_slow_no_parent(self.accounts_id, program_id)
+    }
+
     pub fn get_account_modified_since_parent(&self, pubkey: &Pubkey) -> Option<Account> {
         self.accounts()
             .load_slow_no_parent(self.accounts_id, pubkey)