Skip to main content

miden_client_sqlite_store/account/
accounts.rs

1//! Account-related database operations.
2
3use std::collections::BTreeMap;
4use std::rc::Rc;
5use std::string::{String, ToString};
6use std::sync::{Arc, RwLock};
7use std::vec::Vec;
8
9use miden_client::account::{
10    Account,
11    AccountCode,
12    AccountDelta,
13    AccountHeader,
14    AccountId,
15    AccountStorage,
16    Address,
17    PartialAccount,
18    PartialStorage,
19    PartialStorageMap,
20    StorageMap,
21    StorageMapKey,
22    StorageSlotName,
23    StorageSlotType,
24};
25use miden_client::asset::{Asset, AssetVault, AssetWitness, FungibleAsset};
26use miden_client::store::{
27    AccountRecord,
28    AccountRecordData,
29    AccountSmtForest,
30    AccountStatus,
31    AccountStorageFilter,
32    StoreError,
33};
34use miden_client::sync::NoteTagRecord;
35use miden_client::utils::Serializable;
36use miden_client::{AccountError, Felt, Word};
37use miden_protocol::account::{AccountStorageHeader, StorageMapWitness, StorageSlotHeader};
38use miden_protocol::asset::{AssetVaultKey, PartialVault};
39use miden_protocol::crypto::merkle::MerkleError;
40use rusqlite::types::Value;
41use rusqlite::{Connection, OptionalExtension, Transaction, named_params, params};
42
43use crate::account::helpers::{
44    query_account_addresses,
45    query_account_code,
46    query_historical_account_headers,
47    query_latest_account_headers,
48    query_storage_slots,
49    query_storage_values,
50    query_vault_assets,
51};
52use crate::sql_error::SqlResultExt;
53use crate::sync::{add_note_tag_tx, remove_note_tag_tx};
54use crate::{SqliteStore, column_value_as_u64, insert_sql, subst, u64_to_value};
55
56impl SqliteStore {
57    // READER METHODS
58    // --------------------------------------------------------------------------------------------
59
60    pub(crate) fn get_account_ids(conn: &mut Connection) -> Result<Vec<AccountId>, StoreError> {
61        const QUERY: &str = "SELECT id FROM latest_account_headers";
62
63        conn.prepare_cached(QUERY)
64            .into_store_error()?
65            .query_map([], |row| row.get(0))
66            .expect("no binding parameters used in query")
67            .map(|result| {
68                let id: String = result.map_err(|e| StoreError::ParsingError(e.to_string()))?;
69                Ok(AccountId::from_hex(&id).expect("account id is valid"))
70            })
71            .collect::<Result<Vec<AccountId>, StoreError>>()
72    }
73
74    pub(crate) fn get_account_headers(
75        conn: &mut Connection,
76    ) -> Result<Vec<(AccountHeader, AccountStatus)>, StoreError> {
77        query_latest_account_headers(conn, "1=1 ORDER BY id", params![])
78    }
79
80    pub(crate) fn get_account_header(
81        conn: &mut Connection,
82        account_id: AccountId,
83    ) -> Result<Option<(AccountHeader, AccountStatus)>, StoreError> {
84        Ok(query_latest_account_headers(conn, "id = ?", params![account_id.to_hex()])?.pop())
85    }
86
87    pub(crate) fn get_account_header_by_commitment(
88        conn: &mut Connection,
89        account_commitment: Word,
90    ) -> Result<Option<AccountHeader>, StoreError> {
91        let account_commitment_str: String = account_commitment.to_string();
92        Ok(query_historical_account_headers(
93            conn,
94            "account_commitment = ?",
95            params![account_commitment_str],
96        )?
97        .pop()
98        .map(|(header, _)| header))
99    }
100
101    /// Retrieves a complete account record with full vault and storage data.
102    pub(crate) fn get_account(
103        conn: &mut Connection,
104        account_id: AccountId,
105    ) -> Result<Option<AccountRecord>, StoreError> {
106        let Some((header, status)) = Self::get_account_header(conn, account_id)? else {
107            return Ok(None);
108        };
109
110        let assets = query_vault_assets(conn, account_id)?;
111        let vault = AssetVault::new(&assets)?;
112
113        let slots = query_storage_slots(conn, account_id, &AccountStorageFilter::All)?
114            .into_values()
115            .collect();
116
117        let storage = AccountStorage::new(slots)?;
118
119        let Some(account_code) = query_account_code(conn, header.code_commitment())? else {
120            return Ok(None);
121        };
122
123        let account = Account::new_unchecked(
124            header.id(),
125            vault,
126            storage,
127            account_code,
128            header.nonce(),
129            status.seed().copied(),
130        );
131
132        let account_data = AccountRecordData::Full(account);
133        Ok(Some(AccountRecord::new(account_data, status)))
134    }
135
136    /// Retrieves a minimal partial account record with storage and vault witnesses.
137    pub(crate) fn get_minimal_partial_account(
138        conn: &mut Connection,
139        account_id: AccountId,
140    ) -> Result<Option<AccountRecord>, StoreError> {
141        let Some((header, status)) = Self::get_account_header(conn, account_id)? else {
142            return Ok(None);
143        };
144
145        // Partial vault retrieval
146        let partial_vault = PartialVault::new(header.vault_root());
147
148        // Partial storage retrieval
149        let mut storage_header = Vec::new();
150        let mut maps = vec![];
151
152        let storage_values = query_storage_values(conn, account_id)?;
153
154        // Storage maps are always minimal here (just roots, no entries).
155        // New accounts that need full storage data are handled by the DataStore layer,
156        // which fetches the full account via `get_account()` when nonce == 0.
157        for (slot_name, (slot_type, value)) in storage_values {
158            storage_header.push(StorageSlotHeader::new(slot_name.clone(), slot_type, value));
159            if slot_type == StorageSlotType::Map {
160                maps.push(PartialStorageMap::new(value));
161            }
162        }
163        storage_header.sort_by_key(StorageSlotHeader::id);
164        let storage_header =
165            AccountStorageHeader::new(storage_header).map_err(StoreError::AccountError)?;
166        let partial_storage =
167            PartialStorage::new(storage_header, maps).map_err(StoreError::AccountError)?;
168
169        let Some(account_code) = query_account_code(conn, header.code_commitment())? else {
170            return Ok(None);
171        };
172
173        let partial_account = PartialAccount::new(
174            header.id(),
175            header.nonce(),
176            account_code,
177            partial_storage,
178            partial_vault,
179            status.seed().copied(),
180        )?;
181        let account_record_data = AccountRecordData::Partial(partial_account);
182        Ok(Some(AccountRecord::new(account_record_data, status)))
183    }
184
185    pub fn get_foreign_account_code(
186        conn: &mut Connection,
187        account_ids: Vec<AccountId>,
188    ) -> Result<BTreeMap<AccountId, AccountCode>, StoreError> {
189        let params: Vec<Value> =
190            account_ids.into_iter().map(|id| Value::from(id.to_hex())).collect();
191        const QUERY: &str = "
192            SELECT account_id, code
193            FROM foreign_account_code JOIN account_code ON foreign_account_code.code_commitment = account_code.commitment
194            WHERE account_id IN rarray(?)";
195
196        conn.prepare_cached(QUERY)
197            .into_store_error()?
198            .query_map([Rc::new(params)], |row| Ok((row.get(0)?, row.get(1)?)))
199            .expect("no binding parameters used in query")
200            .map(|result| {
201                result.map_err(|err| StoreError::ParsingError(err.to_string())).and_then(
202                    |(id, code): (String, Vec<u8>)| {
203                        Ok((
204                            AccountId::from_hex(&id).map_err(|err| {
205                                StoreError::AccountError(
206                                    AccountError::FinalAccountHeaderIdParsingFailed(err),
207                                )
208                            })?,
209                            AccountCode::from_bytes(&code).map_err(StoreError::AccountError)?,
210                        ))
211                    },
212                )
213            })
214            .collect::<Result<BTreeMap<AccountId, AccountCode>, _>>()
215    }
216
217    /// Retrieves the full asset vault for a specific account.
218    pub fn get_account_vault(
219        conn: &Connection,
220        account_id: AccountId,
221    ) -> Result<AssetVault, StoreError> {
222        let assets = query_vault_assets(conn, account_id)?;
223        Ok(AssetVault::new(&assets)?)
224    }
225
226    /// Retrieves the full storage for a specific account.
227    pub fn get_account_storage(
228        conn: &Connection,
229        account_id: AccountId,
230        filter: &AccountStorageFilter,
231    ) -> Result<AccountStorage, StoreError> {
232        let slots = query_storage_slots(conn, account_id, filter)?.into_values().collect();
233        Ok(AccountStorage::new(slots)?)
234    }
235
236    /// Fetches a specific asset from the account's vault without the need of loading the entire
237    /// vault. The witness is retrieved from the [`AccountSmtForest`].
238    pub(crate) fn get_account_asset(
239        conn: &mut Connection,
240        smt_forest: &Arc<RwLock<AccountSmtForest>>,
241        account_id: AccountId,
242        vault_key: AssetVaultKey,
243    ) -> Result<Option<(Asset, AssetWitness)>, StoreError> {
244        let header = Self::get_account_header(conn, account_id)?
245            .ok_or(StoreError::AccountDataNotFound(account_id))?
246            .0;
247
248        let smt_forest = smt_forest.read().expect("smt_forest read lock not poisoned");
249        match smt_forest.get_asset_and_witness(header.vault_root(), vault_key) {
250            Ok((asset, witness)) => Ok(Some((asset, witness))),
251            Err(StoreError::MerkleStoreError(MerkleError::UntrackedKey(_))) => Ok(None),
252            Err(err) => Err(err),
253        }
254    }
255
256    /// Retrieves a specific item from the account's storage map without loading the entire storage.
257    /// The witness is retrieved from the [`AccountSmtForest`].
258    pub(crate) fn get_account_map_item(
259        conn: &mut Connection,
260        smt_forest: &Arc<RwLock<AccountSmtForest>>,
261        account_id: AccountId,
262        slot_name: StorageSlotName,
263        key: StorageMapKey,
264    ) -> Result<(Word, StorageMapWitness), StoreError> {
265        let header = Self::get_account_header(conn, account_id)?
266            .ok_or(StoreError::AccountDataNotFound(account_id))?
267            .0;
268
269        let mut storage_values = query_storage_values(conn, 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 = smt_forest.read().expect("smt_forest read lock not poisoned");
278        let witness = smt_forest.get_storage_map_item_witness(map_root, key)?;
279        let item = witness.get(key).unwrap_or(miden_client::EMPTY_WORD);
280
281        Ok((item, witness))
282    }
283
284    pub(crate) fn get_account_addresses(
285        conn: &mut Connection,
286        account_id: AccountId,
287    ) -> Result<Vec<Address>, StoreError> {
288        query_account_addresses(conn, account_id)
289    }
290
291    /// Retrieves the account code for a specific account by ID.
292    pub(crate) fn get_account_code_by_id(
293        conn: &mut Connection,
294        account_id: AccountId,
295    ) -> Result<Option<AccountCode>, StoreError> {
296        let Some((header, _)) =
297            query_latest_account_headers(conn, "id = ?", params![account_id.to_hex()])?
298                .into_iter()
299                .next()
300        else {
301            return Ok(None);
302        };
303
304        query_account_code(conn, header.code_commitment())
305    }
306
307    // MUTATOR/WRITER METHODS
308    // --------------------------------------------------------------------------------------------
309
310    pub(crate) fn insert_account(
311        conn: &mut Connection,
312        smt_forest: &Arc<RwLock<AccountSmtForest>>,
313        account: &Account,
314        initial_address: &Address,
315    ) -> Result<(), StoreError> {
316        let tx = conn.transaction().into_store_error()?;
317
318        Self::insert_account_code(&tx, account.code())?;
319
320        let account_id = account.id();
321
322        Self::insert_storage_slots(&tx, account_id, account.storage().slots().iter())?;
323
324        Self::insert_assets(&tx, account_id, account.vault().assets())?;
325        Self::insert_account_header(&tx, &account.into(), account.seed(), None)?;
326
327        Self::insert_address(&tx, initial_address, account.id())?;
328
329        tx.commit().into_store_error()?;
330
331        let mut smt_forest = smt_forest.write().expect("smt_forest write lock not poisoned");
332        smt_forest.insert_and_register_account_state(
333            account.id(),
334            account.vault(),
335            account.storage(),
336        )?;
337
338        Ok(())
339    }
340
341    pub(crate) fn update_account(
342        conn: &mut Connection,
343        smt_forest: &Arc<RwLock<AccountSmtForest>>,
344        new_account_state: &Account,
345    ) -> Result<(), StoreError> {
346        const QUERY: &str = "SELECT id FROM latest_account_headers WHERE id = ?";
347        if conn
348            .prepare(QUERY)
349            .into_store_error()?
350            .query_map(params![new_account_state.id().to_hex()], |row| row.get(0))
351            .into_store_error()?
352            .map(|result| {
353                result.map_err(|err| StoreError::ParsingError(err.to_string())).and_then(
354                    |id: String| {
355                        AccountId::from_hex(&id).map_err(|err| {
356                            StoreError::AccountError(
357                                AccountError::FinalAccountHeaderIdParsingFailed(err),
358                            )
359                        })
360                    },
361                )
362            })
363            .next()
364            .is_none()
365        {
366            return Err(StoreError::AccountDataNotFound(new_account_state.id()));
367        }
368
369        let mut smt_forest = smt_forest.write().expect("smt_forest write lock not poisoned");
370        let tx = conn.transaction().into_store_error()?;
371        Self::update_account_state(&tx, &mut smt_forest, new_account_state)?;
372        tx.commit().into_store_error()
373    }
374
375    pub fn upsert_foreign_account_code(
376        conn: &mut Connection,
377        account_id: AccountId,
378        code: &AccountCode,
379    ) -> Result<(), StoreError> {
380        let tx = conn.transaction().into_store_error()?;
381
382        Self::insert_account_code(&tx, code)?;
383
384        const QUERY: &str =
385            insert_sql!(foreign_account_code { account_id, code_commitment } | REPLACE);
386
387        tx.execute(QUERY, params![account_id.to_hex(), code.commitment().to_string()])
388            .into_store_error()?;
389
390        Self::insert_account_code(&tx, code)?;
391        tx.commit().into_store_error()
392    }
393
394    pub(crate) fn insert_address(
395        tx: &Transaction<'_>,
396        address: &Address,
397        account_id: AccountId,
398    ) -> Result<(), StoreError> {
399        let derived_note_tag = address.to_note_tag();
400        let note_tag_record = NoteTagRecord::with_account_source(derived_note_tag, account_id);
401
402        add_note_tag_tx(tx, &note_tag_record)?;
403        Self::insert_address_internal(tx, address, account_id)?;
404
405        Ok(())
406    }
407
408    pub(crate) fn remove_address(
409        conn: &mut Connection,
410        address: &Address,
411        account_id: AccountId,
412    ) -> Result<(), StoreError> {
413        let derived_note_tag = address.to_note_tag();
414        let note_tag_record = NoteTagRecord::with_account_source(derived_note_tag, account_id);
415
416        let tx = conn.transaction().into_store_error()?;
417        remove_note_tag_tx(&tx, note_tag_record)?;
418        Self::remove_address_internal(&tx, address)?;
419
420        tx.commit().into_store_error()
421    }
422
423    /// Inserts an [`AccountCode`].
424    pub(crate) fn insert_account_code(
425        tx: &Transaction<'_>,
426        account_code: &AccountCode,
427    ) -> Result<(), StoreError> {
428        const QUERY: &str = insert_sql!(account_code { commitment, code } | IGNORE);
429        tx.execute(QUERY, params![account_code.commitment().to_hex(), account_code.to_bytes()])
430            .into_store_error()?;
431        Ok(())
432    }
433
434    /// Applies the account delta to the account state, updating the vault and storage maps.
435    ///
436    /// Archives old values from latest to historical and updates latest via INSERT OR REPLACE.
437    pub(crate) fn apply_account_delta(
438        tx: &Transaction<'_>,
439        smt_forest: &mut AccountSmtForest,
440        init_account_state: &AccountHeader,
441        final_account_state: &AccountHeader,
442        updated_fungible_assets: BTreeMap<AssetVaultKey, FungibleAsset>,
443        old_map_roots: &BTreeMap<StorageSlotName, Word>,
444        delta: &AccountDelta,
445    ) -> Result<(), StoreError> {
446        let account_id = final_account_state.id();
447
448        // Archive old header and insert the new one
449        Self::insert_account_header(tx, final_account_state, None, Some(init_account_state))?;
450
451        Self::apply_account_vault_delta(
452            tx,
453            smt_forest,
454            account_id,
455            init_account_state,
456            final_account_state,
457            updated_fungible_assets,
458            delta,
459        )?;
460
461        // Build the final roots from the init state's registered roots:
462        // - Replace vault root with the final one
463        // - Replace changed map roots with their new values (done by apply_account_storage_delta)
464        // - Unchanged map roots continue as they were
465        let mut final_roots = smt_forest
466            .get_roots(&init_account_state.id())
467            .cloned()
468            .ok_or(StoreError::AccountDataNotFound(init_account_state.id()))?;
469
470        // First element is always the vault root
471        if let Some(vault_root) = final_roots.first_mut() {
472            *vault_root = final_account_state.vault_root();
473        }
474
475        let default_map_root = StorageMap::default().root();
476        let updated_storage_slots =
477            Self::apply_account_storage_delta(smt_forest, old_map_roots, delta)?;
478
479        // Update map roots in final_roots with new values from the delta
480        for (slot_name, (new_root, slot_type)) in &updated_storage_slots {
481            if *slot_type == StorageSlotType::Map {
482                let old_root = old_map_roots.get(slot_name).copied().unwrap_or(default_map_root);
483                if let Some(root) = final_roots.iter_mut().find(|r| **r == old_root) {
484                    *root = *new_root;
485                } else {
486                    // New map slot not in the old roots — append it
487                    final_roots.push(*new_root);
488                }
489            }
490        }
491
492        Self::write_storage_delta(
493            tx,
494            account_id,
495            final_account_state.nonce().as_canonical_u64(),
496            &updated_storage_slots,
497            delta,
498        )?;
499
500        smt_forest.stage_roots(final_account_state.id(), final_roots);
501
502        Ok(())
503    }
504
505    /// Undoes discarded account states by restoring old values from historical.
506    pub(crate) fn undo_account_state(
507        tx: &Transaction<'_>,
508        smt_forest: &mut AccountSmtForest,
509        discarded_states: &[(AccountId, Word)],
510    ) -> Result<(), StoreError> {
511        if discarded_states.is_empty() {
512            return Ok(());
513        }
514
515        let commitment_params = Rc::new(
516            discarded_states
517                .iter()
518                .map(|(_, commitment)| Value::from(commitment.to_hex()))
519                .collect::<Vec<_>>(),
520        );
521
522        // Step 1: Resolve (account_id, nonce) pairs from both latest and historical headers.
523        // The most recent discarded state is in latest, older ones are in historical.
524        let mut id_nonce_pairs: Vec<(String, u64)> = Vec::new();
525        for query in [
526            "SELECT id, nonce FROM latest_account_headers WHERE account_commitment IN rarray(?)",
527            "SELECT id, nonce FROM historical_account_headers WHERE account_commitment IN rarray(?)",
528        ] {
529            id_nonce_pairs.extend(
530                tx.prepare(query)
531                    .into_store_error()?
532                    .query_map(params![commitment_params.clone()], |row| {
533                        let id: String = row.get(0)?;
534                        let nonce: u64 = column_value_as_u64(row, 1)?;
535                        Ok((id, nonce))
536                    })
537                    .into_store_error()?
538                    .filter_map(Result::ok),
539            );
540        }
541
542        // Step 2: Group nonces by account, sort descending (undo most recent first).
543        // Descending order is needed because each nonce's old value is the state before
544        // that nonce — processing most recent first lets earlier nonces overwrite with
545        // the correct final value.
546        let mut nonces_by_account: BTreeMap<String, Vec<u64>> = BTreeMap::new();
547        for (id, nonce) in &id_nonce_pairs {
548            nonces_by_account.entry(id.clone()).or_default().push(*nonce);
549        }
550        for nonces in nonces_by_account.values_mut() {
551            nonces.sort_unstable();
552            nonces.dedup();
553            nonces.reverse();
554        }
555
556        // Steps 3-5
557        for (account_id_hex, nonces) in &nonces_by_account {
558            Self::undo_account_nonces(tx, account_id_hex, nonces)?;
559        }
560
561        // Step 6: Discard rolled-back states from the in-memory forest
562        for (account_id, _) in discarded_states {
563            smt_forest.discard_roots(*account_id);
564        }
565
566        Ok(())
567    }
568
569    /// Undoes all nonces for a single account: restores old values, restores old header,
570    /// and cleans up consumed historical entries.
571    fn undo_account_nonces(
572        tx: &Transaction<'_>,
573        account_id_hex: &str,
574        nonces: &[u64],
575    ) -> Result<(), StoreError> {
576        // Step 3: Undo each nonce in descending order
577        for &nonce in nonces {
578            let nonce_val = u64_to_value(nonce);
579            Self::restore_old_values_for_nonce(tx, account_id_hex, &nonce_val)?;
580        }
581
582        // Step 4: Restore old header from the earliest discarded nonce
583        let min_nonce = *nonces.last().unwrap();
584        let min_nonce_val = u64_to_value(min_nonce);
585
586        let old_header_exists: bool = tx
587            .query_row(
588                "SELECT COUNT(*) FROM historical_account_headers \
589                 WHERE id = ? AND replaced_at_nonce = ?",
590                params![account_id_hex, &min_nonce_val],
591                |row| row.get::<_, i64>(0),
592            )
593            .into_store_error()?
594            > 0;
595
596        if old_header_exists {
597            tx.execute(
598                "INSERT OR REPLACE INTO latest_account_headers \
599                 (id, account_commitment, code_commitment, storage_commitment, \
600                  vault_root, nonce, account_seed, locked) \
601                 SELECT id, account_commitment, code_commitment, storage_commitment, \
602                        vault_root, nonce, account_seed, locked \
603                 FROM historical_account_headers \
604                 WHERE id = ? AND replaced_at_nonce = ?",
605                params![account_id_hex, &min_nonce_val],
606            )
607            .into_store_error()?;
608        } else {
609            // No previous state — delete the account entirely
610            for table in [
611                "DELETE FROM latest_account_headers WHERE id = ?",
612                "DELETE FROM latest_account_storage WHERE account_id = ?",
613                "DELETE FROM latest_storage_map_entries WHERE account_id = ?",
614                "DELETE FROM latest_account_assets WHERE account_id = ?",
615            ] {
616                tx.execute(table, params![account_id_hex]).into_store_error()?;
617            }
618        }
619
620        // Step 5: Delete all consumed historical entries at the discarded nonces
621        let nonce_params = Rc::new(nonces.iter().map(|n| u64_to_value(*n)).collect::<Vec<_>>());
622        for table in [
623            "historical_account_storage",
624            "historical_storage_map_entries",
625            "historical_account_assets",
626        ] {
627            tx.execute(
628                &format!(
629                    "DELETE FROM {table} WHERE account_id = ? AND replaced_at_nonce IN rarray(?)"
630                ),
631                params![account_id_hex, nonce_params.clone()],
632            )
633            .into_store_error()?;
634        }
635        tx.execute(
636            "DELETE FROM historical_account_headers \
637             WHERE id = ? AND replaced_at_nonce IN rarray(?)",
638            params![account_id_hex, nonce_params],
639        )
640        .into_store_error()?;
641
642        Ok(())
643    }
644
645    /// Restores old values from historical entries for a given nonce.
646    /// Non-NULL old values overwrite latest, NULL old values (new entries) are deleted.
647    fn restore_old_values_for_nonce(
648        tx: &Transaction<'_>,
649        account_id_hex: &str,
650        nonce_val: &rusqlite::types::Value,
651    ) -> Result<(), StoreError> {
652        // Restore storage slots with non-NULL old values
653        tx.execute(
654            "INSERT OR REPLACE INTO latest_account_storage \
655             (account_id, slot_name, slot_value, slot_type) \
656             SELECT account_id, slot_name, old_slot_value, slot_type \
657             FROM historical_account_storage \
658             WHERE account_id = ? AND replaced_at_nonce = ? AND old_slot_value IS NOT NULL",
659            params![account_id_hex, nonce_val],
660        )
661        .into_store_error()?;
662
663        // Delete storage slots that were new (NULL old value)
664        tx.execute(
665            "DELETE FROM latest_account_storage \
666             WHERE account_id = ?1 AND slot_name IN (\
667                 SELECT slot_name FROM historical_account_storage \
668                 WHERE account_id = ?1 AND replaced_at_nonce = ?2 AND old_slot_value IS NULL\
669             )",
670            params![account_id_hex, nonce_val],
671        )
672        .into_store_error()?;
673
674        // Restore map entries with non-NULL old values
675        tx.execute(
676            "INSERT OR REPLACE INTO latest_storage_map_entries \
677             (account_id, slot_name, key, value) \
678             SELECT account_id, slot_name, key, old_value \
679             FROM historical_storage_map_entries \
680             WHERE account_id = ? AND replaced_at_nonce = ? AND old_value IS NOT NULL",
681            params![account_id_hex, nonce_val],
682        )
683        .into_store_error()?;
684
685        // Delete map entries that were new (NULL old value)
686        tx.execute(
687            "DELETE FROM latest_storage_map_entries \
688             WHERE account_id = ?1 AND EXISTS (\
689                 SELECT 1 FROM historical_storage_map_entries h \
690                 WHERE h.account_id = latest_storage_map_entries.account_id \
691                   AND h.slot_name = latest_storage_map_entries.slot_name \
692                   AND h.key = latest_storage_map_entries.key \
693                   AND h.replaced_at_nonce = ?2 AND h.old_value IS NULL\
694             )",
695            params![account_id_hex, nonce_val],
696        )
697        .into_store_error()?;
698
699        // Restore assets with non-NULL old values
700        tx.execute(
701            "INSERT OR REPLACE INTO latest_account_assets \
702             (account_id, vault_key, asset) \
703             SELECT account_id, vault_key, old_asset \
704             FROM historical_account_assets \
705             WHERE account_id = ? AND replaced_at_nonce = ? AND old_asset IS NOT NULL",
706            params![account_id_hex, nonce_val],
707        )
708        .into_store_error()?;
709
710        // Delete assets that were new (NULL old value)
711        tx.execute(
712            "DELETE FROM latest_account_assets \
713             WHERE account_id = ?1 AND vault_key IN (\
714                 SELECT vault_key FROM historical_account_assets \
715                 WHERE account_id = ?1 AND replaced_at_nonce = ?2 AND old_asset IS NULL\
716             )",
717            params![account_id_hex, nonce_val],
718        )
719        .into_store_error()?;
720
721        Ok(())
722    }
723
724    /// Replaces the account state with a completely new one from the network.
725    ///
726    /// Replaces the account state entirely: archives old state to historical,
727    /// clears latest, inserts new state to latest only.
728    pub(crate) fn update_account_state(
729        tx: &Transaction<'_>,
730        smt_forest: &mut AccountSmtForest,
731        new_account_state: &Account,
732    ) -> Result<(), StoreError> {
733        let account_id = new_account_state.id();
734        let account_id_hex = account_id.to_hex();
735        let nonce_val = u64_to_value(new_account_state.nonce().as_canonical_u64());
736
737        // Insert and register account state in the SMT forest (handles old root cleanup)
738        smt_forest.insert_and_register_account_state(
739            account_id,
740            new_account_state.vault(),
741            new_account_state.storage(),
742        )?;
743
744        // Read old header before overwriting
745        let old_header = query_latest_account_headers(tx, "id = ?", params![account_id.to_hex()])?
746            .into_iter()
747            .next()
748            .map(|(header, _)| header);
749
750        // Archive all old entries from latest → historical
751        tx.execute(
752            "INSERT OR REPLACE INTO historical_account_storage \
753             (account_id, replaced_at_nonce, slot_name, old_slot_value, slot_type) \
754             SELECT account_id, ?, slot_name, slot_value, slot_type \
755             FROM latest_account_storage WHERE account_id = ?",
756            params![&nonce_val, &account_id_hex],
757        )
758        .into_store_error()?;
759        tx.execute(
760            "INSERT OR REPLACE INTO historical_storage_map_entries \
761             (account_id, replaced_at_nonce, slot_name, key, old_value) \
762             SELECT account_id, ?, slot_name, key, value \
763             FROM latest_storage_map_entries WHERE account_id = ?",
764            params![&nonce_val, &account_id_hex],
765        )
766        .into_store_error()?;
767        tx.execute(
768            "INSERT OR REPLACE INTO historical_account_assets \
769             (account_id, replaced_at_nonce, vault_key, old_asset) \
770             SELECT account_id, ?, vault_key, asset \
771             FROM latest_account_assets WHERE account_id = ?",
772            params![&nonce_val, &account_id_hex],
773        )
774        .into_store_error()?;
775
776        // Delete all latest entries for this account
777        tx.execute(
778            "DELETE FROM latest_account_storage WHERE account_id = ?",
779            params![&account_id_hex],
780        )
781        .into_store_error()?;
782        tx.execute(
783            "DELETE FROM latest_storage_map_entries WHERE account_id = ?",
784            params![&account_id_hex],
785        )
786        .into_store_error()?;
787        tx.execute(
788            "DELETE FROM latest_account_assets WHERE account_id = ?",
789            params![&account_id_hex],
790        )
791        .into_store_error()?;
792
793        // Insert all new entries into latest only
794        Self::insert_storage_slots(tx, account_id, new_account_state.storage().slots().iter())?;
795        Self::insert_assets(tx, account_id, new_account_state.vault().assets())?;
796
797        // Write NULL historical entries for genuinely new entries that didn't exist
798        // in the old state (INSERT OR IGNORE skips entries already archived above)
799        tx.execute(
800            "INSERT OR IGNORE INTO historical_account_storage \
801             (account_id, replaced_at_nonce, slot_name, old_slot_value, slot_type) \
802             SELECT account_id, ?, slot_name, NULL, slot_type \
803             FROM latest_account_storage WHERE account_id = ?",
804            params![&nonce_val, &account_id_hex],
805        )
806        .into_store_error()?;
807        tx.execute(
808            "INSERT OR IGNORE INTO historical_storage_map_entries \
809             (account_id, replaced_at_nonce, slot_name, key, old_value) \
810             SELECT account_id, ?, slot_name, key, NULL \
811             FROM latest_storage_map_entries WHERE account_id = ?",
812            params![&nonce_val, &account_id_hex],
813        )
814        .into_store_error()?;
815        tx.execute(
816            "INSERT OR IGNORE INTO historical_account_assets \
817             (account_id, replaced_at_nonce, vault_key, old_asset) \
818             SELECT account_id, ?, vault_key, NULL \
819             FROM latest_account_assets WHERE account_id = ?",
820            params![&nonce_val, &account_id_hex],
821        )
822        .into_store_error()?;
823
824        // Insert account header (archives old header to historical)
825        Self::insert_account_header(tx, &new_account_state.into(), None, old_header.as_ref())?;
826
827        Ok(())
828    }
829
830    /// Locks the account if the mismatched digest doesn't belong to a previous account state (stale
831    /// data).
832    pub(crate) fn lock_account_on_unexpected_commitment(
833        tx: &Transaction<'_>,
834        account_id: &AccountId,
835        mismatched_digest: &Word,
836    ) -> Result<(), StoreError> {
837        // Mismatched digests may be due to stale network data. If the mismatched digest is
838        // tracked in the db and corresponds to the mismatched account, it means we
839        // got a past update and shouldn't lock the account.
840        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)";
841        let account_id_hex = account_id.to_hex();
842        let digest_str = mismatched_digest.to_string();
843        let params = named_params! {
844            ":account_id": account_id_hex,
845            ":digest": digest_str
846        };
847
848        let query = format!("UPDATE latest_account_headers SET locked = true {LOCK_CONDITION}");
849        tx.execute(&query, params).into_store_error()?;
850
851        // Also lock historical rows so that undo_account_state preserves the lock.
852        let query = format!("UPDATE historical_account_headers SET locked = true {LOCK_CONDITION}");
853        tx.execute(&query, params).into_store_error()?;
854
855        Ok(())
856    }
857
858    // HELPERS
859    // --------------------------------------------------------------------------------------------
860
861    /// Inserts a new account header into the latest table.
862    ///
863    /// If `old_header` is provided, the old header is archived to the historical table.
864    /// For initial inserts (no previous state), pass `None` for `old_header`.
865    fn insert_account_header(
866        tx: &Transaction<'_>,
867        new_header: &AccountHeader,
868        account_seed: Option<Word>,
869        old_header: Option<&AccountHeader>,
870    ) -> Result<(), StoreError> {
871        // Archive the old header to historical before overwriting latest.
872        if let Some(old) = old_header {
873            let old_id = old.id().to_hex();
874            let old_code_commitment = old.code_commitment().to_string();
875            let old_storage_commitment = old.storage_commitment().to_string();
876            let old_vault_root = old.vault_root().to_string();
877            let old_nonce = u64_to_value(old.nonce().as_canonical_u64());
878            let old_commitment = old.to_commitment().to_string();
879            let replaced_at_nonce = u64_to_value(new_header.nonce().as_canonical_u64());
880
881            // Read the old seed and locked status from latest (if any)
882            let (old_seed, old_locked): (Option<Vec<u8>>, bool) = tx
883                .query_row(
884                    "SELECT account_seed, locked FROM latest_account_headers WHERE id = ?",
885                    params![&old_id],
886                    |row| Ok((row.get(0)?, row.get(1)?)),
887                )
888                .optional()
889                .into_store_error()?
890                .unwrap_or((None, false));
891
892            const HISTORICAL_QUERY: &str = insert_sql!(
893                historical_account_headers {
894                    id,
895                    code_commitment,
896                    storage_commitment,
897                    vault_root,
898                    nonce,
899                    account_seed,
900                    account_commitment,
901                    locked,
902                    replaced_at_nonce
903                } | REPLACE
904            );
905
906            tx.execute(
907                HISTORICAL_QUERY,
908                params![
909                    old_id,
910                    old_code_commitment,
911                    old_storage_commitment,
912                    old_vault_root,
913                    old_nonce,
914                    old_seed,
915                    old_commitment,
916                    old_locked,
917                    replaced_at_nonce,
918                ],
919            )
920            .into_store_error()?;
921        }
922
923        // Write the new header to latest.
924        let id = new_header.id().to_hex();
925        let code_commitment = new_header.code_commitment().to_string();
926        let storage_commitment = new_header.storage_commitment().to_string();
927        let vault_root = new_header.vault_root().to_string();
928        let nonce = u64_to_value(new_header.nonce().as_canonical_u64());
929        let commitment = new_header.to_commitment().to_string();
930        let account_seed = account_seed.map(|seed| seed.to_bytes());
931
932        const LATEST_QUERY: &str = insert_sql!(
933            latest_account_headers {
934                id,
935                code_commitment,
936                storage_commitment,
937                vault_root,
938                nonce,
939                account_seed,
940                account_commitment,
941                locked
942            } | REPLACE
943        );
944
945        tx.execute(
946            LATEST_QUERY,
947            params![
948                id,
949                code_commitment,
950                storage_commitment,
951                vault_root,
952                nonce,
953                account_seed,
954                commitment,
955                false,
956            ],
957        )
958        .into_store_error()?;
959
960        Ok(())
961    }
962
963    fn insert_address_internal(
964        tx: &Transaction<'_>,
965        address: &Address,
966        account_id: AccountId,
967    ) -> Result<(), StoreError> {
968        const QUERY: &str = insert_sql!(addresses { address, account_id } | REPLACE);
969        let serialized_address = address.to_bytes();
970        tx.execute(QUERY, params![serialized_address, account_id.to_hex(),])
971            .into_store_error()?;
972
973        Ok(())
974    }
975
976    /// Prunes historical account states for a single account up to the given nonce.
977    ///
978    /// Deletes all historical entries with `replaced_at_nonce <= up_to_nonce`
979    /// (see DESIGN.md for why this threshold is safe), then removes any account
980    /// code that was only referenced by the deleted headers.
981    pub fn prune_account_history(
982        conn: &mut Connection,
983        account_id: AccountId,
984        up_to_nonce: Felt,
985    ) -> Result<usize, StoreError> {
986        let tx = conn.transaction().into_store_error()?;
987        let account_id_hex = account_id.to_hex();
988        let boundary_val = u64_to_value(up_to_nonce.as_canonical_u64());
989        let mut total_deleted: usize = 0;
990
991        // Collect code commitments from headers we are about to delete.
992        let candidate_code_commitments: Vec<String> = {
993            let mut stmt = tx
994                .prepare(
995                    "SELECT DISTINCT code_commitment FROM historical_account_headers \
996                     WHERE id = ? AND replaced_at_nonce <= ?",
997                )
998                .into_store_error()?;
999            let rows = stmt
1000                .query_map(params![&account_id_hex, &boundary_val], |row| row.get(0))
1001                .into_store_error()?;
1002            rows.collect::<Result<Vec<String>, _>>().into_store_error()?
1003        };
1004
1005        // Delete historical entries.
1006        total_deleted += tx
1007            .execute(
1008                "DELETE FROM historical_account_headers \
1009                 WHERE id = ? AND replaced_at_nonce <= ?",
1010                params![&account_id_hex, &boundary_val],
1011            )
1012            .into_store_error()?;
1013
1014        total_deleted += tx
1015            .execute(
1016                "DELETE FROM historical_account_storage \
1017                 WHERE account_id = ? AND replaced_at_nonce <= ?",
1018                params![&account_id_hex, &boundary_val],
1019            )
1020            .into_store_error()?;
1021
1022        total_deleted += tx
1023            .execute(
1024                "DELETE FROM historical_storage_map_entries \
1025                 WHERE account_id = ? AND replaced_at_nonce <= ?",
1026                params![&account_id_hex, &boundary_val],
1027            )
1028            .into_store_error()?;
1029
1030        total_deleted += tx
1031            .execute(
1032                "DELETE FROM historical_account_assets \
1033                 WHERE account_id = ? AND replaced_at_nonce <= ?",
1034                params![&account_id_hex, &boundary_val],
1035            )
1036            .into_store_error()?;
1037
1038        // Delete orphaned code: only check commitments from the deleted headers,
1039        // and only if they are not referenced by any remaining header or foreign code.
1040        for commitment in &candidate_code_commitments {
1041            let still_referenced: bool = tx
1042                .query_row(
1043                    "SELECT EXISTS(
1044                        SELECT 1 FROM latest_account_headers WHERE code_commitment = ?1
1045                        UNION ALL
1046                        SELECT 1 FROM historical_account_headers WHERE code_commitment = ?1
1047                        UNION ALL
1048                        SELECT 1 FROM foreign_account_code WHERE code_commitment = ?1
1049                    )",
1050                    params![commitment],
1051                    |row| row.get(0),
1052                )
1053                .into_store_error()?;
1054
1055            if !still_referenced {
1056                total_deleted += tx
1057                    .execute("DELETE FROM account_code WHERE commitment = ?", params![commitment])
1058                    .into_store_error()?;
1059            }
1060        }
1061
1062        tx.commit().into_store_error()?;
1063        Ok(total_deleted)
1064    }
1065
1066    fn remove_address_internal(tx: &Transaction<'_>, address: &Address) -> Result<(), StoreError> {
1067        let serialized_address = address.to_bytes();
1068
1069        const DELETE_QUERY: &str = "DELETE FROM addresses WHERE address = ?";
1070        tx.execute(DELETE_QUERY, params![serialized_address]).into_store_error()?;
1071
1072        Ok(())
1073    }
1074}