Skip to main content

miden_client_sqlite_store/account/
accounts.rs

1//! Account-related database operations.
2
3use std::collections::{BTreeMap, BTreeSet};
4use std::string::ToString;
5use std::vec::Vec;
6
7use miden_client::account::{
8    Account,
9    AccountCode,
10    AccountHeader,
11    AccountId,
12    AccountPatch,
13    AccountStorage,
14    Address,
15    PartialAccount,
16    PartialStorage,
17    PartialStorageMap,
18    StorageMapKey,
19    StorageSlotName,
20    StorageSlotType,
21};
22use miden_client::asset::{Asset, AssetVault, AssetWitness};
23use miden_client::store::{
24    AccountRecord,
25    AccountRecordData,
26    AccountStatus,
27    AccountStorageFilter,
28    AccountUpdate,
29    ClientAccountType,
30    StoreError,
31};
32use miden_client::utils::{Deserializable, Serializable};
33use miden_client::{AccountError, Felt, Word};
34use miden_protocol::account::{AccountStorageHeader, StorageMapWitness, StorageSlotHeader};
35use miden_protocol::asset::{AssetId, PartialVault};
36use miden_protocol::crypto::merkle::MerkleError;
37use rusqlite::{Connection, OptionalExtension, Transaction, named_params, params};
38
39use crate::account::rows::{
40    query_account_addresses,
41    query_account_code,
42    query_historical_account_headers,
43    query_latest_account_headers,
44    query_storage_slots,
45    query_storage_values,
46    query_vault_assets,
47};
48use crate::forest::{ScopedAccountForest, SqliteForestBackend, allocate_forest_revision};
49use crate::sql_error::SqlResultExt;
50use crate::{
51    SqliteStore,
52    blob_array,
53    column_value_as_u64,
54    insert_sql,
55    int_array,
56    subst,
57    u64_to_value,
58    with_write_tx,
59};
60
61impl SqliteStore {
62    // READER METHODS
63    // --------------------------------------------------------------------------------------------
64
65    pub(crate) fn get_account_ids(conn: &mut Connection) -> Result<Vec<AccountId>, StoreError> {
66        const QUERY: &str = "SELECT id FROM latest_account_headers";
67
68        conn.prepare_cached(QUERY)
69            .into_store_error()?
70            .query_map([], |row| row.get(0))
71            .expect("no binding parameters used in query")
72            .map(|result| {
73                let id: Vec<u8> = result.into_store_error()?;
74                Ok(AccountId::read_from_bytes(&id)?)
75            })
76            .collect::<Result<Vec<AccountId>, StoreError>>()
77    }
78
79    pub(crate) fn get_account_headers(
80        conn: &mut Connection,
81    ) -> Result<Vec<(AccountHeader, AccountStatus)>, StoreError> {
82        Ok(query_latest_account_headers(conn, "1=1 ORDER BY id", params![])?
83            .into_iter()
84            .map(|(header, status, _)| (header, status))
85            .collect())
86    }
87
88    pub(crate) fn get_account_header(
89        conn: &Connection,
90        account_id: AccountId,
91    ) -> Result<Option<(AccountHeader, AccountStatus)>, StoreError> {
92        Ok(query_latest_account_headers(conn, "id = ?", params![account_id.to_bytes()])?
93            .pop()
94            .map(|(header, status, _)| (header, status)))
95    }
96
97    pub(crate) fn get_account_header_by_commitment(
98        conn: &mut Connection,
99        account_commitment: Word,
100    ) -> Result<Option<AccountHeader>, StoreError> {
101        Ok(query_historical_account_headers(
102            conn,
103            "account_commitment = ?",
104            params![account_commitment.to_bytes()],
105        )?
106        .pop()
107        .map(|(header, _)| header))
108    }
109
110    /// Retrieves a complete account record with full vault and storage data.
111    pub(crate) fn get_account(
112        conn: &mut Connection,
113        account_id: AccountId,
114    ) -> Result<Option<AccountRecord>, StoreError> {
115        let Some((header, status, client_account_type)) =
116            query_latest_account_headers(conn, "id = ?", params![account_id.to_bytes()])?.pop()
117        else {
118            return Ok(None);
119        };
120
121        let assets = query_vault_assets(conn, account_id)?;
122        let vault = AssetVault::new(&assets)?;
123
124        let slots = query_storage_slots(conn, account_id, &AccountStorageFilter::All)?
125            .into_values()
126            .collect();
127
128        let storage = AccountStorage::new(slots)?;
129
130        let Some(account_code) = query_account_code(conn, header.code_commitment())? else {
131            return Ok(None);
132        };
133
134        let account = Account::new_unchecked(
135            header.id(),
136            vault,
137            storage,
138            account_code,
139            header.nonce(),
140            status.seed().copied(),
141        );
142
143        let account_data = AccountRecordData::Full(account);
144        Ok(Some(AccountRecord::new(account_data, status, client_account_type)))
145    }
146
147    /// Retrieves a minimal partial account record with storage and vault witnesses.
148    pub(crate) fn get_minimal_partial_account(
149        conn: &mut Connection,
150        account_id: AccountId,
151    ) -> Result<Option<AccountRecord>, StoreError> {
152        let Some((header, status, client_account_type)) =
153            query_latest_account_headers(conn, "id = ?", params![account_id.to_bytes()])?.pop()
154        else {
155            return Ok(None);
156        };
157
158        // Partial vault retrieval
159        let partial_vault = PartialVault::new(header.vault_root());
160
161        // Partial storage retrieval
162        let mut storage_header = Vec::new();
163        let mut maps = vec![];
164
165        let storage_values = query_storage_values(conn, account_id)?;
166
167        // Storage maps are always minimal here (just roots, no entries). New accounts that need
168        // full storage data are handled by the DataStore layer, which fetches the full account via
169        // `get_account()` when nonce == 0.
170        for (slot_name, (slot_type, value)) in storage_values {
171            storage_header.push(StorageSlotHeader::new(slot_name.clone(), slot_type, value));
172            if slot_type == StorageSlotType::Map {
173                maps.push(PartialStorageMap::new(value));
174            }
175        }
176        storage_header.sort_by_key(StorageSlotHeader::id);
177        let storage_header =
178            AccountStorageHeader::new(storage_header).map_err(StoreError::AccountError)?;
179        let partial_storage =
180            PartialStorage::new(storage_header, maps).map_err(StoreError::AccountError)?;
181
182        let Some(account_code) = query_account_code(conn, header.code_commitment())? else {
183            return Ok(None);
184        };
185
186        let partial_account = PartialAccount::new(
187            header.id(),
188            header.nonce(),
189            account_code,
190            partial_storage,
191            partial_vault,
192            status.seed().copied(),
193        )?;
194        let account_record_data = AccountRecordData::Partial(partial_account);
195        Ok(Some(AccountRecord::new(account_record_data, status, client_account_type)))
196    }
197
198    pub fn get_foreign_account_code(
199        conn: &mut Connection,
200        account_ids: Vec<AccountId>,
201    ) -> Result<BTreeMap<AccountId, AccountCode>, StoreError> {
202        let account_id_list = blob_array(account_ids);
203        const QUERY: &str = "
204            SELECT account_id, code
205            FROM foreign_account_code JOIN account_code ON foreign_account_code.code_commitment = account_code.commitment
206            WHERE account_id IN rarray(?)";
207
208        conn.prepare_cached(QUERY)
209            .into_store_error()?
210            .query_map([account_id_list], |row| Ok((row.get("account_id")?, row.get("code")?)))
211            .into_store_error()?
212            .map(|result| {
213                let (id, code): (Vec<u8>, Vec<u8>) = result.into_store_error()?;
214                Ok((AccountId::read_from_bytes(&id)?, AccountCode::read_from_bytes(&code)?))
215            })
216            .collect::<Result<BTreeMap<AccountId, AccountCode>, _>>()
217    }
218
219    /// Retrieves the full asset vault for a specific account.
220    pub fn get_account_vault(
221        conn: &Connection,
222        account_id: AccountId,
223    ) -> Result<AssetVault, StoreError> {
224        let assets = query_vault_assets(conn, account_id)?;
225        Ok(AssetVault::new(&assets)?)
226    }
227
228    /// Retrieves the full storage for a specific account.
229    pub fn get_account_storage(
230        conn: &Connection,
231        account_id: AccountId,
232        filter: &AccountStorageFilter,
233    ) -> Result<AccountStorage, StoreError> {
234        let slots = query_storage_slots(conn, account_id, filter)?.into_values().collect();
235        Ok(AccountStorage::new(slots)?)
236    }
237
238    /// Fetches a specific asset from the account's vault without the need of loading the entire
239    /// vault. The witness is retrieved from the [`AccountSmtForest`].
240    pub(crate) fn get_account_asset(
241        conn: &mut Connection,
242        account_id: AccountId,
243        asset_id: AssetId,
244    ) -> Result<Option<(Asset, AssetWitness)>, StoreError> {
245        // Begin the transaction first so the header and forest reads share one snapshot.
246        let db_tx = conn.transaction().into_store_error()?;
247        let header = Self::require_latest_account_header(&db_tx, account_id)?;
248        let smt_forest = ScopedAccountForest::new(SqliteForestBackend::new(&db_tx))?;
249
250        match smt_forest.get_asset_and_witness(account_id, header.vault_root(), asset_id) {
251            Ok((asset, witness)) => Ok(Some((asset, witness))),
252            Err(StoreError::VaultKeyNotTracked(..)) => Ok(None),
253            Err(err) => Err(err),
254        }
255    }
256
257    /// Retrieves a specific item from the account's storage map without loading the entire storage.
258    /// The witness is retrieved from the [`AccountSmtForest`].
259    pub(crate) fn get_account_map_item(
260        conn: &mut Connection,
261        account_id: AccountId,
262        slot_name: StorageSlotName,
263        key: StorageMapKey,
264    ) -> Result<(Word, StorageMapWitness), StoreError> {
265        // Begin the transaction first so the slot root and forest reads share one snapshot.
266        let db_tx = conn.transaction().into_store_error()?;
267        let header = Self::require_latest_account_header(&db_tx, account_id)?;
268
269        let mut storage_values = query_storage_values(&db_tx, account_id)?;
270        let (slot_type, map_root) = storage_values
271            .remove(&slot_name)
272            .ok_or(StoreError::AccountStorageRootNotFound(header.storage_commitment()))?;
273        if slot_type != StorageSlotType::Map {
274            return Err(StoreError::AccountError(AccountError::StorageSlotNotMap(slot_name)));
275        }
276
277        let smt_forest = ScopedAccountForest::new(SqliteForestBackend::new(&db_tx))?;
278
279        let witness =
280            smt_forest.get_storage_map_item_witness(account_id, &slot_name, map_root, key)?;
281        let item = witness.get(key).unwrap_or(miden_client::EMPTY_WORD);
282
283        Ok((item, witness))
284    }
285
286    /// Retrieves vault asset witnesses for the given vault keys, including emptiness proofs for
287    /// keys absent from the vault (which the executor needs when an asset is being added).
288    ///
289    /// The witnesses are opened against the account's vault tree in the forest, after verifying
290    /// that its root matches `vault_root` — the committed root the caller expects.
291    pub(crate) fn get_vault_asset_witnesses(
292        conn: &mut Connection,
293        account_id: AccountId,
294        vault_root: Word,
295        asset_ids: BTreeSet<AssetId>,
296    ) -> Result<Vec<AssetWitness>, StoreError> {
297        let db_tx = conn.transaction().into_store_error()?;
298        let smt_forest = ScopedAccountForest::new(SqliteForestBackend::new(&db_tx))?;
299        smt_forest.open_vault_asset_witnesses(account_id, vault_root, asset_ids)
300    }
301
302    pub(crate) fn get_account_addresses(
303        conn: &mut Connection,
304        account_id: AccountId,
305    ) -> Result<Vec<Address>, StoreError> {
306        query_account_addresses(conn, account_id)
307    }
308
309    /// Retrieves the account code for a specific account by ID.
310    pub(crate) fn get_account_code_by_id(
311        conn: &mut Connection,
312        account_id: AccountId,
313    ) -> Result<Option<AccountCode>, StoreError> {
314        let Some((header, ..)) =
315            query_latest_account_headers(conn, "id = ?", params![account_id.to_bytes()])?
316                .into_iter()
317                .next()
318        else {
319            return Ok(None);
320        };
321
322        query_account_code(conn, header.code_commitment())
323    }
324
325    // MUTATOR/WRITER METHODS
326    // --------------------------------------------------------------------------------------------
327
328    pub(crate) fn insert_account(
329        conn: &mut Connection,
330        account: &Account,
331        initial_address: &Address,
332        client_account_type: ClientAccountType,
333    ) -> Result<(), StoreError> {
334        with_write_tx(conn, |tx| {
335            let mut smt_forest = ScopedAccountForest::new(SqliteForestBackend::new(tx))?;
336            Self::insert_account_code(tx, account.code())?;
337
338            let account_id = account.id();
339            Self::insert_storage_slots(tx, account_id, account.storage().slots().iter())?;
340            Self::insert_assets(tx, account_id, account.vault().assets())?;
341            let watched = matches!(client_account_type, ClientAccountType::Watched);
342            Self::insert_new_account_header(tx, &account.into(), account.seed(), watched)?;
343            Self::insert_address_tx(tx, initial_address, account.id())?;
344
345            Self::reconcile_account_forest(
346                tx,
347                &mut smt_forest,
348                account_id,
349                account.vault(),
350                account.storage(),
351            )
352        })
353    }
354
355    pub(crate) fn update_account(
356        conn: &mut Connection,
357        new_account_state: &Account,
358    ) -> Result<(), StoreError> {
359        with_write_tx(conn, |tx| {
360            let mut smt_forest = ScopedAccountForest::new(SqliteForestBackend::new(tx))?;
361            Self::update_account_state(tx, &mut smt_forest, new_account_state)
362        })
363    }
364
365    pub(crate) fn upsert_foreign_account_code(
366        conn: &mut Connection,
367        account_id: AccountId,
368        code: &AccountCode,
369    ) -> Result<(), StoreError> {
370        with_write_tx(conn, |tx| {
371            Self::insert_account_code(tx, code)?;
372
373            const QUERY: &str =
374                insert_sql!(foreign_account_code { account_id, code_commitment } | REPLACE);
375
376            tx.execute(QUERY, params![account_id.to_bytes(), code.commitment().to_bytes()])
377                .into_store_error()?;
378
379            Ok(())
380        })
381    }
382
383    pub(crate) fn insert_address(
384        conn: &mut Connection,
385        address: &Address,
386        account_id: AccountId,
387    ) -> Result<(), StoreError> {
388        with_write_tx(conn, |tx| Self::insert_address_tx(tx, address, account_id))
389    }
390
391    pub(crate) fn insert_address_tx(
392        tx: &Transaction<'_>,
393        address: &Address,
394        account_id: AccountId,
395    ) -> Result<(), StoreError> {
396        const QUERY: &str = insert_sql!(addresses { address, account_id } | REPLACE);
397        let serialized_address = address.to_bytes();
398        tx.execute(QUERY, params![serialized_address, account_id.to_bytes(),])
399            .into_store_error()?;
400
401        Ok(())
402    }
403
404    /// Returns `true` if a row was deleted, `false` if the address wasn't tracked.
405    pub(crate) fn remove_address(
406        conn: &mut Connection,
407        address: &Address,
408    ) -> Result<bool, StoreError> {
409        with_write_tx(conn, |tx| {
410            const DELETE_QUERY: &str = "DELETE FROM addresses WHERE address = ?";
411            let count = tx.execute(DELETE_QUERY, params![address.to_bytes()]).into_store_error()?;
412
413            Ok(count > 0)
414        })
415    }
416
417    /// Inserts an [`AccountCode`].
418    pub(crate) fn insert_account_code(
419        tx: &Transaction<'_>,
420        account_code: &AccountCode,
421    ) -> Result<(), StoreError> {
422        const QUERY: &str = insert_sql!(account_code { commitment, code } | IGNORE);
423        tx.execute(QUERY, params![account_code.commitment().to_bytes(), account_code.to_bytes()])
424            .into_store_error()?;
425        Ok(())
426    }
427
428    /// Applies the account patch to the account state, updating the vault and storage maps.
429    ///
430    /// Archives old values from latest to historical and updates latest via INSERT OR REPLACE.
431    pub(crate) fn apply_account_patch(
432        tx: &Transaction<'_>,
433        smt_forest: &mut ScopedAccountForest<'_, '_>,
434        init_account_state: &AccountHeader,
435        final_account_state: &AccountHeader,
436        patch: &AccountPatch,
437    ) -> Result<(), StoreError> {
438        let account_id = final_account_state.id();
439
440        // Reject patches for accounts the store does not track (forest updates for unknown accounts
441        // would silently create partial state from empty trees), and stale or replayed patches
442        // whose initial state does not match the stored latest state (they would overwrite newer
443        // state and archive incorrect history).
444        let stored_header = Self::require_latest_account_header(tx, account_id)?;
445        if stored_header.to_commitment() != init_account_state.to_commitment() {
446            return Err(StoreError::DatabaseError(format!(
447                "apply_account_patch: stored state {} for account {} does not match the patch's \
448                 initial state {}",
449                stored_header.to_commitment(),
450                account_id,
451                init_account_state.to_commitment(),
452            )));
453        }
454
455        // Archive old header and insert the new one
456        Self::replace_account_header(tx, final_account_state, init_account_state, None)?;
457
458        Self::apply_account_vault_patch(tx, account_id, final_account_state, patch.vault())?;
459
460        // Build one forest update covering the vault and every changed map slot, and apply it at a
461        // freshly allocated revision.
462        let mut update = AccountUpdate::new();
463        update.vault_patch(account_id, patch.vault(), final_account_state.vault_root());
464        update.storage_patch(account_id, patch.storage());
465
466        let revision = allocate_forest_revision(tx).into_store_error()?;
467        smt_forest.apply(revision, update)?;
468
469        Self::write_storage_patch(
470            tx,
471            smt_forest,
472            account_id,
473            final_account_state.nonce().as_canonical_u64(),
474            patch.storage(),
475        )?;
476        Self::verify_storage_commitment(tx, account_id, final_account_state.storage_commitment())?;
477
478        Ok(())
479    }
480
481    /// Reconciles the account's forest lineages to exactly match the provided full state.
482    ///
483    /// Map slots that disappeared from the state are enumerated from the latest storage tables, so
484    /// this must run before those rows are replaced.
485    pub(crate) fn reconcile_account_forest(
486        tx: &Transaction<'_>,
487        smt_forest: &mut ScopedAccountForest<'_, '_>,
488        account_id: AccountId,
489        vault: &AssetVault,
490        storage: &AccountStorage,
491    ) -> Result<(), StoreError> {
492        let mut update = AccountUpdate::new();
493        update.full_state(account_id, vault.assets(), storage.slots().iter());
494
495        // Slots that still have stored rows but are absent from the new state must be emptied
496        // rather than keeping their old entries. Slots the new state does repopulate keep the
497        // entries recorded above; this only marks the lineage exhaustive.
498        for slot_name in Self::query_map_slot_names(tx, account_id)? {
499            update.clear_map(account_id, &slot_name);
500        }
501
502        Self::apply_forest_update(tx, smt_forest, update)
503    }
504
505    /// Reconciles the account's forest lineages to the state currently stored in the latest account
506    /// tables. Used after rows are restored from historical during undo.
507    ///
508    /// `extra_map_slots` lists map slots that may hold forest entries even though they have no rows
509    /// anymore (captured before the tables were rewritten); their lineages are reset to the empty
510    /// tree unless the restored state repopulates them.
511    fn reconcile_account_forest_from_tables(
512        tx: &Transaction<'_>,
513        smt_forest: &mut ScopedAccountForest<'_, '_>,
514        account_id: AccountId,
515        extra_map_slots: &[StorageSlotName],
516    ) -> Result<(), StoreError> {
517        let assets = query_vault_assets(tx, account_id)?;
518        let slots = query_storage_slots(tx, account_id, &AccountStorageFilter::All)?;
519
520        let mut update = AccountUpdate::new();
521        update.full_state(account_id, assets.into_iter(), slots.values());
522        for slot_name in extra_map_slots {
523            update.clear_map(account_id, slot_name);
524        }
525
526        Self::apply_forest_update(tx, smt_forest, update)
527    }
528
529    /// Verifies that the persisted top-level storage slots match the expected commitment.
530    ///
531    /// This runs after the storage patch is written so create, update, and removal semantics have a
532    /// single source of truth. A mismatch rolls back together with the rest of the transaction.
533    fn verify_storage_commitment(
534        tx: &Transaction<'_>,
535        account_id: AccountId,
536        expected: Word,
537    ) -> Result<(), StoreError> {
538        let mut slot_headers: Vec<StorageSlotHeader> = query_storage_values(tx, account_id)?
539            .into_iter()
540            .map(|(slot_name, (slot_type, value))| {
541                StorageSlotHeader::new(slot_name, slot_type, value)
542            })
543            .collect();
544        slot_headers.sort_by_key(StorageSlotHeader::id);
545
546        let actual = AccountStorageHeader::new(slot_headers)
547            .map_err(StoreError::AccountError)?
548            .to_commitment();
549        if actual != expected {
550            return Err(StoreError::MerkleStoreError(MerkleError::ConflictingRoots {
551                expected_root: expected,
552                actual_root: actual,
553            }));
554        }
555
556        Ok(())
557    }
558
559    /// Applies a recorded forest update at a freshly allocated revision.
560    fn apply_forest_update(
561        tx: &Transaction<'_>,
562        smt_forest: &mut ScopedAccountForest<'_, '_>,
563        update: AccountUpdate,
564    ) -> Result<(), StoreError> {
565        let revision = allocate_forest_revision(tx).into_store_error()?;
566        smt_forest.apply(revision, update)
567    }
568
569    /// Returns the stored latest header of an account, or [`StoreError::AccountDataNotFound`] if
570    /// the store does not track it.
571    fn require_latest_account_header(
572        tx: &Transaction<'_>,
573        account_id: AccountId,
574    ) -> Result<AccountHeader, StoreError> {
575        query_latest_account_headers(tx, "id = ?", params![account_id.to_bytes()])?
576            .into_iter()
577            .next()
578            .map(|(header, ..)| header)
579            .ok_or(StoreError::AccountDataNotFound(account_id))
580    }
581
582    /// Returns the names of the map slots that currently have entries stored for an account.
583    fn query_map_slot_names(
584        tx: &Transaction<'_>,
585        account_id: AccountId,
586    ) -> Result<Vec<StorageSlotName>, StoreError> {
587        let mut stmt = tx
588            .prepare(
589                "SELECT DISTINCT slot_name FROM latest_storage_map_entries WHERE account_id = ?",
590            )
591            .into_store_error()?;
592        let rows = stmt
593            .query_map(params![account_id.to_bytes()], |row| row.get::<_, String>(0))
594            .into_store_error()?;
595
596        rows.map(|row| {
597            StorageSlotName::new(row.into_store_error()?)
598                .map_err(|e| StoreError::ParsingError(e.to_string()))
599        })
600        .collect()
601    }
602
603    /// Undoes discarded account states by restoring old values from historical.
604    pub(crate) fn undo_account_state(
605        tx: &Transaction<'_>,
606        smt_forest: &mut ScopedAccountForest<'_, '_>,
607        discarded_states: &[(AccountId, Word)],
608    ) -> Result<(), StoreError> {
609        if discarded_states.is_empty() {
610            return Ok(());
611        }
612
613        let commitment_params =
614            blob_array(discarded_states.iter().map(|(_, commitment)| commitment));
615
616        // Resolve (account_id, nonce) pairs from both latest and historical headers, and group the
617        // nonces by account. The most recent discarded state is in latest, older ones are in
618        // historical.
619        let mut nonces_by_account: BTreeMap<Vec<u8>, BTreeSet<u64>> = BTreeMap::new();
620        for query in [
621            "SELECT id, nonce FROM latest_account_headers WHERE account_commitment IN rarray(?)",
622            "SELECT id, nonce FROM historical_account_headers WHERE account_commitment IN rarray(?)",
623        ] {
624            for row in tx
625                .prepare(query)
626                .into_store_error()?
627                .query_map(params![commitment_params.clone()], |row| {
628                    let id: Vec<u8> = row.get("id")?;
629                    let nonce: u64 = column_value_as_u64(row, "nonce")?;
630                    Ok((id, nonce))
631                })
632                .into_store_error()?
633            {
634                let (id, nonce) = row.into_store_error()?;
635                nonces_by_account.entry(id).or_default().insert(nonce);
636            }
637        }
638
639        // Undo one account at a time. Read the account's map slot names before the undo rewrites
640        // the latest tables, then reconcile its forest lineages to the restored state.
641        for (account_id_bytes, nonces) in &nonces_by_account {
642            let account_id = AccountId::read_from_bytes(account_id_bytes)?;
643            let stale_map_slots = Self::query_map_slot_names(tx, account_id)?;
644            Self::undo_account_nonces(tx, account_id_bytes, nonces)?;
645            Self::reconcile_account_forest_from_tables(
646                tx,
647                smt_forest,
648                account_id,
649                &stale_map_slots,
650            )?;
651        }
652
653        Ok(())
654    }
655
656    /// Undoes all nonces for a single account: restores old values, restores old header, and cleans
657    /// up consumed historical entries.
658    fn undo_account_nonces(
659        tx: &Transaction<'_>,
660        account_id_bytes: &[u8],
661        nonces: &BTreeSet<u64>,
662    ) -> Result<(), StoreError> {
663        // Undo each nonce in descending order. Each nonce's old value is the state before that
664        // nonce, so the most recent nonce must be undone first. Earlier nonces then overwrite it
665        // with the correct final value.
666        for &nonce in nonces.iter().rev() {
667            let nonce_val = u64_to_value(nonce);
668            Self::restore_old_values_for_nonce(tx, account_id_bytes, &nonce_val)?;
669        }
670
671        // Restore the old header from the earliest discarded nonce. The set always holds at least
672        // one nonce, because an entry is added to the map only when a nonce is inserted.
673        let min_nonce_val = u64_to_value(*nonces.first().expect("nonces is not empty"));
674
675        let old_header_exists: bool = tx
676            .query_row(
677                "SELECT COUNT(*) FROM historical_account_headers \
678                 WHERE id = ? AND replaced_at_nonce = ?",
679                params![account_id_bytes, &min_nonce_val],
680                |row| row.get::<_, i64>(0),
681            )
682            .into_store_error()?
683            > 0;
684
685        if old_header_exists {
686            // `watched` is not carried in historical_account_headers, so this restore resets it to
687            // the column default (FALSE). This is safe because undo only fires for discarded local
688            // transactions, and watched accounts have none.
689            tx.execute(
690                "INSERT OR REPLACE INTO latest_account_headers \
691                 (id, account_commitment, code_commitment, storage_commitment, \
692                  vault_root, nonce, account_seed, locked) \
693                 SELECT id, account_commitment, code_commitment, storage_commitment, \
694                        vault_root, nonce, account_seed, locked \
695                 FROM historical_account_headers \
696                 WHERE id = ? AND replaced_at_nonce = ?",
697                params![account_id_bytes, &min_nonce_val],
698            )
699            .into_store_error()?;
700        } else {
701            // No previous state — delete the account entirely
702            for table in [
703                "DELETE FROM latest_account_headers WHERE id = ?",
704                "DELETE FROM latest_account_storage WHERE account_id = ?",
705                "DELETE FROM latest_storage_map_entries WHERE account_id = ?",
706                "DELETE FROM latest_account_assets WHERE account_id = ?",
707            ] {
708                tx.execute(table, params![account_id_bytes]).into_store_error()?;
709            }
710        }
711
712        // Delete all consumed historical entries at the discarded nonces.
713        let nonce_params = int_array(nonces.iter().copied());
714        for table in [
715            "historical_account_storage",
716            "historical_storage_map_entries",
717            "historical_account_assets",
718        ] {
719            tx.execute(
720                &format!(
721                    "DELETE FROM {table} WHERE account_id = ? AND replaced_at_nonce IN rarray(?)"
722                ),
723                params![account_id_bytes, nonce_params.clone()],
724            )
725            .into_store_error()?;
726        }
727        tx.execute(
728            "DELETE FROM historical_account_headers \
729             WHERE id = ? AND replaced_at_nonce IN rarray(?)",
730            params![account_id_bytes, nonce_params],
731        )
732        .into_store_error()?;
733
734        Ok(())
735    }
736
737    /// Restores old values from historical entries for a given nonce. Non-NULL old values overwrite
738    /// latest, NULL old values (new entries) are deleted.
739    fn restore_old_values_for_nonce(
740        tx: &Transaction<'_>,
741        account_id_bytes: &[u8],
742        nonce_val: &rusqlite::types::Value,
743    ) -> Result<(), StoreError> {
744        // Restore storage slots with non-NULL old values
745        tx.execute(
746            "INSERT OR REPLACE INTO latest_account_storage \
747             (account_id, slot_name, slot_value, slot_type) \
748             SELECT account_id, slot_name, old_slot_value, slot_type \
749             FROM historical_account_storage \
750             WHERE account_id = ? AND replaced_at_nonce = ? AND old_slot_value IS NOT NULL",
751            params![account_id_bytes, nonce_val],
752        )
753        .into_store_error()?;
754
755        // Delete storage slots that were new (NULL old value)
756        tx.execute(
757            "DELETE FROM latest_account_storage \
758             WHERE account_id = ?1 AND slot_name IN (\
759                 SELECT slot_name FROM historical_account_storage \
760                 WHERE account_id = ?1 AND replaced_at_nonce = ?2 AND old_slot_value IS NULL\
761             )",
762            params![account_id_bytes, nonce_val],
763        )
764        .into_store_error()?;
765
766        // Restore map entries with non-NULL old values
767        tx.execute(
768            "INSERT OR REPLACE INTO latest_storage_map_entries \
769             (account_id, slot_name, key, value) \
770             SELECT account_id, slot_name, key, old_value \
771             FROM historical_storage_map_entries \
772             WHERE account_id = ? AND replaced_at_nonce = ? AND old_value IS NOT NULL",
773            params![account_id_bytes, nonce_val],
774        )
775        .into_store_error()?;
776
777        // Delete map entries that were new (NULL old value)
778        tx.execute(
779            "DELETE FROM latest_storage_map_entries \
780             WHERE account_id = ?1 AND EXISTS (\
781                 SELECT 1 FROM historical_storage_map_entries h \
782                 WHERE h.account_id = latest_storage_map_entries.account_id \
783                   AND h.slot_name = latest_storage_map_entries.slot_name \
784                   AND h.key = latest_storage_map_entries.key \
785                   AND h.replaced_at_nonce = ?2 AND h.old_value IS NULL\
786             )",
787            params![account_id_bytes, nonce_val],
788        )
789        .into_store_error()?;
790
791        // Restore assets with non-NULL old values
792        tx.execute(
793            "INSERT OR REPLACE INTO latest_account_assets \
794             (account_id, asset_id, asset) \
795             SELECT account_id, asset_id, old_asset \
796             FROM historical_account_assets \
797             WHERE account_id = ? AND replaced_at_nonce = ? AND old_asset IS NOT NULL",
798            params![account_id_bytes, nonce_val],
799        )
800        .into_store_error()?;
801
802        // Delete assets that were new (NULL old value)
803        tx.execute(
804            "DELETE FROM latest_account_assets \
805             WHERE account_id = ?1 AND asset_id IN (\
806             SELECT asset_id FROM historical_account_assets \
807                 WHERE account_id = ?1 AND replaced_at_nonce = ?2 AND old_asset IS NULL\
808             )",
809            params![account_id_bytes, nonce_val],
810        )
811        .into_store_error()?;
812
813        Ok(())
814    }
815
816    /// Replaces the account state with a completely new one from the network.
817    ///
818    /// Replaces the account state entirely: archives old state to historical, clears latest,
819    /// inserts new state to latest only. Preserves the `watched` flag.
820    pub(crate) fn update_account_state(
821        tx: &Transaction<'_>,
822        smt_forest: &mut ScopedAccountForest<'_, '_>,
823        new_account_state: &Account,
824    ) -> Result<(), StoreError> {
825        let account_id = new_account_state.id();
826        let account_id_bytes = account_id.to_bytes();
827
828        // Read old header before mutating the SMT snapshot or database rows. Sync filters stale
829        // full-account snapshots; if one still reaches storage, reject it before mutating.
830        let old_header = Self::require_latest_account_header(tx, account_id)?;
831
832        if new_account_state.nonce().as_canonical_u64() < old_header.nonce().as_canonical_u64() {
833            return Err(StoreError::DatabaseError(format!(
834                "update_account_state: new nonce {} is less than old nonce {} for account {}",
835                new_account_state.nonce().as_canonical_u64(),
836                old_header.nonce().as_canonical_u64(),
837                account_id,
838            )));
839        }
840
841        let nonce_val = u64_to_value(new_account_state.nonce().as_canonical_u64());
842
843        // Reconcile the forest to the new full state before the latest tables are replaced below.
844        Self::reconcile_account_forest(
845            tx,
846            smt_forest,
847            account_id,
848            new_account_state.vault(),
849            new_account_state.storage(),
850        )?;
851
852        // Archive all old entries from latest → historical
853        tx.execute(
854            "INSERT OR REPLACE INTO historical_account_storage \
855             (account_id, replaced_at_nonce, slot_name, old_slot_value, slot_type) \
856             SELECT account_id, ?, slot_name, slot_value, slot_type \
857             FROM latest_account_storage WHERE account_id = ?",
858            params![&nonce_val, &account_id_bytes],
859        )
860        .into_store_error()?;
861        tx.execute(
862            "INSERT OR REPLACE INTO historical_storage_map_entries \
863             (account_id, replaced_at_nonce, slot_name, key, old_value) \
864             SELECT account_id, ?, slot_name, key, value \
865             FROM latest_storage_map_entries WHERE account_id = ?",
866            params![&nonce_val, &account_id_bytes],
867        )
868        .into_store_error()?;
869        tx.execute(
870            "INSERT OR REPLACE INTO historical_account_assets \
871             (account_id, replaced_at_nonce, asset_id, old_asset) \
872             SELECT account_id, ?, asset_id, asset \
873             FROM latest_account_assets WHERE account_id = ?",
874            params![&nonce_val, &account_id_bytes],
875        )
876        .into_store_error()?;
877
878        // Delete all latest entries for this account
879        tx.execute(
880            "DELETE FROM latest_account_storage WHERE account_id = ?",
881            params![&account_id_bytes],
882        )
883        .into_store_error()?;
884        tx.execute(
885            "DELETE FROM latest_storage_map_entries WHERE account_id = ?",
886            params![&account_id_bytes],
887        )
888        .into_store_error()?;
889        tx.execute(
890            "DELETE FROM latest_account_assets WHERE account_id = ?",
891            params![&account_id_bytes],
892        )
893        .into_store_error()?;
894
895        // Insert all new entries into latest only
896        Self::insert_storage_slots(tx, account_id, new_account_state.storage().slots().iter())?;
897        Self::insert_assets(tx, account_id, new_account_state.vault().assets())?;
898
899        // Write NULL historical entries for genuinely new entries that didn't exist in the old
900        // state (INSERT OR IGNORE skips entries already archived above)
901        tx.execute(
902            "INSERT OR IGNORE INTO historical_account_storage \
903             (account_id, replaced_at_nonce, slot_name, old_slot_value, slot_type) \
904             SELECT account_id, ?, slot_name, NULL, slot_type \
905             FROM latest_account_storage WHERE account_id = ?",
906            params![&nonce_val, &account_id_bytes],
907        )
908        .into_store_error()?;
909        tx.execute(
910            "INSERT OR IGNORE INTO historical_storage_map_entries \
911             (account_id, replaced_at_nonce, slot_name, key, old_value) \
912             SELECT account_id, ?, slot_name, key, NULL \
913             FROM latest_storage_map_entries WHERE account_id = ?",
914            params![&nonce_val, &account_id_bytes],
915        )
916        .into_store_error()?;
917        tx.execute(
918            "INSERT OR IGNORE INTO historical_account_assets \
919             (account_id, replaced_at_nonce, asset_id, old_asset) \
920             SELECT account_id, ?, asset_id, NULL \
921             FROM latest_account_assets WHERE account_id = ?",
922            params![&nonce_val, &account_id_bytes],
923        )
924        .into_store_error()?;
925
926        // Archive the old header to historical and write the new one to latest. A state that is
927        // still undeployed keeps its seed
928        let new_seed = new_account_state.seed().filter(|_| new_account_state.is_new());
929        Self::replace_account_header(tx, &new_account_state.into(), &old_header, new_seed)?;
930
931        Ok(())
932    }
933
934    /// Applies an incremental patch to a public account's state during sync.
935    pub(crate) fn apply_sync_account_patch(
936        tx: &Transaction<'_>,
937        smt_forest: &mut ScopedAccountForest<'_, '_>,
938        new_header: &AccountHeader,
939        patch: &AccountPatch,
940    ) -> Result<(), StoreError> {
941        let account_id = new_header.id();
942
943        // Read current header from the store.
944        let init_header = Self::require_latest_account_header(tx, account_id)?;
945
946        if new_header.nonce().as_canonical_u64() <= init_header.nonce().as_canonical_u64() {
947            return Err(StoreError::DatabaseError(format!(
948                "apply_sync_account_patch: new nonce {} is not greater than local nonce {} for account {}",
949                new_header.nonce().as_canonical_u64(),
950                init_header.nonce().as_canonical_u64(),
951                account_id,
952            )));
953        }
954
955        // Transaction derefs to Connection, so we can pass it where Connection is expected.
956
957        Self::apply_account_patch(tx, smt_forest, &init_header, new_header, patch)
958    }
959
960    /// Locks the account if the mismatched digest doesn't belong to a previous account state (stale
961    /// data).
962    pub(crate) fn lock_account_on_unexpected_commitment(
963        tx: &Transaction<'_>,
964        account_id: &AccountId,
965        mismatched_digest: &Word,
966    ) -> Result<(), StoreError> {
967        // Mismatched digests may be due to stale network data. If the mismatched digest is tracked
968        // in the db and corresponds to the mismatched account, it means we got a past update and
969        // shouldn't lock the account.
970        const LOCK_CONDITION: &str = "WHERE id = :account_id AND NOT EXISTS (SELECT 1 FROM historical_account_headers WHERE id = :account_id AND account_commitment = :digest)";
971        let account_id_bytes = account_id.to_bytes();
972        let digest_bytes = mismatched_digest.to_bytes();
973        let params = named_params! {
974            ":account_id": account_id_bytes,
975            ":digest": digest_bytes
976        };
977
978        let query = format!("UPDATE latest_account_headers SET locked = true {LOCK_CONDITION}");
979        tx.execute(&query, params).into_store_error()?;
980
981        // Also lock historical rows so that undo_account_state preserves the lock.
982        let query = format!("UPDATE historical_account_headers SET locked = true {LOCK_CONDITION}");
983        tx.execute(&query, params).into_store_error()?;
984
985        Ok(())
986    }
987
988    // HELPERS
989    // --------------------------------------------------------------------------------------------
990
991    /// Writes a new row into `latest_account_headers`.
992    ///
993    /// Does not archive any previous state, use [`Self::replace_account_header`] when a row for
994    /// this account already exists. If a row does exist it will be overwritten with the provided
995    /// `watched` value and no historical row added.
996    fn insert_new_account_header(
997        tx: &Transaction<'_>,
998        new_header: &AccountHeader,
999        account_seed: Option<Word>,
1000        watched: bool,
1001    ) -> Result<(), StoreError> {
1002        let id = new_header.id().to_bytes();
1003        let code_commitment = new_header.code_commitment().to_bytes();
1004        let storage_commitment = new_header.storage_commitment().to_bytes();
1005        let vault_root = new_header.vault_root().to_bytes();
1006        let nonce = u64_to_value(new_header.nonce().as_canonical_u64());
1007        let commitment = new_header.to_commitment().to_bytes();
1008        let account_seed = account_seed.map(|seed| seed.to_bytes());
1009
1010        const LATEST_QUERY: &str = insert_sql!(
1011            latest_account_headers {
1012                id,
1013                code_commitment,
1014                storage_commitment,
1015                vault_root,
1016                nonce,
1017                account_seed,
1018                account_commitment,
1019                locked,
1020                watched
1021            } | REPLACE
1022        );
1023
1024        tx.execute(
1025            LATEST_QUERY,
1026            params![
1027                id,
1028                code_commitment,
1029                storage_commitment,
1030                vault_root,
1031                nonce,
1032                account_seed,
1033                commitment,
1034                false,
1035                watched,
1036            ],
1037        )
1038        .into_store_error()?;
1039
1040        Ok(())
1041    }
1042
1043    /// Replaces an account's latest header, archiving the previous one to historical.
1044    ///
1045    /// Preserves the `watched` flag from the existing latest row (mode is a per-account property,
1046    /// not per-state). The new latest row is written with `account_seed = new_seed` and `locked =
1047    /// false`; the previous seed and lock state move into the historical row. `new_seed` is only
1048    /// `Some` while the new state is still undeployed (nonce zero), since a deployed account no
1049    /// longer needs its seed.
1050    fn replace_account_header(
1051        tx: &Transaction<'_>,
1052        new_header: &AccountHeader,
1053        old_header: &AccountHeader,
1054        new_seed: Option<Word>,
1055    ) -> Result<(), StoreError> {
1056        if new_header.id() != old_header.id() {
1057            return Err(StoreError::DatabaseError(format!(
1058                "replace_account_header: account id mismatch (new: {}, old: {})",
1059                new_header.id(),
1060                old_header.id(),
1061            )));
1062        }
1063        if new_header.nonce().as_canonical_u64() < old_header.nonce().as_canonical_u64() {
1064            return Err(StoreError::DatabaseError(format!(
1065                "replace_account_header: new nonce {} is less than old nonce {} for account {}",
1066                new_header.nonce().as_canonical_u64(),
1067                old_header.nonce().as_canonical_u64(),
1068                new_header.id(),
1069            )));
1070        }
1071
1072        let id_bytes = new_header.id().to_bytes();
1073
1074        // `AccountHeader` doesn't carry the seed or per-account flags, so read them from the row
1075        // we're about to overwrite: `account_seed`/`locked` get archived into the historical row,
1076        // `watched` is carried into the new latest row.
1077        let (old_seed, old_locked, old_watched): (Option<Vec<u8>>, bool, bool) = tx
1078            .query_row(
1079                "SELECT account_seed, locked, watched FROM latest_account_headers WHERE id = ?",
1080                params![&id_bytes],
1081                |row| Ok((row.get("account_seed")?, row.get("locked")?, row.get("watched")?)),
1082            )
1083            .optional()
1084            .into_store_error()?
1085            .unwrap_or((None, false, false));
1086
1087        // Archive the old header to historical.
1088        let old_id = old_header.id().to_bytes();
1089        let old_code_commitment = old_header.code_commitment().to_bytes();
1090        let old_storage_commitment = old_header.storage_commitment().to_bytes();
1091        let old_vault_root = old_header.vault_root().to_bytes();
1092        let old_nonce = u64_to_value(old_header.nonce().as_canonical_u64());
1093        let old_commitment = old_header.to_commitment().to_bytes();
1094        let replaced_at_nonce = u64_to_value(new_header.nonce().as_canonical_u64());
1095
1096        const HISTORICAL_QUERY: &str = insert_sql!(
1097            historical_account_headers {
1098                id,
1099                code_commitment,
1100                storage_commitment,
1101                vault_root,
1102                nonce,
1103                account_seed,
1104                account_commitment,
1105                locked,
1106                replaced_at_nonce
1107            } | REPLACE
1108        );
1109
1110        tx.execute(
1111            HISTORICAL_QUERY,
1112            params![
1113                old_id,
1114                old_code_commitment,
1115                old_storage_commitment,
1116                old_vault_root,
1117                old_nonce,
1118                old_seed,
1119                old_commitment,
1120                old_locked,
1121                replaced_at_nonce,
1122            ],
1123        )
1124        .into_store_error()?;
1125
1126        // Write the new latest row.
1127        Self::insert_new_account_header(tx, new_header, new_seed, old_watched)
1128    }
1129
1130    /// Prunes historical account states for a single account up to the given nonce.
1131    ///
1132    /// Deletes all historical entries with `replaced_at_nonce <= up_to_nonce` (see DESIGN.md for
1133    /// why this threshold is safe), then removes any account code that was only referenced by the
1134    /// deleted headers.
1135    pub(crate) fn prune_account_history(
1136        conn: &mut Connection,
1137        account_id: AccountId,
1138        up_to_nonce: Felt,
1139    ) -> Result<usize, StoreError> {
1140        with_write_tx(conn, |tx| {
1141            let account_id_bytes = account_id.to_bytes();
1142            let boundary_val = u64_to_value(up_to_nonce.as_canonical_u64());
1143            let mut total_deleted: usize = 0;
1144
1145            // Collect code commitments from headers we are about to delete.
1146            let candidate_code_commitments: Vec<Vec<u8>> = {
1147                let mut stmt = tx
1148                    .prepare(
1149                        "SELECT DISTINCT code_commitment FROM historical_account_headers \
1150                     WHERE id = ? AND replaced_at_nonce <= ?",
1151                    )
1152                    .into_store_error()?;
1153                let rows = stmt
1154                    .query_map(params![&account_id_bytes, &boundary_val], |row| row.get(0))
1155                    .into_store_error()?;
1156                rows.collect::<Result<Vec<Vec<u8>>, _>>().into_store_error()?
1157            };
1158
1159            // Delete historical entries. The headers table names the account column `id`.
1160            for (table, account_column) in [
1161                ("historical_account_headers", "id"),
1162                ("historical_account_storage", "account_id"),
1163                ("historical_storage_map_entries", "account_id"),
1164                ("historical_account_assets", "account_id"),
1165            ] {
1166                let query = format!(
1167                    "DELETE FROM {table} WHERE {account_column} = ? AND replaced_at_nonce <= ?"
1168                );
1169                total_deleted += tx
1170                    .execute(&query, params![&account_id_bytes, &boundary_val])
1171                    .into_store_error()?;
1172            }
1173
1174            // Delete orphaned code: only check commitments from the deleted headers, and only if
1175            // they are not referenced by any remaining header or foreign code.
1176            for commitment in &candidate_code_commitments {
1177                let still_referenced: bool = tx
1178                    .query_row(
1179                        "SELECT EXISTS(
1180                        SELECT 1 FROM latest_account_headers WHERE code_commitment = ?1
1181                        UNION ALL
1182                        SELECT 1 FROM historical_account_headers WHERE code_commitment = ?1
1183                        UNION ALL
1184                        SELECT 1 FROM foreign_account_code WHERE code_commitment = ?1
1185                    )",
1186                        params![commitment],
1187                        |row| row.get(0),
1188                    )
1189                    .into_store_error()?;
1190
1191                if !still_referenced {
1192                    total_deleted += tx
1193                        .execute(
1194                            "DELETE FROM account_code WHERE commitment = ?",
1195                            params![commitment],
1196                        )
1197                        .into_store_error()?;
1198                }
1199            }
1200
1201            Ok(total_deleted)
1202        })
1203    }
1204}