Simplify State/AccountsStore lifecycle by avoiding default->replace

rs/backend/src/accounts_store.rs

@@ -1,6 +1,7 @@
 //! User accounts and transactions.
+use crate::accounts_store::schema::accounts_in_unbounded_stable_btree_map::AccountsDbAsUnboundedStableBTreeMap;
 use crate::multi_part_transactions_processor::MultiPartTransactionsProcessor;
-use crate::state::StableState;
+use crate::state::{partitions::PartitionType, with_partitions, StableState};
 use crate::stats::Stats;
 use candid::CandidType;
 use dfn_candid::Candid;
@@ -17,7 +18,6 @@ use std::collections::{BTreeMap, HashMap, VecDeque};
 use std::fmt;
 use std::ops::RangeBounds;
 
-pub mod constructors;
 pub mod histogram;
 pub mod schema;
 use schema::{
@@ -318,6 +318,21 @@ pub enum DetachCanisterResponse {
 }
 
 impl AccountsStore {
+    /// Creates a new `AccountsStore`. Should be called during canister `init`.
+    #[must_use]
+    pub fn new() -> Self {
+        let accounts_partition = with_partitions(|p| p.get(PartitionType::Accounts.memory_id()));
+        let accounts_db = AccountsDbAsProxy::from(AccountsDb::UnboundedStableBTreeMap(
+            AccountsDbAsUnboundedStableBTreeMap::new(accounts_partition),
+        ));
+        let accounts_db_stats = AccountsDbStats::default();
+
+        Self {
+            accounts_db,
+            accounts_db_stats,
+        }
+    }
+
     /// Determines whether a migration is being performed.
     #[must_use]
     #[allow(dead_code)]
@@ -777,12 +792,13 @@ impl StableState for AccountsStore {
         let Some(accounts_db_stats) = accounts_db_stats_maybe else {
             return Err("Accounts DB stats should be present since the stable structures migration.".to_string());
         };
+        let accounts_partition = with_partitions(|partitions| partitions.get(PartitionType::Accounts.memory_id()));
+        let accounts_db = AccountsDbAsProxy::from(AccountsDb::UnboundedStableBTreeMap(
+            AccountsDbAsUnboundedStableBTreeMap::load(accounts_partition),
+        ));
 
         Ok(AccountsStore {
-            // Because the stable structures migration is finished, accounts_db
-            // will be replaced with an AccountsDbAsUnboundedStableBTreeMap in
-            // State::from(Partitions) so it doesn't matter what we set here.
-            accounts_db: AccountsDbAsProxy::default(),
+            accounts_db,
             accounts_db_stats,
         })
     }

rs/backend/src/state.rs

@@ -1,18 +1,15 @@
 pub mod partitions;
 #[cfg(test)]
 pub mod tests;
-mod with_accounts_in_stable_memory;
 
-use self::partitions::{PartitionType, Partitions};
-use crate::accounts_store::schema::accounts_in_unbounded_stable_btree_map::AccountsDbAsUnboundedStableBTreeMap;
-use crate::accounts_store::schema::proxy::AccountsDb;
 use crate::accounts_store::AccountsStore;
 use crate::assets::AssetHashes;
 use crate::assets::Assets;
 use crate::perf::PerformanceCounts;
 use crate::tvl::state::TvlState;
 
 use dfn_candid::Candid;
+use dfn_core::api::trap_with;
 use ic_cdk::println;
 use ic_stable_structures::DefaultMemoryImpl;
 use on_wire::{FromWire, IntoWire};
@@ -122,12 +119,8 @@ impl State {
     /// Creates new state. Should be called in `init`.
     #[must_use]
     pub fn new() -> Self {
-        let accounts_partition = with_partitions(|p| p.get(PartitionType::Accounts.memory_id()));
-        let accounts_store = AccountsStore::from(AccountsDb::UnboundedStableBTreeMap(
-            AccountsDbAsUnboundedStableBTreeMap::new(accounts_partition),
-        ));
         State {
-            accounts_store,
+            accounts_store: AccountsStore::new(),
             assets: Assets::default(),
             asset_hashes: AssetHashes::default(),
             performance: PerformanceCounts::default(),
@@ -139,19 +132,18 @@ impl State {
     #[must_use]
     pub fn new_restored() -> Self {
         println!("START state::new_restored: ())");
-        let accounts_partition = with_partitions(|p| p.get(PartitionType::Accounts.memory_id()));
-        let mut state = Self::recover_heap_from_managed_memory();
-        let accounts_db =
-            AccountsDb::UnboundedStableBTreeMap(AccountsDbAsUnboundedStableBTreeMap::load(accounts_partition));
-        // Replace the default accountsdb created by `serde` with the one from stable memory.
-        let _deserialized_accounts_db = state.accounts_store.replace_accounts_db(accounts_db);
+        let bytes = with_partitions(Partitions::read_bytes_from_managed_memory);
+        let state =
+            State::decode(bytes).unwrap_or_else(|e| trap_with(&format!("Decoding stable memory failed. Error: {e:?}")));
         println!("END   state::new_restored: ()");
         state
     }
 
     /// Saves the state to stable memory. Should be called in `pre_upgrade`.
     pub fn save(&self) {
-        self.save_heap_to_managed_memory();
+        println!("START state::save_heap: ()");
+        let bytes = self.encode();
+        with_partitions(|p| p.write_bytes_to_managed_memory(&bytes));
     }
 }