replace program account

runtime/src/bank.rs

@@ -8070,32 +8070,46 @@ impl Bank {
         new_address: &Pubkey,
         datapoint_name: &'static str,
     ) {
-        if let Some(old_account) = self.get_account_with_fixed_root(old_address) {
-            if let Some(new_account) = self.get_account_with_fixed_root(new_address) {
-                datapoint_info!(datapoint_name, ("slot", self.slot, i64));
-
-                // Burn lamports in the old account
-                self.capitalization
-                    .fetch_sub(old_account.lamports(), Relaxed);
-
-                // Transfer new account to old account
-                self.store_account(old_address, &new_account);
-
-                // Clear new account
-                self.store_account(new_address, &AccountSharedData::default());
-
-                // Unload a program from the bank's cache
-                self.loaded_programs_cache
-                    .write()
-                    .unwrap()
-                    .remove_programs([*old_address].into_iter());
+        let program_data_address = |address: &Pubkey| {
+            Pubkey::find_program_address(&[address.as_ref()], &bpf_loader_upgradeable::id()).0
+        };
+        let maybe_replace_program_account =
+            |old_address: &Pubkey, new_address: &Pubkey, is_program: bool| {
+                if let Some(old_account) = self.get_account_with_fixed_root(old_address) {
+                    if let Some(new_account) = self.get_account_with_fixed_root(new_address) {
+                        datapoint_info!(datapoint_name, ("slot", self.slot, i64));
+
+                        // Burn lamports in the old account
+                        self.capitalization
+                            .fetch_sub(old_account.lamports(), Relaxed);
+
+                        // Transfer new account to old account
+                        self.store_account(old_address, &new_account);
+
+                        // Clear new account
+                        self.store_account(new_address, &AccountSharedData::default());
+
+                        // Unload a program from the bank's cache
+                        if is_program {
+                            self.loaded_programs_cache
+                                .write()
+                                .unwrap()
+                                .remove_programs([*old_address].into_iter());
+                        }
 
-                self.calculate_and_update_accounts_data_size_delta_off_chain(
-                    old_account.data().len(),
-                    new_account.data().len(),
-                );
-            }
-        }
+                        self.calculate_and_update_accounts_data_size_delta_off_chain(
+                            old_account.data().len(),
+                            new_account.data().len(),
+                        );
+                    }
+                }
+            };
+        maybe_replace_program_account(
+            &program_data_address(old_address),
+            &program_data_address(new_address),
+            false, /* is_program */
+        );
+        maybe_replace_program_account(old_address, new_address, true /* is_program */);
     }
 
     /// Get all the accounts for this bank and calculate stats

runtime/src/bank/tests.rs

@@ -8013,42 +8013,55 @@ fn test_compute_active_feature_set() {
     assert!(bank.feature_set.is_active(&test_feature));
 }
 
+fn set_up_account_with_bank(bank: &mut Bank, pubkey: &Pubkey, lamports: u64) -> AccountSharedData {
+    let new_account = AccountSharedData::from(Account {
+        lamports,
+        ..Account::default()
+    });
+    bank.store_account_and_update_capitalization(pubkey, &new_account);
+    assert_eq!(bank.get_balance(pubkey), lamports);
+    new_account
+}
+
 #[test]
 fn test_program_replacement() {
     let mut bank = create_simple_test_bank(0);
 
-    // Setup original program account
+    // Setup original program accounts
     let old_address = Pubkey::new_unique();
+    set_up_account_with_bank(&mut bank, &old_address, 100);
+
+    let (old_data_address, _) =
+        Pubkey::find_program_address(&[old_address.as_ref()], &bpf_loader_upgradeable::id());
+    set_up_account_with_bank(&mut bank, &old_data_address, 102);
+
+    // Setup new program accounts
     let new_address = Pubkey::new_unique();
-    bank.store_account_and_update_capitalization(
-        &old_address,
-        &AccountSharedData::from(Account {
-            lamports: 100,
-            ..Account::default()
-        }),
-    );
-    assert_eq!(bank.get_balance(&old_address), 100);
-
-    // Setup new program account
-    let new_program_account = AccountSharedData::from(Account {
-        lamports: 123,
-        ..Account::default()
-    });
-    bank.store_account_and_update_capitalization(&new_address, &new_program_account);
-    assert_eq!(bank.get_balance(&new_address), 123);
+    let new_program_account = set_up_account_with_bank(&mut bank, &new_address, 123);
+
+    let (new_data_address, _) =
+        Pubkey::find_program_address(&[new_address.as_ref()], &bpf_loader_upgradeable::id());
+    let new_program_data_account = set_up_account_with_bank(&mut bank, &new_data_address, 125);
 
     let original_capitalization = bank.capitalization();
 
     bank.replace_program_account(&old_address, &new_address, "bank-apply_program_replacement");
 
-    // New program account is now empty
+    // New program accounts are now empty
     assert_eq!(bank.get_balance(&new_address), 0);
+    assert_eq!(bank.get_balance(&new_data_address), 0);
 
     // Old program account holds the new program account
     assert_eq!(bank.get_account(&old_address), Some(new_program_account));
 
-    // Lamports in the old token account were burnt
-    assert_eq!(bank.capitalization(), original_capitalization - 100);
+    // Old program data account holds the new program data account
+    assert_eq!(
+        bank.get_account(&old_data_address),
+        Some(new_program_data_account)
+    );
+
+    // Lamports in the old token accounts were burnt
+    assert_eq!(bank.capitalization(), original_capitalization - 100 - 102);
 }
 
 fn min_rent_exempt_balance_for_sysvars(bank: &Bank, sysvar_ids: &[Pubkey]) -> u64 {