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: String,
35}
36struct SerializedPartialBlockchainNodeParts {
37 id: u64,
38 node: String,
39}
40
41impl SqliteStore {
42 pub(crate) fn insert_block_header(
43 conn: &mut Connection,
44 block_header: &BlockHeader,
45 has_client_notes: bool,
46 ) -> Result<(), StoreError> {
47 let tx = conn.transaction().into_store_error()?;
48
49 Self::insert_block_header_tx(&tx, block_header, has_client_notes)?;
50
51 tx.commit().into_store_error()?;
52 Ok(())
53 }
54
55 pub(crate) fn get_block_headers(
56 conn: &mut Connection,
57 block_numbers: &BTreeSet<BlockNumber>,
58 ) -> Result<Vec<(BlockHeader, BlockRelevance)>, StoreError> {
59 let block_number_list = block_numbers
60 .iter()
61 .map(|block_number| Value::Integer(i64::from(block_number.as_u32())))
62 .collect::<Vec<Value>>();
63
64 const QUERY: &str = "SELECT block_num, header, has_client_notes FROM block_headers WHERE block_num IN rarray(?)";
65
66 conn.prepare(QUERY)
67 .into_store_error()?
68 .query_map(params![Rc::new(block_number_list)], parse_block_headers_columns)
69 .into_store_error()?
70 .map(|result| {
71 let serialized_block_header_parts: SerializedBlockHeaderParts =
72 result.into_store_error()?;
73 parse_block_header(&serialized_block_header_parts)
74 })
75 .collect()
76 }
77
78 pub(crate) fn get_tracked_block_headers(
79 conn: &mut Connection,
80 ) -> Result<Vec<BlockHeader>, StoreError> {
81 const QUERY: &str = "SELECT block_num, header, has_client_notes FROM block_headers WHERE has_client_notes=true";
82 conn.prepare(QUERY)
83 .into_store_error()?
84 .query_map(params![], parse_block_headers_columns)
85 .into_store_error()?
86 .map(|result| {
87 let serialized_block_header_parts: SerializedBlockHeaderParts =
88 result.into_store_error()?;
89 parse_block_header(&serialized_block_header_parts).map(|(block, _)| block)
90 })
91 .collect()
92 }
93
94 pub(crate) fn get_tracked_block_header_numbers(
95 conn: &mut Connection,
96 ) -> Result<BTreeSet<usize>, StoreError> {
97 const QUERY: &str = "SELECT block_num FROM block_headers WHERE has_client_notes=true";
98 conn.prepare(QUERY)
99 .into_store_error()?
100 .query_map(params![], |row| row.get::<_, u32>(0))
101 .into_store_error()?
102 .map(|result| {
103 let block_num: u32 = result.into_store_error()?;
104 Ok(block_num as usize)
105 })
106 .collect()
107 }
108
109 pub(crate) fn get_partial_blockchain_nodes(
110 conn: &mut Connection,
111 filter: &PartialBlockchainFilter,
112 ) -> Result<BTreeMap<InOrderIndex, Word>, StoreError> {
113 match filter {
114 PartialBlockchainFilter::All => query_partial_blockchain_nodes(
115 conn,
116 "SELECT id, node FROM partial_blockchain_nodes",
117 params![],
118 ),
119
120 PartialBlockchainFilter::List(ids) if ids.is_empty() => Ok(BTreeMap::new()),
121 PartialBlockchainFilter::List(ids) => {
122 let id_values = ids
123 .iter()
124 .map(|id| Value::Integer(i64::try_from(id.inner()).expect("id is a valid i64")))
125 .collect::<Vec<_>>();
126
127 query_partial_blockchain_nodes(
128 conn,
129 "SELECT id, node FROM partial_blockchain_nodes WHERE id IN rarray(?)",
130 params_from_iter([Rc::new(id_values)]),
131 )
132 },
133
134 PartialBlockchainFilter::Forest(forest) if forest.is_empty() => Ok(BTreeMap::new()),
135 PartialBlockchainFilter::Forest(forest) => {
136 let max_index = i64::try_from(forest.rightmost_in_order_index().inner())
137 .expect("id is a valid i64");
138
139 query_partial_blockchain_nodes(
140 conn,
141 "SELECT id, node FROM partial_blockchain_nodes WHERE id <= ?",
142 params![max_index],
143 )
144 },
145 }
146 }
147
148 pub(crate) fn get_current_blockchain_peaks(
149 conn: &mut Connection,
150 ) -> Result<MmrPeaks, StoreError> {
151 const QUERY: &str =
152 "SELECT block_num, partial_blockchain_peaks FROM blockchain_checkpoint LIMIT 1";
153
154 let row: Option<(u32, Vec<u8>)> = conn
155 .prepare(QUERY)
156 .into_store_error()?
157 .query_row(params![], |row| Ok((row.get(0)?, row.get(1)?)))
158 .optional()
159 .into_store_error()?;
160
161 match row {
162 Some((block_num, peaks_bytes)) if !peaks_bytes.is_empty() => {
163 parse_partial_blockchain_peaks(block_num, &peaks_bytes)
164 },
165 _ => Ok(MmrPeaks::new(Forest::empty(), vec![])?),
166 }
167 }
168
169 pub fn insert_partial_blockchain_nodes(
170 conn: &mut Connection,
171 nodes: &[(InOrderIndex, Word)],
172 ) -> Result<(), StoreError> {
173 let tx = conn.transaction().into_store_error()?;
174
175 Self::insert_partial_blockchain_nodes_tx(&tx, nodes)?;
176 tx.commit().into_store_error()?;
177 Ok(())
178 }
179
180 pub(crate) fn insert_partial_blockchain_nodes_tx(
182 tx: &Transaction<'_>,
183 nodes: &[(InOrderIndex, Word)],
184 ) -> Result<(), StoreError> {
185 for (index, node) in nodes {
186 insert_partial_blockchain_node(tx, *index, *node)?;
187 }
188 Ok(())
189 }
190
191 pub(crate) fn insert_block_header_tx(
196 tx: &Transaction<'_>,
197 block_header: &BlockHeader,
198 has_client_notes: bool,
199 ) -> Result<(), StoreError> {
200 let SerializedBlockHeaderData { block_num, header, has_client_notes } =
201 serialize_block_header(block_header, has_client_notes);
202 const QUERY: &str =
203 insert_sql!(block_headers { block_num, header, has_client_notes } | IGNORE);
204 tx.execute(QUERY, params![block_num, header, has_client_notes])
205 .into_store_error()?;
206
207 set_block_header_has_client_notes(tx, u64::from(block_num), has_client_notes)?;
208 Ok(())
209 }
210
211 pub fn prune_irrelevant_blocks(
219 conn: &mut Connection,
220 blocks_to_untrack: &[BlockNumber],
221 node_indices_to_remove: &[InOrderIndex],
222 ) -> Result<(), StoreError> {
223 let tx = conn.transaction().into_store_error()?;
224
225 if !node_indices_to_remove.is_empty() {
227 let id_values = node_indices_to_remove
228 .iter()
229 .map(|id| Value::Integer(i64::try_from(id.inner()).expect("id is a valid i64")))
230 .collect::<Vec<_>>();
231
232 tx.execute(
233 "DELETE FROM partial_blockchain_nodes WHERE id IN rarray(?)",
234 params![Rc::new(id_values)],
235 )
236 .into_store_error()?;
237 }
238
239 if !blocks_to_untrack.is_empty() {
241 let block_values = blocks_to_untrack
242 .iter()
243 .map(|b| Value::Integer(i64::from(b.as_u32())))
244 .collect::<Vec<_>>();
245
246 tx.execute(
247 "UPDATE block_headers SET has_client_notes = 0 WHERE block_num IN rarray(?)",
248 params![Rc::new(block_values)],
249 )
250 .into_store_error()?;
251 }
252
253 let genesis: u32 = BlockNumber::GENESIS.as_u32();
255
256 let sync_block: Option<u32> = tx
257 .query_row("SELECT block_num FROM blockchain_checkpoint LIMIT 1", [], |r| r.get(0))
258 .optional()
259 .into_store_error()?;
260
261 if let Some(sync_height) = sync_block {
262 tx.execute(
263 "DELETE FROM block_headers \
264 WHERE has_client_notes = 0 \
265 AND block_num > ?1 \
266 AND block_num < ?2",
267 rusqlite::params![genesis, sync_height],
268 )
269 .into_store_error()?;
270 }
271
272 tx.commit().into_store_error()
273 }
274}
275
276fn insert_partial_blockchain_node(
281 tx: &Transaction<'_>,
282 id: InOrderIndex,
283 node: Word,
284) -> Result<(), StoreError> {
285 let SerializedPartialBlockchainNodeData { id, node } =
286 serialize_partial_blockchain_node(id, node);
287 const QUERY: &str = insert_sql!(partial_blockchain_nodes { id, node } | IGNORE);
288 tx.execute(QUERY, params![id, node]).into_store_error()?;
289 Ok(())
290}
291
292fn query_partial_blockchain_nodes<P: rusqlite::Params>(
293 conn: &mut Connection,
294 sql: &str,
295 params: P,
296) -> Result<BTreeMap<InOrderIndex, Word>, StoreError> {
297 let mut stmt = conn.prepare_cached(sql).into_store_error()?;
298
299 stmt.query_map(params, parse_partial_blockchain_nodes_columns)
300 .into_store_error()?
301 .map(|row_res| {
302 let parts: SerializedPartialBlockchainNodeParts = row_res.into_store_error()?;
303 parse_partial_blockchain_nodes(&parts)
304 })
305 .collect()
306}
307
308fn parse_partial_blockchain_peaks(forest: u32, peaks_nodes: &[u8]) -> Result<MmrPeaks, StoreError> {
309 let mmr_peaks_nodes = Vec::<Word>::read_from_bytes(peaks_nodes)?;
310
311 let forest_size = usize::try_from(forest).expect("u64 should fit in usize");
312 let forest = Forest::new(forest_size).map_err(|err| {
313 StoreError::DataDeserializationError(
314 miden_protocol::utils::serde::DeserializationError::InvalidValue(format!(
315 "invalid forest size {forest_size}: {err}"
316 )),
317 )
318 })?;
319 MmrPeaks::new(forest, mmr_peaks_nodes).map_err(StoreError::MmrError)
320}
321
322fn serialize_block_header(
323 block_header: &BlockHeader,
324 has_client_notes: bool,
325) -> SerializedBlockHeaderData {
326 let block_num = block_header.block_num();
327 let header = block_header.to_bytes();
328
329 SerializedBlockHeaderData {
330 block_num: block_num.as_u32(),
331 header,
332 has_client_notes,
333 }
334}
335
336fn parse_block_headers_columns(
337 row: &rusqlite::Row<'_>,
338) -> Result<SerializedBlockHeaderParts, rusqlite::Error> {
339 let block_num: u32 = row.get(0)?;
340 let header: Vec<u8> = row.get(1)?;
341 let has_client_notes: bool = row.get(2)?;
342
343 Ok(SerializedBlockHeaderParts {
344 _block_num: u64::from(block_num),
345 header,
346 has_client_notes,
347 })
348}
349
350fn parse_block_header(
351 serialized_block_header_parts: &SerializedBlockHeaderParts,
352) -> Result<(BlockHeader, BlockRelevance), StoreError> {
353 Ok((
354 BlockHeader::read_from_bytes(&serialized_block_header_parts.header)?,
355 serialized_block_header_parts.has_client_notes.into(),
356 ))
357}
358
359fn serialize_partial_blockchain_node(
360 id: InOrderIndex,
361 node: Word,
362) -> SerializedPartialBlockchainNodeData {
363 let id = i64::try_from(id.inner()).expect("id is a valid i64");
364 let node = node.to_hex();
365 SerializedPartialBlockchainNodeData { id, node }
366}
367
368fn parse_partial_blockchain_nodes_columns(
369 row: &rusqlite::Row<'_>,
370) -> Result<SerializedPartialBlockchainNodeParts, rusqlite::Error> {
371 let id: u64 = row.get(0)?;
372 let node = row.get(1)?;
373 Ok(SerializedPartialBlockchainNodeParts { id, node })
374}
375
376fn parse_partial_blockchain_nodes(
377 serialized_partial_blockchain_node_parts: &SerializedPartialBlockchainNodeParts,
378) -> Result<(InOrderIndex, Word), StoreError> {
379 let id = InOrderIndex::new(
380 NonZeroUsize::new(
381 usize::try_from(serialized_partial_blockchain_node_parts.id)
382 .expect("id is u64, should not fail"),
383 )
384 .unwrap(),
385 );
386 let node: Word = Word::try_from(&serialized_partial_blockchain_node_parts.node)?;
387 Ok((id, node))
388}
389
390pub(crate) fn set_block_header_has_client_notes(
391 tx: &Transaction<'_>,
392 block_num: u64,
393 has_client_notes: bool,
394) -> Result<(), StoreError> {
395 const QUERY: &str = "\
397 UPDATE block_headers
398 SET has_client_notes=?
399 WHERE block_num=? AND has_client_notes=FALSE;";
400 tx.execute(QUERY, params![has_client_notes, block_num]).into_store_error()?;
401 Ok(())
402}
403
404#[cfg(test)]
405mod test {
406 use std::collections::{BTreeMap, BTreeSet};
407 use std::vec::Vec;
408
409 use miden_client::Word;
410 use miden_client::block::BlockHeader;
411 use miden_client::crypto::{Forest, InOrderIndex, MmrPeaks};
412 use miden_client::note::BlockNumber;
413 use miden_client::store::Store;
414 use miden_client::utils::Serializable;
415 use miden_protocol::crypto::merkle::mmr::Mmr;
416 use miden_protocol::transaction::TransactionKernel;
417 use rusqlite::params;
418
419 use crate::SqliteStore;
420 use crate::tests::create_test_store;
421
422 async fn insert_dummy_block_headers(store: &mut SqliteStore) -> Vec<BlockHeader> {
423 let block_headers: Vec<BlockHeader> = (0..5)
424 .map(|block_num| {
425 BlockHeader::mock(block_num, None, None, &[], TransactionKernel.to_commitment())
426 })
427 .collect();
428
429 let block_headers_clone = block_headers.clone();
430 store
431 .interact_with_connection(move |conn| {
432 let tx = conn.transaction().unwrap();
433 (0..5).for_each(|block_num| {
434 SqliteStore::insert_block_header_tx(
435 &tx,
436 &block_headers_clone[block_num],
437 false,
438 )
439 .unwrap();
440 });
441 tx.commit().unwrap();
442 Ok(())
443 })
444 .await
445 .unwrap();
446
447 block_headers
448 }
449
450 #[tokio::test]
451 async fn insert_and_get_block_headers_by_number() {
452 let mut store = create_test_store().await;
453 let block_headers = insert_dummy_block_headers(&mut store).await;
454
455 let block_header = Store::get_block_header_by_num(&store, 3.into()).await.unwrap().unwrap();
456 assert_eq!(block_headers[3], block_header.0);
457 }
458
459 #[tokio::test]
460 async fn insert_and_get_block_headers_by_list() {
461 let mut store = create_test_store().await;
462 let mock_block_headers = insert_dummy_block_headers(&mut store).await;
463
464 let block_headers: Vec<BlockHeader> =
465 Store::get_block_headers(&store, &[1.into(), 3.into()].into_iter().collect())
466 .await
467 .unwrap()
468 .into_iter()
469 .map(|(block_header, _has_notes)| block_header)
470 .collect();
471 assert_eq!(
472 &[mock_block_headers[1].clone(), mock_block_headers[3].clone()],
473 &block_headers[..]
474 );
475 }
476
477 #[tokio::test]
479 async fn partial_mmr_reconstructs_after_multiple_prune() {
480 let store = create_test_store().await;
484 const TOTAL_BLOCKS: usize = 7300;
485
486 let tx_kernel_commitment = TransactionKernel.to_commitment();
487 let block_headers: Vec<BlockHeader> = (0..TOTAL_BLOCKS)
488 .map(|block_num| {
489 BlockHeader::mock(
490 u32::try_from(block_num).unwrap(),
491 None,
492 None,
493 &[],
494 tx_kernel_commitment,
495 )
496 })
497 .collect();
498
499 let mut mmr = Mmr::default();
500 for header in &block_headers {
501 mmr.add(header.commitment()).expect("valid MMR append");
502 }
503
504 let mut tracked_set: BTreeSet<usize> = (0..(TOTAL_BLOCKS - 1)).step_by(97).collect();
505 tracked_set.insert(TOTAL_BLOCKS - 2);
506 let tracked_blocks: Vec<usize> = tracked_set.iter().copied().collect();
507
508 let mut tracked_nodes: BTreeMap<InOrderIndex, Word> = BTreeMap::new();
509 for &block_num in &tracked_blocks {
510 let header = &block_headers[block_num];
511 tracked_nodes.insert(InOrderIndex::from_leaf_pos(block_num), header.commitment());
512
513 let proof = mmr.open(block_num).expect("valid proof");
514 let mut idx = InOrderIndex::from_leaf_pos(block_num);
515 for node in proof.merkle_path().nodes() {
516 tracked_nodes.insert(idx.sibling(), *node);
517 idx = idx.parent();
518 }
519 }
520 let tracked_nodes: Vec<(InOrderIndex, Word)> = tracked_nodes.into_iter().collect();
521
522 let peaks_by_block: Vec<MmrPeaks> = (0..TOTAL_BLOCKS)
523 .map(|block_num| {
524 mmr.peaks_at(Forest::new(block_num).expect("valid forest"))
525 .expect("valid peaks")
526 })
527 .collect();
528
529 store
531 .interact_with_connection(move |conn| {
532 let tx = conn.transaction().unwrap();
533 for (block_num, block_header) in block_headers.iter().enumerate() {
534 let has_notes = tracked_set.contains(&block_num);
535 SqliteStore::insert_block_header_tx(&tx, block_header, has_notes).unwrap();
536 }
537
538 SqliteStore::insert_partial_blockchain_nodes_tx(&tx, &tracked_nodes).unwrap();
539 tx.commit().unwrap();
540 Ok(())
541 })
542 .await
543 .unwrap();
544
545 let prune_heights = [
546 TOTAL_BLOCKS / 5,
547 (TOTAL_BLOCKS * 2) / 5,
548 (TOTAL_BLOCKS * 3) / 5,
549 TOTAL_BLOCKS - 1,
550 ];
551
552 let mut previous_remaining: Option<i64> = None;
556 for height in prune_heights {
557 let height_i64 = i64::try_from(height).expect("fits in i64");
558 let peaks_bytes = peaks_by_block[height].peaks().to_vec().to_bytes();
559
560 store
562 .interact_with_connection(move |conn| {
563 conn.execute(
564 "UPDATE blockchain_checkpoint SET block_num = ?, partial_blockchain_peaks = ?",
565 params![height_i64, peaks_bytes],
566 )
567 .unwrap();
568 Ok(())
569 })
570 .await
571 .unwrap();
572
573 store.untrack_and_prune_irrelevant_blocks(&[], &[]).await.unwrap();
575
576 let remaining_headers: i64 = store
578 .interact_with_connection(|conn| {
579 let count = conn
580 .query_row("SELECT COUNT(*) FROM block_headers", [], |row| row.get(0))
581 .unwrap();
582 Ok(count)
583 })
584 .await
585 .unwrap();
586 if let Some(previous) = previous_remaining {
587 assert!(remaining_headers < previous);
588 } else {
589 assert!(remaining_headers < i64::try_from(TOTAL_BLOCKS).unwrap());
590 }
591 previous_remaining = Some(remaining_headers);
592 }
593
594 let partial_mmr = Store::get_current_partial_mmr(&store).await.unwrap();
596 assert_eq!(partial_mmr.peaks().hash_peaks(), mmr.peaks().hash_peaks());
597
598 for block_num in tracked_blocks {
599 let partial_proof = partial_mmr.open(block_num).expect("partial mmr query succeeds");
600 assert!(partial_proof.is_some());
601 assert_eq!(
602 partial_proof.unwrap().merkle_path(),
603 mmr.open(block_num).unwrap().merkle_path()
604 );
605 }
606 }
607
608 fn collect_auth_nodes(
610 mmr: &Mmr,
611 block_headers: &[BlockHeader],
612 tracked: &BTreeSet<usize>,
613 ) -> Vec<(InOrderIndex, Word)> {
614 let mut nodes: BTreeMap<InOrderIndex, Word> = BTreeMap::new();
615 for &block_num in tracked {
616 nodes.insert(
617 InOrderIndex::from_leaf_pos(block_num),
618 block_headers[block_num].commitment(),
619 );
620 let proof = mmr.open(block_num).expect("valid proof");
621 let mut idx = InOrderIndex::from_leaf_pos(block_num);
622 for node in proof.merkle_path().nodes() {
623 nodes.insert(idx.sibling(), *node);
624 idx = idx.parent();
625 }
626 }
627 nodes.into_iter().collect()
628 }
629
630 #[tokio::test]
633 async fn prune_irrelevant_blocks_removes_redundant_auth_nodes() {
634 let store = create_test_store().await;
635 const TOTAL_BLOCKS: usize = 16;
636 let tx_kernel = TransactionKernel.to_commitment();
637
638 let headers: Vec<BlockHeader> = (0..TOTAL_BLOCKS)
639 .map(|n| BlockHeader::mock(u32::try_from(n).unwrap(), None, None, &[], tx_kernel))
640 .collect();
641 let mut mmr = Mmr::default();
642 for h in &headers {
643 mmr.add(h.commitment()).expect("valid MMR append");
644 }
645
646 let tracked: BTreeSet<usize> = [3, 10].into();
648 let auth_nodes = collect_auth_nodes(&mmr, &headers, &tracked);
649 let tip_peaks_bytes = mmr
650 .peaks_at(Forest::new(TOTAL_BLOCKS - 1).expect("valid forest"))
651 .unwrap()
652 .peaks()
653 .to_vec()
654 .to_bytes();
655
656 let headers_clone = headers.clone();
658 store
659 .interact_with_connection(move |conn| {
660 let tx = conn.transaction().unwrap();
661 for (i, header) in headers_clone.iter().enumerate().take(TOTAL_BLOCKS) {
662 SqliteStore::insert_block_header_tx(&tx, header, tracked.contains(&i)).unwrap();
663 }
664 SqliteStore::insert_partial_blockchain_nodes_tx(&tx, &auth_nodes).unwrap();
665 tx.execute(
666 "UPDATE blockchain_checkpoint SET block_num = ?, partial_blockchain_peaks = ?",
667 params![i64::try_from(TOTAL_BLOCKS - 1).unwrap(), tip_peaks_bytes],
668 )
669 .unwrap();
670 tx.commit().unwrap();
671 Ok(())
672 })
673 .await
674 .unwrap();
675
676 let mut partial_mmr = Store::get_current_partial_mmr(&store).await.unwrap();
678 let removed: Vec<InOrderIndex> =
679 partial_mmr.untrack(3).into_iter().map(|(idx, _)| idx).collect();
680 assert!(!removed.is_empty(), "untracking should remove at least one node");
681
682 store
683 .untrack_and_prune_irrelevant_blocks(&[BlockNumber::from(3u32)], &removed)
684 .await
685 .unwrap();
686
687 let rebuilt = Store::get_current_partial_mmr(&store).await.unwrap();
689 assert_eq!(rebuilt.peaks().hash_peaks(), mmr.peaks().hash_peaks());
690
691 let proof_10 = rebuilt.open(10).expect("open succeeds");
692 assert!(proof_10.is_some(), "block 10 should still be provable");
693
694 let proof_3 = rebuilt.open(3).expect("open succeeds");
695 assert!(proof_3.is_none(), "block 3 should no longer be provable");
696 }
697}