modify for all cases

runtime/src/bank.rs

@@ -8062,6 +8062,26 @@ impl Bank {
         }
     }
 
+    fn replace_account(
+        &mut self,
+        old_address: &Pubkey,
+        new_address: &Pubkey,
+        old_account: Option<&AccountSharedData>,
+        new_account: &AccountSharedData,
+    ) {
+        let (old_lamports, old_len) = match old_account {
+            Some(old_account) => (old_account.lamports(), old_account.data().len()),
+            None => (0, 0),
+        };
+        self.capitalization.fetch_sub(old_lamports, Relaxed);
+        self.store_account(old_address, new_account);
+        self.store_account(new_address, &AccountSharedData::default());
+        self.calculate_and_update_accounts_data_size_delta_off_chain(
+            old_len,
+            new_account.data().len(),
+        );
+    }
+
     /// Use to replace programs by feature activation
     #[allow(dead_code)]
     fn replace_program_account(
@@ -8070,46 +8090,95 @@ impl Bank {
         new_address: &Pubkey,
         datapoint_name: &'static str,
     ) {
-        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());
-                        }
+        // Both program accounts must exist in order to attempt a replacement
+        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) {
+                // Both exist, so we can proceed with the replacement
+                datapoint_info!(datapoint_name, ("slot", self.slot, i64));
+
+                // Derive each program's data account address (PDA)
+                let (old_data_address, _) = Pubkey::find_program_address(
+                    &[old_address.as_ref()],
+                    &bpf_loader_upgradeable::id(),
+                );
+                let (new_data_address, _) = Pubkey::find_program_address(
+                    &[new_address.as_ref()],
+                    &bpf_loader_upgradeable::id(),
+                );
 
-                        self.calculate_and_update_accounts_data_size_delta_off_chain(
-                            old_account.data().len(),
-                            new_account.data().len(),
+                // If a data account is also provided with this new program
+                // account, then we want to update the existing data account
+                if let Some(new_data_account) = self.get_account_with_fixed_root(&new_data_address)
+                {
+                    // A data account exists for the new program
+                    // Check if the old program account has a data account
+                    if let Some(old_data_account) =
+                        self.get_account_with_fixed_root(&old_data_address)
+                    {
+                        // It does. Replace it with the new data account
+                        self.replace_account(
+                            &old_data_address,
+                            &new_data_address,
+                            Some(&old_data_account),
+                            &new_data_account,
+                        );
+                        // The old program account will already house the PDA
+                        // of the data account
+                    } else {
+                        // It does _not_. Create it with the new data account
+                        self.replace_account(
+                            &old_data_address,
+                            &new_data_address,
+                            None,
+                            &new_data_account,
                         );
+                        // Update the old program account to house the PDA of
+                        // the data account
+                        {
+                            let mut account = Account::from(old_account.clone());
+                            account.data = old_data_address.as_ref().to_vec();
+                            let account_shared_data = AccountSharedData::from(account);
+                            self.store_account(old_address, &account_shared_data);
+                        }
                     }
+                    // We only swapped the data accounts, so now we need to
+                    // clear the new program account
+                    self.capitalization
+                        .fetch_sub(new_account.lamports(), Relaxed);
+                    self.store_account(new_address, &AccountSharedData::default());
+                } else if let Some(old_data_account) =
+                    self.get_account_with_fixed_root(&old_data_address)
+                {
+                    // A data account exists for the old program, but not the
+                    // new program
+                    // Swap program accounts and delete the old data account
+                    self.replace_account(
+                        old_address,
+                        new_address,
+                        Some(&old_account),
+                        &new_account,
+                    );
+                    self.capitalization
+                        .fetch_sub(old_data_account.lamports(), Relaxed);
+                    self.store_account(&old_data_address, &AccountSharedData::default());
+                } else {
+                    // A data account does not exist for the new program
+                    // Swap program accounts only
+                    self.replace_account(
+                        old_address,
+                        new_address,
+                        Some(&old_account),
+                        &new_account,
+                    );
                 }
-            };
-        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 */);
+
+                // Unload a program from the bank's cache
+                self.loaded_programs_cache
+                    .write()
+                    .unwrap()
+                    .remove_programs([*old_address].into_iter());
+            }
+        }
     }
 
     /// Get all the accounts for this bank and calculate stats

runtime/src/bank/tests.rs

