use callback for account and account_default_if_zero_lamport (#819)

use callback for `account` and `account_default_if_zero_lamport`

accounts-db/src/account_storage/meta.rs

@@ -73,8 +73,8 @@ impl<'a: 'b, 'b, U: StorableAccounts<'a>, V: Borrow<AccountHash>>
             let item = self.hashes.as_ref().unwrap();
             item[index].borrow()
         };
-        let account = self.accounts.account_default_if_zero_lamport(index);
-        callback((account, pubkey, hash))
+        self.accounts
+            .account_default_if_zero_lamport(index, |account| callback((account, pubkey, hash)))
     }
 
     /// None if account at index has lamports == 0
@@ -83,9 +83,10 @@ impl<'a: 'b, 'b, U: StorableAccounts<'a>, V: Borrow<AccountHash>>
     pub fn account<Ret>(
         &self,
         index: usize,
-        mut callback: impl FnMut(Option<AccountForStorage<'a>>) -> Ret,
+        callback: impl FnMut(Option<AccountForStorage<'a>>) -> Ret,
     ) -> Ret {
-        callback(self.accounts.account_default_if_zero_lamport(index))
+        self.accounts
+            .account_default_if_zero_lamport(index, callback)
     }
 
     /// # accounts to write

accounts-db/src/accounts_db.rs

@@ -6397,29 +6397,31 @@ impl AccountsDb {
         txn_iter
             .enumerate()
             .map(|(i, txn)| {
-                let account = accounts_and_meta_to_store
-                    .account_default_if_zero_lamport(i)
-                    .map(|account| account.to_account_shared_data())
-                    .unwrap_or_default();
-                let pubkey = accounts_and_meta_to_store.pubkey(i);
-                let account_info = AccountInfo::new(StorageLocation::Cached, account.lamports());
+                let mut account_info = AccountInfo::default();
+                accounts_and_meta_to_store.account_default_if_zero_lamport(i, |account| {
+                    let account = account
+                        .map(|account| account.to_account_shared_data())
+                        .unwrap_or_default();
+                    let pubkey = accounts_and_meta_to_store.pubkey(i);
+                    account_info = AccountInfo::new(StorageLocation::Cached, account.lamports());
 
-                self.notify_account_at_accounts_update(
-                    slot,
-                    &account,
-                    txn,
-                    pubkey,
-                    &mut write_version_producer,
-                );
+                    self.notify_account_at_accounts_update(
+                        slot,
+                        &account,
+                        txn,
+                        pubkey,
+                        &mut write_version_producer,
+                    );
 
-                let cached_account = self.accounts_cache.store(slot, pubkey, account);
-                // hash this account in the bg
-                match &self.sender_bg_hasher {
-                    Some(ref sender) => {
-                        let _ = sender.send(cached_account);
-                    }
-                    None => (),
-                };
+                    let cached_account = self.accounts_cache.store(slot, pubkey, account);
+                    // hash this account in the bg
+                    match &self.sender_bg_hasher {
+                        Some(ref sender) => {
+                            let _ = sender.send(cached_account);
+                        }
+                        None => (),
+                    };
+                });
                 account_info
             })
             .collect()
@@ -6484,10 +6486,11 @@ impl AccountsDb {
                             let len = accounts.len();
                             let mut hashes = Vec::with_capacity(len);
                             for index in 0..accounts.len() {
-                                let (pubkey, account) =
-                                    (accounts.pubkey(index), accounts.account(index));
-                                let hash = Self::hash_account(&account, pubkey);
-                                hashes.push(hash);
+                                let pubkey = accounts.pubkey(index);
+                                accounts.account(index, |account| {
+                                    let hash = Self::hash_account(&account, pubkey);
+                                    hashes.push(hash);
+                                });
                             }
                             hash_time.stop();
                             self.stats
@@ -7785,19 +7788,20 @@ impl AccountsDb {
 
             (start..end).for_each(|i| {
                 let info = infos[i];
-                let pubkey_account = (accounts.pubkey(i), accounts.account(i));
-                let pubkey = pubkey_account.0;
-                let old_slot = accounts.slot(i);
-                self.accounts_index.upsert(
-                    target_slot,
-                    old_slot,
-                    pubkey,
-                    &pubkey_account.1,
-                    &self.account_indexes,
-                    info,
-                    &mut reclaims,
-                    reclaim,
-                );
+                let pubkey = accounts.pubkey(i);
+                accounts.account(i, |account| {
+                    let old_slot = accounts.slot(i);
+                    self.accounts_index.upsert(
+                        target_slot,
+                        old_slot,
+                        pubkey,
+                        &account,
+                        &self.account_indexes,
+                        info,
+                        &mut reclaims,
+                        reclaim,
+                    );
+                });
             });
             reclaims
         };
@@ -8213,9 +8217,10 @@ impl AccountsDb {
         let mut stats = BankHashStats::default();
         let mut total_data = 0;
         (0..accounts.len()).for_each(|index| {
-            let account = accounts.account(index);
-            total_data += account.data().len();
-            stats.update(&account);
+            accounts.account(index, |account| {
+                total_data += account.data().len();
+                stats.update(&account);
+            })
         });
 
         self.stats
@@ -9637,8 +9642,12 @@ pub mod tests {
         fn pubkey(&self, index: usize) -> &Pubkey {
             self.1[index].0
         }
-        fn account(&self, index: usize) -> AccountForStorage<'a> {
-            self.1[index].1.into()
+        fn account<Ret>(
+            &self,
+            index: usize,
+            mut callback: impl FnMut(AccountForStorage<'a>) -> Ret,
+        ) -> Ret {
+            callback(self.1[index].1.into())
         }
         fn slot(&self, index: usize) -> Slot {
             // note that this could be different than 'target_slot()' PER account

accounts-db/src/stake_rewards.rs

@@ -25,9 +25,13 @@ impl<'a> StorableAccounts<'a> for (Slot, &'a [StakeReward]) {
     fn pubkey(&self, index: usize) -> &Pubkey {
         &self.1[index].stake_pubkey
     }
-    fn account(&self, index: usize) -> AccountForStorage<'a> {
+    fn account<Ret>(
+        &self,
+        index: usize,
+        mut callback: impl FnMut(AccountForStorage<'a>) -> Ret,
+    ) -> Ret {
         let entry = &self.1[index];
-        (&entry.stake_account).into()
+        callback((&entry.stake_account).into())
     }
     fn slot(&self, _index: usize) -> Slot {
         // per-index slot is not unique per slot when per-account slot is not included in the source data

accounts-db/src/storable_accounts.rs

@@ -84,11 +84,17 @@ pub trait StorableAccounts<'a>: Sync {
     /// pubkey at 'index'
     fn pubkey(&self, index: usize) -> &Pubkey;
     /// account at 'index'
-    fn account(&self, index: usize) -> AccountForStorage<'a>;
+    fn account<Ret>(&self, index: usize, callback: impl FnMut(AccountForStorage<'a>) -> Ret)
+        -> Ret;
     /// None if account is zero lamports
-    fn account_default_if_zero_lamport(&self, index: usize) -> Option<AccountForStorage<'a>> {
-        let account = self.account(index);
-        (account.lamports() != 0).then_some(account)
+    fn account_default_if_zero_lamport<Ret>(
+        &self,
+        index: usize,
+        mut callback: impl FnMut(Option<AccountForStorage<'a>>) -> Ret,
+    ) -> Ret {
+        self.account(index, |account| {
+            callback((!account.is_zero_lamport()).then_some(account))
+        })
     }
     // current slot for account at 'index'
     fn slot(&self, index: usize) -> Slot;
@@ -139,8 +145,12 @@ where
     fn pubkey(&self, index: usize) -> &Pubkey {
         self.accounts[index].0
     }
-    fn account(&self, index: usize) -> AccountForStorage<'a> {
-        self.accounts[index].1.into()
+    fn account<Ret>(
+        &self,
+        index: usize,
+        mut callback: impl FnMut(AccountForStorage<'a>) -> Ret,
+    ) -> Ret {
+        callback(self.accounts[index].1.into())
     }
     fn slot(&self, _index: usize) -> Slot {
         // per-index slot is not unique per slot, but it is different than 'target_slot'
@@ -162,8 +172,12 @@ where
     fn pubkey(&self, index: usize) -> &Pubkey {
         self.1[index].0
     }
-    fn account(&self, index: usize) -> AccountForStorage<'a> {
-        self.1[index].1.into()
+    fn account<Ret>(
+        &self,
+        index: usize,
+        mut callback: impl FnMut(AccountForStorage<'a>) -> Ret,
+    ) -> Ret {
+        callback(self.1[index].1.into())
     }
     fn slot(&self, _index: usize) -> Slot {
         // per-index slot is not unique per slot when per-account slot is not included in the source data
@@ -183,8 +197,12 @@ where
     fn pubkey(&self, index: usize) -> &Pubkey {
         &self.1[index].0
     }
-    fn account(&self, index: usize) -> AccountForStorage<'a> {
-        (&self.1[index].1).into()
+    fn account<Ret>(
+        &self,
+        index: usize,
+        mut callback: impl FnMut(AccountForStorage<'a>) -> Ret,
+    ) -> Ret {
+        callback((&self.1[index].1).into())
     }
     fn slot(&self, _index: usize) -> Slot {
         // per-index slot is not unique per slot when per-account slot is not included in the source data
@@ -202,8 +220,12 @@ impl<'a> StorableAccounts<'a> for (Slot, &'a [&'a StoredAccountMeta<'a>]) {
     fn pubkey(&self, index: usize) -> &Pubkey {
         self.1[index].pubkey()
     }
-    fn account(&self, index: usize) -> AccountForStorage<'a> {
-        self.1[index].into()
+    fn account<Ret>(
+        &self,
+        index: usize,
+        mut callback: impl FnMut(AccountForStorage<'a>) -> Ret,
+    ) -> Ret {
+        callback(self.1[index].into())
     }
     fn slot(&self, _index: usize) -> Slot {
         // per-index slot is not unique per slot when per-account slot is not included in the source data
@@ -291,9 +313,13 @@ impl<'a> StorableAccounts<'a> for StorableAccountsBySlot<'a> {
         let indexes = self.find_internal_index(index);
         self.slots_and_accounts[indexes.0].1[indexes.1].pubkey()
     }
-    fn account(&self, index: usize) -> AccountForStorage<'a> {
+    fn account<Ret>(
+        &self,
+        index: usize,
+        mut callback: impl FnMut(AccountForStorage<'a>) -> Ret,
+    ) -> Ret {
         let indexes = self.find_internal_index(index);
-        self.slots_and_accounts[indexes.0].1[indexes.1].into()
+        callback(self.slots_and_accounts[indexes.0].1[indexes.1].into())
     }
     fn slot(&self, index: usize) -> Slot {
         let indexes = self.find_internal_index(index);
@@ -323,8 +349,12 @@ impl<'a> StorableAccounts<'a> for (Slot, &'a [&'a StoredAccountMeta<'a>], Slot)
     fn pubkey(&self, index: usize) -> &Pubkey {
         self.1[index].pubkey()
     }
-    fn account(&self, index: usize) -> AccountForStorage<'a> {
-        self.1[index].into()
+    fn account<Ret>(
+        &self,
+        index: usize,
+        mut callback: impl FnMut(AccountForStorage<'a>) -> Ret,
+    ) -> Ret {
+        callback(self.1[index].into())
     }
     fn slot(&self, _index: usize) -> Slot {
         // same other slot for all accounts
@@ -364,7 +394,11 @@ pub mod tests {
         assert_eq!(a.is_empty(), b.is_empty());
         (0..a.len()).for_each(|i| {
             assert_eq!(a.pubkey(i), b.pubkey(i));
-            assert!(accounts_equal(&a.account(i), &b.account(i)));
+            b.account(i, |account| {
+                a.account(i, |account_a| {
+                    assert!(accounts_equal(&account_a, &account));
+                });
+            });
         })
     }
 
@@ -486,7 +520,9 @@ pub mod tests {
                     compare(&test2, &test_moving_slots2);
                     for (i, raw) in raw.iter().enumerate() {
                         assert_eq!(raw.0, *test3.pubkey(i));
-                        assert!(accounts_equal(&raw.1, &test3.account(i)));
+                        test3.account(i, |account| {
+                            assert!(accounts_equal(&raw.1, &account));
+                        });
                         assert_eq!(raw.2, test3.slot(i));
                         assert_eq!(target_slot, test4.slot(i));
                         assert_eq!(target_slot, test2.slot(i));
@@ -587,7 +623,9 @@ pub mod tests {
                         assert_eq!(entries0 != entries, storable.contains_multiple_slots());
                         (0..entries).for_each(|index| {
                             let index = index as usize;
-                            assert!(accounts_equal(&storable.account(index), &raw2[index]));
+                            storable.account(index, |account| {
+                                assert!(accounts_equal(&account, &raw2[index]));
+                            });
                             assert_eq!(storable.pubkey(index), raw2[index].pubkey());
                             assert_eq!(storable.hash(index), raw2[index].hash());
                             assert_eq!(storable.slot(index), expected_slots[index]);

runtime/src/bank.rs

@@ -5045,11 +5045,13 @@ impl Bank {
         let mut m = Measure::start("stakes_cache.check_and_store");
         let new_warmup_cooldown_rate_epoch = self.new_warmup_cooldown_rate_epoch();
         (0..accounts.len()).for_each(|i| {
-            self.stakes_cache.check_and_store(
-                accounts.pubkey(i),
-                &accounts.account(i),
-                new_warmup_cooldown_rate_epoch,
-            )
+            accounts.account(i, |account| {
+                self.stakes_cache.check_and_store(
+                    accounts.pubkey(i),
+                    &account,
+                    new_warmup_cooldown_rate_epoch,
+                )
+            })
         });
         self.rc.accounts.store_accounts_cached(accounts);
         m.stop();