Skip to main content

miden_client_sqlite_store/
chain_data.rs

1#![allow(clippy::items_after_statements)]
2
3use std::collections::{BTreeMap, BTreeSet};
4use std::num::NonZeroUsize;
5use std::rc::Rc;
6use std::vec::Vec;
7
8use miden_client::Word;
9use miden_client::block::BlockHeader;
10use miden_client::crypto::{Forest, InOrderIndex, MmrPeaks};
11use miden_client::note::BlockNumber;
12use miden_client::store::{BlockRelevance, PartialBlockchainFilter, StoreError};
13use miden_client::utils::{Deserializable, Serializable};
14use rusqlite::types::Value;
15use rusqlite::{Connection, OptionalExtension, Transaction, params, params_from_iter};
16
17use super::SqliteStore;
18use crate::sql_error::SqlResultExt;
19use crate::{insert_sql, subst};
20
21struct SerializedBlockHeaderData {
22    block_num: u32,
23    header: Vec<u8>,
24    has_client_notes: bool,
25}
26struct SerializedBlockHeaderParts {
27    _block_num: u64,
28    header: Vec<u8>,
29    has_client_notes: bool,
30}
31
32struct SerializedPartialBlockchainNodeData {
33    id: i64,
34    node: Vec<u8>,
35}
36struct SerializedPartialBlockchainNodeParts {
37    id: u64,
38    node: Vec<u8>,
39}
40
41impl SqliteStore {
42    pub(crate) fn get_block_headers(
43        conn: &mut Connection,
44        block_numbers: &BTreeSet<BlockNumber>,
45    ) -> Result<Vec<(BlockHeader, BlockRelevance)>, StoreError> {
46        let block_number_list = block_numbers
47            .iter()
48            .map(|block_number| Value::Integer(i64::from(block_number.as_u32())))
49            .collect::<Vec<Value>>();
50
51        const QUERY: &str = "SELECT block_num, header, has_client_notes FROM block_headers WHERE block_num IN rarray(?)";
52
53        conn.prepare(QUERY)
54            .into_store_error()?
55            .query_map(params![Rc::new(block_number_list)], parse_block_headers_columns)
56            .into_store_error()?
57            .map(|result| {
58                let serialized_block_header_parts: SerializedBlockHeaderParts =
59                    result.into_store_error()?;
60                parse_block_header(&serialized_block_header_parts)
61            })
62            .collect()
63    }
64
65    pub(crate) fn get_tracked_block_headers(
66        conn: &mut Connection,
67    ) -> Result<Vec<BlockHeader>, StoreError> {
68        // `idx_block_headers_has_notes` is declared `WHERE has_client_notes = 1`, and SQLite
69        // matches a partial index only when the predicate is spelled the same way.
70        const QUERY: &str = "SELECT block_num, header, has_client_notes FROM block_headers WHERE has_client_notes=1";
71        conn.prepare(QUERY)
72            .into_store_error()?
73            .query_map(params![], parse_block_headers_columns)
74            .into_store_error()?
75            .map(|result| {
76                let serialized_block_header_parts: SerializedBlockHeaderParts =
77                    result.into_store_error()?;
78                parse_block_header(&serialized_block_header_parts).map(|(block, _)| block)
79            })
80            .collect()
81    }
82
83    pub(crate) fn get_tracked_block_header_numbers(
84        conn: &mut Connection,
85    ) -> Result<BTreeSet<usize>, StoreError> {
86        const QUERY: &str = "SELECT block_num FROM block_headers WHERE has_client_notes=1";
87        conn.prepare(QUERY)
88            .into_store_error()?
89            .query_map(params![], |row| row.get::<_, u32>(0))
90            .into_store_error()?
91            .map(|result| {
92                let block_num: u32 = result.into_store_error()?;
93                Ok(block_num as usize)
94            })
95            .collect()
96    }
97
98    pub(crate) fn get_partial_blockchain_nodes(
99        conn: &mut Connection,
100        filter: &PartialBlockchainFilter,
101    ) -> Result<BTreeMap<InOrderIndex, Word>, StoreError> {
102        match filter {
103            PartialBlockchainFilter::All => query_partial_blockchain_nodes(
104                conn,
105                "SELECT id, node FROM partial_blockchain_nodes",
106                params![],
107            ),
108
109            PartialBlockchainFilter::List(ids) if ids.is_empty() => Ok(BTreeMap::new()),
110            PartialBlockchainFilter::List(ids) => {
111                let id_values = ids
112                    .iter()
113                    .map(|id| Value::Integer(i64::try_from(id.inner()).expect("id is a valid i64")))
114                    .collect::<Vec<_>>();
115
116                query_partial_blockchain_nodes(
117                    conn,
118                    "SELECT id, node FROM partial_blockchain_nodes WHERE id IN rarray(?)",
119                    params_from_iter([Rc::new(id_values)]),
120                )
121            },
122
123            PartialBlockchainFilter::Forest(forest) if forest.is_empty() => Ok(BTreeMap::new()),
124            PartialBlockchainFilter::Forest(forest) => {
125                let max_index = i64::try_from(forest.rightmost_in_order_index().inner())
126                    .expect("id is a valid i64");
127
128                query_partial_blockchain_nodes(
129                    conn,
130                    "SELECT id, node FROM partial_blockchain_nodes WHERE id <= ?",
131                    params![max_index],
132                )
133            },
134        }
135    }
136
137    pub(crate) fn get_current_blockchain_peaks(
138        conn: &mut Connection,
139    ) -> Result<MmrPeaks, StoreError> {
140        const QUERY: &str =
141            "SELECT block_num, partial_blockchain_peaks FROM blockchain_checkpoint LIMIT 1";
142
143        let row: Option<(u32, Vec<u8>)> = conn
144            .prepare(QUERY)
145            .into_store_error()?
146            .query_row(params![], |row| Ok((row.get(0)?, row.get(1)?)))
147            .optional()
148            .into_store_error()?;
149
150        match row {
151            Some((block_num, peaks_bytes)) if !peaks_bytes.is_empty() => {
152                parse_partial_blockchain_peaks(block_num, &peaks_bytes)
153            },
154            _ => Ok(MmrPeaks::new(Forest::empty(), vec![])?),
155        }
156    }
157
158    pub(crate) fn insert_block_header(
159        conn: &mut Connection,
160        block_header: &BlockHeader,
161        nodes: &[(InOrderIndex, Word)],
162        has_client_notes: bool,
163    ) -> Result<(), StoreError> {
164        let tx = conn.transaction().into_store_error()?;
165
166        Self::insert_block_header_tx(&tx, block_header, has_client_notes)?;
167        Self::insert_partial_blockchain_nodes_tx(&tx, nodes)?;
168        tx.commit().into_store_error()?;
169        Ok(())
170    }
171
172    /// Inserts a list of MMR authentication nodes to the Partial Blockchain nodes table.
173    pub(crate) fn insert_partial_blockchain_nodes_tx(
174        tx: &Transaction<'_>,
175        nodes: &[(InOrderIndex, Word)],
176    ) -> Result<(), StoreError> {
177        for (index, node) in nodes {
178            insert_partial_blockchain_node(tx, *index, *node)?;
179        }
180        Ok(())
181    }
182
183    /// Inserts a block header using a [`rusqlite::Transaction`].
184    ///
185    /// If the block header exists and `has_client_notes` is `true` then the `has_client_notes`
186    /// column is updated to `true` to signify that the block now contains a relevant note.
187    pub(crate) fn insert_block_header_tx(
188        tx: &Transaction<'_>,
189        block_header: &BlockHeader,
190        has_client_notes: bool,
191    ) -> Result<(), StoreError> {
192        let SerializedBlockHeaderData { block_num, header, has_client_notes } =
193            serialize_block_header(block_header, has_client_notes);
194        const QUERY: &str =
195            insert_sql!(block_headers { block_num, header, has_client_notes } | IGNORE);
196        tx.execute(QUERY, params![block_num, header, has_client_notes])
197            .into_store_error()?;
198
199        set_block_header_has_client_notes(tx, u64::from(block_num), has_client_notes)?;
200        Ok(())
201    }
202
203    /// Prunes irrelevant block data from the store.
204    ///
205    /// This performs three operations in a single transaction:
206    /// 1. Deletes MMR authentication nodes at the given `node_indices`.
207    /// 2. Sets `has_client_notes = false` for `blocks_to_untrack`.
208    /// 3. Deletes block headers with `has_client_notes = false` that are not the genesis or
209    ///    sync-height block.
210    pub fn prune_irrelevant_blocks(
211        conn: &mut Connection,
212        blocks_to_untrack: &[BlockNumber],
213        node_indices_to_remove: &[InOrderIndex],
214    ) -> Result<(), StoreError> {
215        let tx = conn.transaction().into_store_error()?;
216
217        // 1. Delete stale MMR authentication nodes.
218        if !node_indices_to_remove.is_empty() {
219            let id_values = node_indices_to_remove
220                .iter()
221                .map(|id| Value::Integer(i64::try_from(id.inner()).expect("id is a valid i64")))
222                .collect::<Vec<_>>();
223
224            tx.execute(
225                "DELETE FROM partial_blockchain_nodes WHERE id IN rarray(?)",
226                params![Rc::new(id_values)],
227            )
228            .into_store_error()?;
229        }
230
231        // 2. Mark untracked blocks as irrelevant.
232        if !blocks_to_untrack.is_empty() {
233            let block_values = blocks_to_untrack
234                .iter()
235                .map(|b| Value::Integer(i64::from(b.as_u32())))
236                .collect::<Vec<_>>();
237
238            tx.execute(
239                "UPDATE block_headers SET has_client_notes = 0 WHERE block_num IN rarray(?)",
240                params![Rc::new(block_values)],
241            )
242            .into_store_error()?;
243        }
244
245        // 3. Delete irrelevant block headers.
246        let genesis: u32 = BlockNumber::GENESIS.as_u32();
247
248        let sync_block: Option<u32> = tx
249            .query_row("SELECT block_num FROM blockchain_checkpoint LIMIT 1", [], |r| r.get(0))
250            .optional()
251            .into_store_error()?;
252
253        if let Some(sync_height) = sync_block {
254            tx.execute(
255                "DELETE FROM block_headers \
256                 WHERE has_client_notes = 0 \
257                 AND block_num > ?1 \
258                 AND block_num < ?2",
259                rusqlite::params![genesis, sync_height],
260            )
261            .into_store_error()?;
262        }
263
264        tx.commit().into_store_error()
265    }
266}
267
268// HELPERS
269// ================================================================================================
270
271/// Inserts a node represented by its in-order index and the node value.
272fn insert_partial_blockchain_node(
273    tx: &Transaction<'_>,
274    id: InOrderIndex,
275    node: Word,
276) -> Result<(), StoreError> {
277    let SerializedPartialBlockchainNodeData { id, node } =
278        serialize_partial_blockchain_node(id, node);
279    const QUERY: &str = insert_sql!(partial_blockchain_nodes { id, node } | IGNORE);
280    tx.execute(QUERY, params![id, node]).into_store_error()?;
281    Ok(())
282}
283
284fn query_partial_blockchain_nodes<P: rusqlite::Params>(
285    conn: &mut Connection,
286    sql: &str,
287    params: P,
288) -> Result<BTreeMap<InOrderIndex, Word>, StoreError> {
289    let mut stmt = conn.prepare_cached(sql).into_store_error()?;
290
291    stmt.query_map(params, parse_partial_blockchain_nodes_columns)
292        .into_store_error()?
293        .map(|row_res| {
294            let parts: SerializedPartialBlockchainNodeParts = row_res.into_store_error()?;
295            parse_partial_blockchain_nodes(&parts)
296        })
297        .collect()
298}
299
300fn parse_partial_blockchain_peaks(forest: u32, peaks_nodes: &[u8]) -> Result<MmrPeaks, StoreError> {
301    let mmr_peaks_nodes = Vec::<Word>::read_from_bytes(peaks_nodes)?;
302
303    let forest_size = usize::try_from(forest).expect("u64 should fit in usize");
304    let forest = Forest::new(forest_size).map_err(|err| {
305        StoreError::DataDeserializationError(
306            miden_protocol::utils::serde::DeserializationError::InvalidValue(format!(
307                "invalid forest size {forest_size}: {err}"
308            )),
309        )
310    })?;
311    MmrPeaks::new(forest, mmr_peaks_nodes).map_err(StoreError::MmrError)
312}
313
314fn serialize_block_header(
315    block_header: &BlockHeader,
316    has_client_notes: bool,
317) -> SerializedBlockHeaderData {
318    let block_num = block_header.block_num();
319    let header = block_header.to_bytes();
320
321    SerializedBlockHeaderData {
322        block_num: block_num.as_u32(),
323        header,
324        has_client_notes,
325    }
326}
327
328fn parse_block_headers_columns(
329    row: &rusqlite::Row<'_>,
330) -> Result<SerializedBlockHeaderParts, rusqlite::Error> {
331    let block_num: u32 = row.get(0)?;
332    let header: Vec<u8> = row.get(1)?;
333    let has_client_notes: bool = row.get(2)?;
334
335    Ok(SerializedBlockHeaderParts {
336        _block_num: u64::from(block_num),
337        header,
338        has_client_notes,
339    })
340}
341
342fn parse_block_header(
343    serialized_block_header_parts: &SerializedBlockHeaderParts,
344) -> Result<(BlockHeader, BlockRelevance), StoreError> {
345    Ok((
346        BlockHeader::read_from_bytes(&serialized_block_header_parts.header)?,
347        serialized_block_header_parts.has_client_notes.into(),
348    ))
349}
350
351fn serialize_partial_blockchain_node(
352    id: InOrderIndex,
353    node: Word,
354) -> SerializedPartialBlockchainNodeData {
355    let id = i64::try_from(id.inner()).expect("id is a valid i64");
356    let node = node.to_bytes();
357    SerializedPartialBlockchainNodeData { id, node }
358}
359
360fn parse_partial_blockchain_nodes_columns(
361    row: &rusqlite::Row<'_>,
362) -> Result<SerializedPartialBlockchainNodeParts, rusqlite::Error> {
363    let id: u64 = row.get(0)?;
364    let node = row.get(1)?;
365    Ok(SerializedPartialBlockchainNodeParts { id, node })
366}
367
368fn parse_partial_blockchain_nodes(
369    serialized_partial_blockchain_node_parts: &SerializedPartialBlockchainNodeParts,
370) -> Result<(InOrderIndex, Word), StoreError> {
371    let id = InOrderIndex::new(
372        NonZeroUsize::new(
373            usize::try_from(serialized_partial_blockchain_node_parts.id)
374                .expect("id is u64, should not fail"),
375        )
376        .unwrap(),
377    );
378    let node: Word = Word::read_from_bytes(&serialized_partial_blockchain_node_parts.node)?;
379    Ok((id, node))
380}
381
382pub(crate) fn set_block_header_has_client_notes(
383    tx: &Transaction<'_>,
384    block_num: u64,
385    has_client_notes: bool,
386) -> Result<(), StoreError> {
387    // Only update to change has_client_notes to true if it was false previously
388    const QUERY: &str = "\
389        UPDATE block_headers
390        SET has_client_notes=?
391        WHERE block_num=? AND has_client_notes=0;";
392    tx.execute(QUERY, params![has_client_notes, block_num]).into_store_error()?;
393    Ok(())
394}
395
396#[cfg(test)]
397mod test {
398    use std::collections::{BTreeMap, BTreeSet};
399    use std::vec::Vec;
400
401    use miden_client::Word;
402    use miden_client::block::BlockHeader;
403    use miden_client::crypto::{Forest, InOrderIndex, MmrPeaks};
404    use miden_client::note::BlockNumber;
405    use miden_client::store::{PartialBlockchainFilter, Store};
406    use miden_client::utils::Serializable;
407    use miden_protocol::crypto::merkle::mmr::Mmr;
408    use miden_protocol::transaction::TransactionKernel;
409    use rusqlite::params;
410
411    use crate::SqliteStore;
412    use crate::tests::create_test_store;
413
414    async fn insert_dummy_block_headers(store: &mut SqliteStore) -> Vec<BlockHeader> {
415        let block_headers: Vec<BlockHeader> = (0..5)
416            .map(|block_num| {
417                BlockHeader::mock(block_num, None, None, &[], TransactionKernel.to_commitment())
418            })
419            .collect();
420
421        let block_headers_clone = block_headers.clone();
422        store
423            .interact_with_connection(move |conn| {
424                let tx = conn.transaction().unwrap();
425                (0..5).for_each(|block_num| {
426                    SqliteStore::insert_block_header_tx(
427                        &tx,
428                        &block_headers_clone[block_num],
429                        false,
430                    )
431                    .unwrap();
432                });
433                tx.commit().unwrap();
434                Ok(())
435            })
436            .await
437            .unwrap();
438
439        block_headers
440    }
441
442    #[tokio::test]
443    async fn insert_and_get_block_headers_by_number() {
444        let mut store = create_test_store().await;
445        let block_headers = insert_dummy_block_headers(&mut store).await;
446
447        let block_header = Store::get_block_header_by_num(&store, 3.into()).await.unwrap().unwrap();
448        assert_eq!(block_headers[3], block_header.0);
449    }
450
451    #[tokio::test]
452    async fn insert_and_get_block_headers_by_list() {
453        let mut store = create_test_store().await;
454        let mock_block_headers = insert_dummy_block_headers(&mut store).await;
455
456        let block_headers: Vec<BlockHeader> =
457            Store::get_block_headers(&store, &[1.into(), 3.into()].into_iter().collect())
458                .await
459                .unwrap()
460                .into_iter()
461                .map(|(block_header, _has_notes)| block_header)
462                .collect();
463        assert_eq!(
464            &[mock_block_headers[1].clone(), mock_block_headers[3].clone()],
465            &block_headers[..]
466        );
467    }
468
469    /// Tests that `insert_block_header` persists the tracked header and its MMR
470    /// authentication nodes in the same call, so both are retrievable afterwards.
471    #[tokio::test]
472    async fn insert_block_header_stores_header_and_nodes() {
473        let store = create_test_store().await;
474        const TOTAL_BLOCKS: usize = 8;
475        let tx_kernel = TransactionKernel.to_commitment();
476
477        let headers: Vec<BlockHeader> = (0..TOTAL_BLOCKS)
478            .map(|n| BlockHeader::mock(u32::try_from(n).unwrap(), None, None, &[], tx_kernel))
479            .collect();
480        let mut mmr = Mmr::default();
481        for header in &headers {
482            mmr.add(header.commitment()).expect("valid MMR append");
483        }
484
485        let tracked: BTreeSet<usize> = [5].into();
486        let auth_nodes = collect_auth_nodes(&mmr, &headers, &tracked);
487        let header = headers[5].clone();
488
489        Store::insert_block_header(&store, &header, &auth_nodes, true).await.unwrap();
490
491        // The header is stored and marked as tracked.
492        let stored = Store::get_block_headers(&store, &[5.into()].into_iter().collect())
493            .await
494            .unwrap();
495        assert_eq!(stored.len(), 1);
496        assert_eq!(stored[0].0, header);
497
498        let tracked = Store::get_tracked_block_header_numbers(&store).await.unwrap();
499        assert_eq!(tracked, [5].into());
500
501        // Every authentication node was stored by the same call.
502        let stored_nodes =
503            Store::get_partial_blockchain_nodes(&store, PartialBlockchainFilter::All)
504                .await
505                .unwrap();
506        let expected: BTreeMap<InOrderIndex, Word> = auth_nodes.iter().copied().collect();
507        assert_eq!(stored_nodes, expected);
508    }
509
510    /// Tests that a failure inserting the MMR nodes rolls back the block header written in the
511    /// same call, proving both land in a single transaction.
512    #[tokio::test]
513    async fn insert_block_header_rolls_back_header_when_nodes_fail() {
514        let store = create_test_store().await;
515        let header = BlockHeader::mock(5, None, None, &[], TransactionKernel.to_commitment());
516        // One node so the node insert actually runs (an empty slice would be a no-op).
517        let nodes = [(InOrderIndex::from_leaf_pos(5), header.commitment())];
518
519        // Force the node insert (the second statement in the transaction) to fail.
520        store
521            .interact_with_connection(|conn| {
522                conn.execute("DROP TABLE partial_blockchain_nodes", []).unwrap();
523                Ok(())
524            })
525            .await
526            .unwrap();
527
528        let result = Store::insert_block_header(&store, &header, &nodes, true).await;
529        assert!(result.is_err(), "node insert must fail against the dropped table");
530
531        // The header must not survive: a non-atomic two-transaction insert would leave it behind.
532        let stored = Store::get_block_headers(&store, &[5.into()].into_iter().collect())
533            .await
534            .unwrap();
535        assert!(stored.is_empty(), "header must roll back when the node insert fails");
536    }
537
538    /// Tests that large stored MMRs are built consistently throughout multiple prunes
539    #[tokio::test]
540    async fn partial_mmr_reconstructs_after_multiple_prune() {
541        // Setup (mock a large MMR to work with, with a partial tracked set)
542        // ----------------------------------------------------------------------------------------
543
544        let store = create_test_store().await;
545        const TOTAL_BLOCKS: usize = 7300;
546
547        let tx_kernel_commitment = TransactionKernel.to_commitment();
548        let block_headers: Vec<BlockHeader> = (0..TOTAL_BLOCKS)
549            .map(|block_num| {
550                BlockHeader::mock(
551                    u32::try_from(block_num).unwrap(),
552                    None,
553                    None,
554                    &[],
555                    tx_kernel_commitment,
556                )
557            })
558            .collect();
559
560        let mut mmr = Mmr::default();
561        for header in &block_headers {
562            mmr.add(header.commitment()).expect("valid MMR append");
563        }
564
565        let mut tracked_set: BTreeSet<usize> = (0..(TOTAL_BLOCKS - 1)).step_by(97).collect();
566        tracked_set.insert(TOTAL_BLOCKS - 2);
567        let tracked_blocks: Vec<usize> = tracked_set.iter().copied().collect();
568
569        let mut tracked_nodes: BTreeMap<InOrderIndex, Word> = BTreeMap::new();
570        for &block_num in &tracked_blocks {
571            let header = &block_headers[block_num];
572            tracked_nodes.insert(InOrderIndex::from_leaf_pos(block_num), header.commitment());
573
574            let proof = mmr.open(block_num).expect("valid proof");
575            let mut idx = InOrderIndex::from_leaf_pos(block_num);
576            for node in proof.merkle_path().nodes() {
577                tracked_nodes.insert(idx.sibling(), *node);
578                idx = idx.parent();
579            }
580        }
581        let tracked_nodes: Vec<(InOrderIndex, Word)> = tracked_nodes.into_iter().collect();
582
583        let peaks_by_block: Vec<MmrPeaks> = (0..TOTAL_BLOCKS)
584            .map(|block_num| {
585                mmr.peaks_at(Forest::new(block_num).expect("valid forest"))
586                    .expect("valid peaks")
587            })
588            .collect();
589
590        // Save blocks and nodes
591        store
592            .interact_with_connection(move |conn| {
593                let tx = conn.transaction().unwrap();
594                for (block_num, block_header) in block_headers.iter().enumerate() {
595                    let has_notes = tracked_set.contains(&block_num);
596                    SqliteStore::insert_block_header_tx(&tx, block_header, has_notes).unwrap();
597                }
598
599                SqliteStore::insert_partial_blockchain_nodes_tx(&tx, &tracked_nodes).unwrap();
600                tx.commit().unwrap();
601                Ok(())
602            })
603            .await
604            .unwrap();
605
606        let prune_heights = [
607            TOTAL_BLOCKS / 5,
608            (TOTAL_BLOCKS * 2) / 5,
609            (TOTAL_BLOCKS * 3) / 5,
610            TOTAL_BLOCKS - 1,
611        ];
612
613        // Tests/assertions
614        // ----------------------------------------------------------------------------------------
615
616        let mut previous_remaining: Option<i64> = None;
617        for height in prune_heights {
618            let height_i64 = i64::try_from(height).expect("fits in i64");
619            let peaks_bytes = peaks_by_block[height].peaks().to_vec().to_bytes();
620
621            // Update sync height (and the matching MMR peaks) to simulate having synced further
622            store
623                .interact_with_connection(move |conn| {
624                    conn.execute(
625                        "UPDATE blockchain_checkpoint SET block_num = ?, partial_blockchain_peaks = ?",
626                        params![height_i64, peaks_bytes],
627                    )
628                    .unwrap();
629                    Ok(())
630                })
631                .await
632                .unwrap();
633
634            // Prune
635            store.untrack_and_prune_irrelevant_blocks(&[], &[]).await.unwrap();
636
637            // Assert blocks
638            let remaining_headers: i64 = store
639                .interact_with_connection(|conn| {
640                    let count = conn
641                        .query_row("SELECT COUNT(*) FROM block_headers", [], |row| row.get(0))
642                        .unwrap();
643                    Ok(count)
644                })
645                .await
646                .unwrap();
647            if let Some(previous) = previous_remaining {
648                assert!(remaining_headers < previous);
649            } else {
650                assert!(remaining_headers < i64::try_from(TOTAL_BLOCKS).unwrap());
651            }
652            previous_remaining = Some(remaining_headers);
653        }
654
655        // Try build MMR
656        let partial_mmr = Store::get_current_partial_mmr(&store).await.unwrap();
657        assert_eq!(partial_mmr.peaks().hash_peaks(), mmr.peaks().hash_peaks());
658
659        for block_num in tracked_blocks {
660            let partial_proof = partial_mmr.open(block_num).expect("partial mmr query succeeds");
661            assert!(partial_proof.is_some());
662            assert_eq!(
663                partial_proof.unwrap().merkle_path(),
664                mmr.open(block_num).unwrap().merkle_path()
665            );
666        }
667    }
668
669    /// Collects authentication nodes for a set of tracked leaves in an MMR.
670    fn collect_auth_nodes(
671        mmr: &Mmr,
672        block_headers: &[BlockHeader],
673        tracked: &BTreeSet<usize>,
674    ) -> Vec<(InOrderIndex, Word)> {
675        let mut nodes: BTreeMap<InOrderIndex, Word> = BTreeMap::new();
676        for &block_num in tracked {
677            nodes.insert(
678                InOrderIndex::from_leaf_pos(block_num),
679                block_headers[block_num].commitment(),
680            );
681            let proof = mmr.open(block_num).expect("valid proof");
682            let mut idx = InOrderIndex::from_leaf_pos(block_num);
683            for node in proof.merkle_path().nodes() {
684                nodes.insert(idx.sibling(), *node);
685                idx = idx.parent();
686            }
687        }
688        nodes.into_iter().collect()
689    }
690
691    /// Tests that `untrack_and_prune_irrelevant_blocks` removes redundant authentication nodes
692    /// for untracked blocks while preserving nodes needed by blocks that remain tracked.
693    #[tokio::test]
694    async fn prune_irrelevant_blocks_removes_redundant_auth_nodes() {
695        let store = create_test_store().await;
696        const TOTAL_BLOCKS: usize = 16;
697        let tx_kernel = TransactionKernel.to_commitment();
698
699        let headers: Vec<BlockHeader> = (0..TOTAL_BLOCKS)
700            .map(|n| BlockHeader::mock(u32::try_from(n).unwrap(), None, None, &[], tx_kernel))
701            .collect();
702        let mut mmr = Mmr::default();
703        for h in &headers {
704            mmr.add(h.commitment()).expect("valid MMR append");
705        }
706
707        // Track blocks 3 and 10; we will untrack 3 later.
708        let tracked: BTreeSet<usize> = [3, 10].into();
709        let auth_nodes = collect_auth_nodes(&mmr, &headers, &tracked);
710        let tip_peaks_bytes = mmr
711            .peaks_at(Forest::new(TOTAL_BLOCKS - 1).expect("valid forest"))
712            .unwrap()
713            .peaks()
714            .to_vec()
715            .to_bytes();
716
717        // Persist everything.
718        let headers_clone = headers.clone();
719        store
720            .interact_with_connection(move |conn| {
721                let tx = conn.transaction().unwrap();
722                for (i, header) in headers_clone.iter().enumerate().take(TOTAL_BLOCKS) {
723                    SqliteStore::insert_block_header_tx(&tx, header, tracked.contains(&i)).unwrap();
724                }
725                SqliteStore::insert_partial_blockchain_nodes_tx(&tx, &auth_nodes).unwrap();
726                tx.execute(
727                    "UPDATE blockchain_checkpoint SET block_num = ?, partial_blockchain_peaks = ?",
728                    params![i64::try_from(TOTAL_BLOCKS - 1).unwrap(), tip_peaks_bytes],
729                )
730                .unwrap();
731                tx.commit().unwrap();
732                Ok(())
733            })
734            .await
735            .unwrap();
736
737        // Untrack block 3 via the PartialMmr, then prune.
738        let mut partial_mmr = Store::get_current_partial_mmr(&store).await.unwrap();
739        let removed: Vec<InOrderIndex> =
740            partial_mmr.untrack(3).into_iter().map(|(idx, _)| idx).collect();
741        assert!(!removed.is_empty(), "untracking should remove at least one node");
742
743        store
744            .untrack_and_prune_irrelevant_blocks(&[BlockNumber::from(3u32)], &removed)
745            .await
746            .unwrap();
747
748        // Block 3 header should be deleted, block 10 should still be provable.
749        let rebuilt = Store::get_current_partial_mmr(&store).await.unwrap();
750        assert_eq!(rebuilt.peaks().hash_peaks(), mmr.peaks().hash_peaks());
751
752        let proof_10 = rebuilt.open(10).expect("open succeeds");
753        assert!(proof_10.is_some(), "block 10 should still be provable");
754
755        let proof_3 = rebuilt.open(3).expect("open succeeds");
756        assert!(proof_3.is_none(), "block 3 should no longer be provable");
757    }
758}