@@ -8013,9 +8013,15 @@ 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 {
+fn set_up_account_with_bank(
+    bank: &mut Bank,
+    pubkey: &Pubkey,
+    lamports: u64,
+    data: Vec<u8>,
+) -> AccountSharedData {
     let new_account = AccountSharedData::from(Account {
         lamports,
+        data,
         ..Account::default()
     });
     bank.store_account_and_update_capitalization(pubkey, &new_account);
@@ -8024,41 +8030,180 @@ fn set_up_account_with_bank(bank: &mut Bank, pubkey: &Pubkey, lamports: u64) ->
 }
 
 #[test]
-fn test_program_replacement() {
+fn test_non_upgradable_program_replace() {
+    // Non-Upgradable program
+    // - Old:     [Old program data]
+    // - New:     [*New program data]
+    //
+    // Should replace the old program account with the new program account:
+    // - Old:     [*New program data]
     let mut bank = create_simple_test_bank(0);
 
-    // Setup original program accounts
-    let old_address = Pubkey::new_unique();
-    set_up_account_with_bank(&mut bank, &old_address, 100);
+    let old = Pubkey::new_unique();
+    set_up_account_with_bank(&mut bank, &old, 100, vec![0, 0, 0, 0]);
 
-    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);
+    let new = Pubkey::new_unique();
+    let new_program_bytes = vec![6, 5, 4, 3, 2, 1, 0];
+    set_up_account_with_bank(&mut bank, &new, 100, new_program_bytes.clone());
 
-    // Setup new program accounts
-    let new_address = Pubkey::new_unique();
-    let new_program_account = set_up_account_with_bank(&mut bank, &new_address, 123);
+    let original_capitalization = bank.capitalization();
+
+    bank.replace_program_account(&old, &new, "bank-apply_program_replacement");
+
+    // Old program account balance is unchanged
+    assert_eq!(bank.get_balance(&old), 100);
+
+    // New program account is now empty
+    assert_eq!(bank.get_balance(&new), 0);
+
+    // Old program account now holds the new program data, ie:
+    // - Old:     [*New program data]
+    let old_account = bank.get_account(&old).unwrap();
+    assert_eq!(old_account.data(), &new_program_bytes,);
+
+    // Lamports in the old token account was burnt
+    assert_eq!(bank.capitalization(), original_capitalization - 100);
+}
+
+#[test]
+fn test_non_upgradable_program_replace_with_data() {
+    // Non-Upgradable program
+    // - Old:     [Old program data]
+    // - New:     PDA(NewData)
+    // - NewData: [*New program data]
+    //
+    // Should replace the program account with the PDA of the data account,
+    // and create the data account:
+    // - Old:     PDA(OldData)
+    // - OldData: [*New program data]
+    let bpf_id = bpf_loader_upgradeable::id();
+    let mut bank = create_simple_test_bank(0);
+
+    let old = Pubkey::new_unique();
+    let (old_data, _) = Pubkey::find_program_address(&[old.as_ref()], &bpf_id);
+    set_up_account_with_bank(&mut bank, &old, 100, vec![0, 0, 0, 0]);
 
-    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 new = Pubkey::new_unique();
+    let (new_data, _) = Pubkey::find_program_address(&[new.as_ref()], &bpf_id);
+    let new_program_bytes = vec![6, 5, 4, 3, 2, 1, 0];
+    set_up_account_with_bank(&mut bank, &new, 100, new_data.to_bytes().to_vec());
+    set_up_account_with_bank(&mut bank, &new_data, 102, new_program_bytes.clone());
 
     let original_capitalization = bank.capitalization();
 
-    bank.replace_program_account(&old_address, &new_address, "bank-apply_program_replacement");
+    bank.replace_program_account(&old, &new, "bank-apply_program_replacement");
+
+    // Old program account balances are unchanged
+    assert_eq!(bank.get_balance(&old), 100);
+    assert_eq!(bank.get_balance(&old_data), 102);
 
     // New program accounts are now empty
-    assert_eq!(bank.get_balance(&new_address), 0);
-    assert_eq!(bank.get_balance(&new_data_address), 0);
+    assert_eq!(bank.get_balance(&new), 0);
+    assert_eq!(bank.get_balance(&new_data), 0);
 
-    // Old program account holds the new program account
-    assert_eq!(bank.get_account(&old_address), Some(new_program_account));
+    // Old program account now holds the PDA, ie:
+    // - Old:     PDA(OldData)
+    let old_account = bank.get_account(&old).unwrap();
+    assert_eq!(old_account.data(), &old_data.to_bytes().to_vec(),);
 
-    // Old program data account holds the new program data account
-    assert_eq!(
-        bank.get_account(&old_data_address),
-        Some(new_program_data_account)
-    );
+    // Old program data account has been created & now holds the new data, ie:
+    // - OldData: [*New program data]
+    let old_data_account = bank.get_account(&old_data).unwrap();
+    assert_eq!(old_data_account.data(), &new_program_bytes,);
+
+    // Lamports in the old token accounts were burnt
+    assert_eq!(bank.capitalization(), original_capitalization - 100);
+}
+
+#[test]
+fn test_upgradable_program_replace() {
+    // Upgradable program
+    // - Old:     PDA(OldData)
+    // - OldData: [Old program data]
+    // - New:     PDA(NewData)
+    // - NewData: [*New program data]
+    //
+    // Should _only_ replace the data account, not the program account:
+    // - Old:     PDA(OldData)
+    // - OldData: [*New program data]
+    let bpf_id = bpf_loader_upgradeable::id();
+    let mut bank = create_simple_test_bank(0);
+
+    let old = Pubkey::new_unique();
+    let (old_data, _) = Pubkey::find_program_address(&[old.as_ref()], &bpf_id);
+    set_up_account_with_bank(&mut bank, &old, 100, old_data.to_bytes().to_vec());
+    set_up_account_with_bank(&mut bank, &old_data, 102, vec![0, 1, 2, 3, 4, 5, 6]);
+
+    let new = Pubkey::new_unique();
+    let (new_data, _) = Pubkey::find_program_address(&[new.as_ref()], &bpf_id);
+    let new_program_bytes = vec![6, 5, 4, 3, 2, 1, 0];
+    set_up_account_with_bank(&mut bank, &new, 100, new_data.to_bytes().to_vec());
+    set_up_account_with_bank(&mut bank, &new_data, 102, new_program_bytes.clone());
+
+    let original_capitalization = bank.capitalization();
+
+    bank.replace_program_account(&old, &new, "bank-apply_program_replacement");
+
+    // Old program account balances are unchanged
+    assert_eq!(bank.get_balance(&old), 100);
+    assert_eq!(bank.get_balance(&old_data), 102);
+
+    // New program accounts are now empty
+    assert_eq!(bank.get_balance(&new), 0);
+    assert_eq!(bank.get_balance(&new_data), 0);
+
+    // Old program account still holds the same PDA, ie:
+    // - Old:     PDA(OldData)
+    let old_account = bank.get_account(&old).unwrap();
+    assert_eq!(old_account.data(), &old_data.to_bytes().to_vec(),);
+
+    // Old program data account now holds the new data, ie:
+    // - OldData: [*New program data]
+    let old_data_account = bank.get_account(&old_data).unwrap();
+    assert_eq!(old_data_account.data(), &new_program_bytes,);
+
+    // Lamports in the old token accounts were burnt
+    assert_eq!(bank.capitalization(), original_capitalization - 100 - 102);
+}
+
+#[test]
+fn test_upgradable_program_replace_with_no_data() {
+    // Upgradable program
+    // - Old:     PDA(OldData)
+    // - OldData: [Old program data]
+    // - New:     [*New program data]
+    //
+    // Should replace the program account and delete the data account:
+    // - Old:     [*New program data]
+    let bpf_id = bpf_loader_upgradeable::id();
+    let mut bank = create_simple_test_bank(0);
+
+    let old = Pubkey::new_unique();
+    let (old_data, _) = Pubkey::find_program_address(&[old.as_ref()], &bpf_id);
+    set_up_account_with_bank(&mut bank, &old, 100, old_data.to_bytes().to_vec());
+    set_up_account_with_bank(&mut bank, &old_data, 102, vec![0, 1, 2, 3, 4, 5, 6]);
+
+    let new = Pubkey::new_unique();
+    let new_program_bytes = vec![6, 5, 4, 3, 2, 1, 0];
+    set_up_account_with_bank(&mut bank, &new, 100, new_program_bytes.clone());
+
+    let original_capitalization = bank.capitalization();
+
+    bank.replace_program_account(&old, &new, "bank-apply_program_replacement");
+
+    // Old program account balance is unchanged
+    assert_eq!(bank.get_balance(&old), 100);
+
+    // Old data account is now empty
+    assert_eq!(bank.get_balance(&old_data), 0);
+
+    // New program account is now empty
+    assert_eq!(bank.get_balance(&new), 0);
+
+    // Old program account now holds the new program bytes
+    // - Old:     [*New program data]
+    let old_account = bank.get_account(&old).unwrap();
+    assert_eq!(old_account.data(), &new_program_bytes,);
 
     // Lamports in the old token accounts were burnt
     assert_eq!(bank.capitalization(), original_capitalization - 100 - 102);