Skip to main content

zcash_client_sqlite/wallet/
commitment_tree.rs

1//! SQLite-backed implementation of [`ShardStore`] for note commitment trees.
2
3use rusqlite::{self, OptionalExtension, named_params};
4use std::{
5    collections::BTreeSet,
6    error, fmt,
7    io::{self, Cursor},
8    marker::PhantomData,
9    num::NonZeroU32,
10    ops::Range,
11    sync::Arc,
12};
13
14use incrementalmerkletree::{Address, Hashable, Level, Position, Retention};
15use shardtree::{
16    LocatedPrunableTree, LocatedTree, PrunableTree, RetentionFlags,
17    error::{QueryError, ShardTreeError},
18    store::{Checkpoint, ShardStore, TreeState},
19};
20
21use zcash_client_backend::{
22    data_api::{chain::CommitmentTreeRoot, wallet::TargetHeight},
23    serialization::shardtree::{read_shard, write_shard},
24};
25use zcash_primitives::merkle_tree::HashSer;
26use zcash_protocol::{ShieldedPool, consensus::BlockHeight};
27
28use crate::{error::SqliteClientError, sapling_tree};
29
30#[cfg(feature = "orchard")]
31use {
32    crate::{IRONWOOD_TABLES_PREFIX, ORCHARD_TABLES_PREFIX, ironwood_tree, orchard_tree},
33    incrementalmerkletree::Marking,
34    shardtree::{ShardTree, store::memory::MemoryShardStore},
35    zcash_client_backend::data_api::ORCHARD_SHARD_HEIGHT,
36};
37
38use super::common::{TableConstants, table_constants};
39
40/// Errors that can appear in SQLite-back [`ShardStore`] implementation operations.
41#[derive(Debug)]
42#[non_exhaustive]
43pub enum Error {
44    /// Errors in deserializing stored shard data
45    Serialization(io::Error),
46    /// Errors encountered querying stored shard data
47    Query(rusqlite::Error),
48    /// Raised when the caller attempts to add a checkpoint at a block height where a checkpoint
49    /// already exists, but the tree state being checkpointed or the marks removed at that
50    /// checkpoint conflict with the existing tree state.
51    CheckpointConflict {
52        /// The block height of the conflicting checkpoint.
53        checkpoint_id: BlockHeight,
54        /// The checkpoint data that was attempted to be inserted.
55        checkpoint: Checkpoint,
56        /// The tree state already stored at that checkpoint height.
57        extant_tree_state: TreeState,
58        /// The marks-removed set already stored at that checkpoint height, if any.
59        extant_marks_removed: Option<BTreeSet<Position>>,
60    },
61    /// Raised when attempting to add shard roots to the database that
62    /// are discontinuous with the existing roots in the database.
63    SubtreeDiscontinuity {
64        /// The index range of the subtree roots that were attempted to be inserted.
65        attempted_insertion_range: Range<u64>,
66        /// The index range of subtree roots already present in the database.
67        existing_range: Range<u64>,
68    },
69}
70
71impl fmt::Display for Error {
72    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
73        match &self {
74            Error::Serialization(err) => write!(f, "Commitment tree serialization error: {err}"),
75            Error::Query(err) => write!(f, "Commitment tree query or update error: {err}"),
76            Error::CheckpointConflict {
77                checkpoint_id,
78                checkpoint,
79                extant_tree_state,
80                extant_marks_removed,
81            } => {
82                write!(
83                    f,
84                    "Conflict at checkpoint id {checkpoint_id}, tried to insert {checkpoint:?}, which is incompatible with existing state ({extant_tree_state:?}, {extant_marks_removed:?})"
85                )
86            }
87            Error::SubtreeDiscontinuity {
88                attempted_insertion_range,
89                existing_range,
90            } => {
91                write!(
92                    f,
93                    "Attempted to write subtree roots with indices {attempted_insertion_range:?} which is discontinuous with existing subtree range {existing_range:?}",
94                )
95            }
96        }
97    }
98}
99
100impl error::Error for Error {
101    fn source(&self) -> Option<&(dyn error::Error + 'static)> {
102        match &self {
103            Error::Serialization(e) => Some(e),
104            Error::Query(e) => Some(e),
105            Error::CheckpointConflict { .. } => None,
106            Error::SubtreeDiscontinuity { .. } => None,
107        }
108    }
109}
110
111/// An implementation of [`ShardStore`] backed by an SQLite database.
112pub struct SqliteShardStore<C, H, const SHARD_HEIGHT: u8> {
113    pub(crate) conn: C,
114    table_prefix: &'static str,
115    _hash_type: PhantomData<H>,
116}
117
118impl<C, H, const SHARD_HEIGHT: u8> SqliteShardStore<C, H, SHARD_HEIGHT> {
119    const SHARD_ROOT_LEVEL: Level = Level::new(SHARD_HEIGHT);
120
121    pub(crate) fn from_connection(
122        conn: C,
123        table_prefix: &'static str,
124    ) -> Result<Self, rusqlite::Error> {
125        Ok(SqliteShardStore {
126            conn,
127            table_prefix,
128            _hash_type: PhantomData,
129        })
130    }
131}
132
133impl<'conn, 'a: 'conn, H: HashSer, const SHARD_HEIGHT: u8> ShardStore
134    for SqliteShardStore<&'a rusqlite::Transaction<'conn>, H, SHARD_HEIGHT>
135{
136    type H = H;
137    type CheckpointId = BlockHeight;
138    type Error = Error;
139
140    fn get_shard(
141        &self,
142        shard_root: Address,
143    ) -> Result<Option<LocatedPrunableTree<Self::H>>, Self::Error> {
144        get_shard(self.conn, self.table_prefix, shard_root)
145    }
146
147    fn last_shard(&self) -> Result<Option<LocatedPrunableTree<Self::H>>, Self::Error> {
148        last_shard(self.conn, self.table_prefix, Self::SHARD_ROOT_LEVEL)
149    }
150
151    fn put_shard(&mut self, subtree: LocatedPrunableTree<Self::H>) -> Result<(), Self::Error> {
152        put_shard(self.conn, self.table_prefix, subtree)
153    }
154
155    fn get_shard_roots(&self) -> Result<Vec<Address>, Self::Error> {
156        get_shard_roots(self.conn, self.table_prefix, Self::SHARD_ROOT_LEVEL)
157    }
158
159    fn truncate_shards(&mut self, shard_index: u64) -> Result<(), Self::Error> {
160        truncate_shards(self.conn, self.table_prefix, shard_index)
161    }
162
163    fn get_cap(&self) -> Result<PrunableTree<Self::H>, Self::Error> {
164        get_cap(self.conn, self.table_prefix)
165    }
166
167    fn put_cap(&mut self, cap: PrunableTree<Self::H>) -> Result<(), Self::Error> {
168        put_cap(self.conn, self.table_prefix, cap)
169    }
170
171    fn min_checkpoint_id(&self) -> Result<Option<Self::CheckpointId>, Self::Error> {
172        min_checkpoint_id(self.conn, self.table_prefix)
173    }
174
175    fn max_checkpoint_id(&self) -> Result<Option<Self::CheckpointId>, Self::Error> {
176        max_checkpoint_id(self.conn, self.table_prefix)
177    }
178
179    fn add_checkpoint(
180        &mut self,
181        checkpoint_id: Self::CheckpointId,
182        checkpoint: Checkpoint,
183    ) -> Result<(), Self::Error> {
184        add_checkpoint(self.conn, self.table_prefix, checkpoint_id, checkpoint)
185    }
186
187    fn checkpoint_count(&self) -> Result<usize, Self::Error> {
188        checkpoint_count(self.conn, self.table_prefix)
189    }
190
191    fn get_checkpoint_at_depth(
192        &self,
193        checkpoint_depth: usize,
194    ) -> Result<Option<(Self::CheckpointId, Checkpoint)>, Self::Error> {
195        get_checkpoint_at_depth(self.conn, self.table_prefix, checkpoint_depth)
196            .map_err(Error::Query)
197    }
198
199    fn get_checkpoint(
200        &self,
201        checkpoint_id: &Self::CheckpointId,
202    ) -> Result<Option<Checkpoint>, Self::Error> {
203        get_checkpoint(self.conn, self.table_prefix, *checkpoint_id)
204    }
205
206    fn with_checkpoints<F>(&mut self, limit: usize, callback: F) -> Result<(), Self::Error>
207    where
208        F: FnMut(&Self::CheckpointId, &Checkpoint) -> Result<(), Self::Error>,
209    {
210        with_checkpoints(self.conn, self.table_prefix, limit, callback)
211    }
212
213    fn for_each_checkpoint<F>(&self, limit: usize, callback: F) -> Result<(), Self::Error>
214    where
215        F: FnMut(&Self::CheckpointId, &Checkpoint) -> Result<(), Self::Error>,
216    {
217        with_checkpoints(self.conn, self.table_prefix, limit, callback)
218    }
219
220    fn update_checkpoint_with<F>(
221        &mut self,
222        checkpoint_id: &Self::CheckpointId,
223        update: F,
224    ) -> Result<bool, Self::Error>
225    where
226        F: Fn(&mut Checkpoint) -> Result<(), Self::Error>,
227    {
228        update_checkpoint_with(self.conn, self.table_prefix, *checkpoint_id, update)
229    }
230
231    fn remove_checkpoint(&mut self, checkpoint_id: &Self::CheckpointId) -> Result<(), Self::Error> {
232        remove_checkpoint(self.conn, self.table_prefix, *checkpoint_id)
233    }
234
235    fn add_retained_checkpoint(
236        &mut self,
237        checkpoint_id: Self::CheckpointId,
238    ) -> Result<(), Self::Error> {
239        add_retained_checkpoint(self.conn, self.table_prefix, checkpoint_id)
240    }
241
242    fn remove_retained_checkpoint(
243        &mut self,
244        checkpoint_id: &Self::CheckpointId,
245    ) -> Result<(), Self::Error> {
246        remove_retained_checkpoint(self.conn, self.table_prefix, *checkpoint_id)
247    }
248
249    fn retained_checkpoints(&self) -> Result<BTreeSet<Self::CheckpointId>, Self::Error> {
250        retained_checkpoints(self.conn, self.table_prefix)
251    }
252
253    fn truncate_checkpoints_retaining(
254        &mut self,
255        checkpoint_id: &Self::CheckpointId,
256    ) -> Result<(), Self::Error> {
257        truncate_checkpoints_retaining(self.conn, self.table_prefix, *checkpoint_id)
258    }
259}
260
261impl<H: HashSer, const SHARD_HEIGHT: u8> ShardStore
262    for SqliteShardStore<rusqlite::Connection, H, SHARD_HEIGHT>
263{
264    type H = H;
265    type CheckpointId = BlockHeight;
266    type Error = Error;
267
268    fn get_shard(
269        &self,
270        shard_root: Address,
271    ) -> Result<Option<LocatedPrunableTree<Self::H>>, Self::Error> {
272        get_shard(&self.conn, self.table_prefix, shard_root)
273    }
274
275    fn last_shard(&self) -> Result<Option<LocatedPrunableTree<Self::H>>, Self::Error> {
276        last_shard(&self.conn, self.table_prefix, Self::SHARD_ROOT_LEVEL)
277    }
278
279    fn put_shard(&mut self, subtree: LocatedPrunableTree<Self::H>) -> Result<(), Self::Error> {
280        let tx = self.conn.transaction().map_err(Error::Query)?;
281        put_shard(&tx, self.table_prefix, subtree)?;
282        tx.commit().map_err(Error::Query)?;
283        Ok(())
284    }
285
286    fn get_shard_roots(&self) -> Result<Vec<Address>, Self::Error> {
287        get_shard_roots(&self.conn, self.table_prefix, Self::SHARD_ROOT_LEVEL)
288    }
289
290    fn truncate_shards(&mut self, shard_index: u64) -> Result<(), Self::Error> {
291        truncate_shards(&self.conn, self.table_prefix, shard_index)
292    }
293
294    fn get_cap(&self) -> Result<PrunableTree<Self::H>, Self::Error> {
295        get_cap(&self.conn, self.table_prefix)
296    }
297
298    fn put_cap(&mut self, cap: PrunableTree<Self::H>) -> Result<(), Self::Error> {
299        put_cap(&self.conn, self.table_prefix, cap)
300    }
301
302    fn min_checkpoint_id(&self) -> Result<Option<Self::CheckpointId>, Self::Error> {
303        min_checkpoint_id(&self.conn, self.table_prefix)
304    }
305
306    fn max_checkpoint_id(&self) -> Result<Option<Self::CheckpointId>, Self::Error> {
307        max_checkpoint_id(&self.conn, self.table_prefix)
308    }
309
310    fn add_checkpoint(
311        &mut self,
312        checkpoint_id: Self::CheckpointId,
313        checkpoint: Checkpoint,
314    ) -> Result<(), Self::Error> {
315        let tx = self.conn.transaction().map_err(Error::Query)?;
316        add_checkpoint(&tx, self.table_prefix, checkpoint_id, checkpoint)?;
317        tx.commit().map_err(Error::Query)
318    }
319
320    fn checkpoint_count(&self) -> Result<usize, Self::Error> {
321        checkpoint_count(&self.conn, self.table_prefix)
322    }
323
324    fn get_checkpoint_at_depth(
325        &self,
326        checkpoint_depth: usize,
327    ) -> Result<Option<(Self::CheckpointId, Checkpoint)>, Self::Error> {
328        get_checkpoint_at_depth(&self.conn, self.table_prefix, checkpoint_depth)
329            .map_err(Error::Query)
330    }
331
332    fn get_checkpoint(
333        &self,
334        checkpoint_id: &Self::CheckpointId,
335    ) -> Result<Option<Checkpoint>, Self::Error> {
336        get_checkpoint(&self.conn, self.table_prefix, *checkpoint_id)
337    }
338
339    fn with_checkpoints<F>(&mut self, limit: usize, callback: F) -> Result<(), Self::Error>
340    where
341        F: FnMut(&Self::CheckpointId, &Checkpoint) -> Result<(), Self::Error>,
342    {
343        let tx = self.conn.transaction().map_err(Error::Query)?;
344        with_checkpoints(&tx, self.table_prefix, limit, callback)?;
345        tx.commit().map_err(Error::Query)
346    }
347
348    fn for_each_checkpoint<F>(&self, limit: usize, callback: F) -> Result<(), Self::Error>
349    where
350        F: FnMut(&Self::CheckpointId, &Checkpoint) -> Result<(), Self::Error>,
351    {
352        let tx = self.conn.unchecked_transaction().map_err(Error::Query)?;
353        with_checkpoints(&tx, self.table_prefix, limit, callback)?;
354        // Here, we use `tx.rollback` as the semantics of this method is that the callback must
355        // not mutate the data store.
356        tx.rollback().map_err(Error::Query)
357    }
358
359    fn update_checkpoint_with<F>(
360        &mut self,
361        checkpoint_id: &Self::CheckpointId,
362        update: F,
363    ) -> Result<bool, Self::Error>
364    where
365        F: Fn(&mut Checkpoint) -> Result<(), Self::Error>,
366    {
367        let tx = self.conn.transaction().map_err(Error::Query)?;
368        let result = update_checkpoint_with(&tx, self.table_prefix, *checkpoint_id, update)?;
369        tx.commit().map_err(Error::Query)?;
370        Ok(result)
371    }
372
373    fn remove_checkpoint(&mut self, checkpoint_id: &Self::CheckpointId) -> Result<(), Self::Error> {
374        let tx = self.conn.transaction().map_err(Error::Query)?;
375        remove_checkpoint(&tx, self.table_prefix, *checkpoint_id)?;
376        tx.commit().map_err(Error::Query)
377    }
378
379    fn add_retained_checkpoint(
380        &mut self,
381        checkpoint_id: Self::CheckpointId,
382    ) -> Result<(), Self::Error> {
383        let tx = self.conn.transaction().map_err(Error::Query)?;
384        add_retained_checkpoint(&tx, self.table_prefix, checkpoint_id)?;
385        tx.commit().map_err(Error::Query)
386    }
387
388    fn remove_retained_checkpoint(
389        &mut self,
390        checkpoint_id: &Self::CheckpointId,
391    ) -> Result<(), Self::Error> {
392        let tx = self.conn.transaction().map_err(Error::Query)?;
393        remove_retained_checkpoint(&tx, self.table_prefix, *checkpoint_id)?;
394        tx.commit().map_err(Error::Query)
395    }
396
397    fn retained_checkpoints(&self) -> Result<BTreeSet<Self::CheckpointId>, Self::Error> {
398        retained_checkpoints(&self.conn, self.table_prefix)
399    }
400
401    fn truncate_checkpoints_retaining(
402        &mut self,
403        checkpoint_id: &Self::CheckpointId,
404    ) -> Result<(), Self::Error> {
405        let tx = self.conn.transaction().map_err(Error::Query)?;
406        truncate_checkpoints_retaining(&tx, self.table_prefix, *checkpoint_id)?;
407        tx.commit().map_err(Error::Query)
408    }
409}
410
411/// Returns the stored root hash of the completed subtree with the given index, as most
412/// recently written by [`put_shard_roots`] or recorded by [`put_shard`], or `Ok(None)` if
413/// no row exists for the subtree or its `root_hash` is unknown.
414pub(crate) fn get_subtree_root<H: HashSer>(
415    conn: &rusqlite::Connection,
416    table_prefix: &'static str,
417    index: u64,
418) -> Result<Option<H>, Error> {
419    conn.query_row(
420        &format!(
421            "SELECT root_hash
422             FROM {table_prefix}_tree_shards
423             WHERE shard_index = :shard_index"
424        ),
425        named_params![":shard_index": index],
426        |row| row.get::<_, Option<Vec<u8>>>(0),
427    )
428    .optional()
429    .map_err(Error::Query)?
430    .flatten()
431    .map(|bytes| H::read(Cursor::new(bytes)).map_err(Error::Serialization))
432    .transpose()
433}
434
435pub(crate) fn get_shard<H: HashSer>(
436    conn: &rusqlite::Connection,
437    table_prefix: &'static str,
438    shard_root_addr: Address,
439) -> Result<Option<LocatedPrunableTree<H>>, Error> {
440    conn.query_row(
441        &format!(
442            "SELECT shard_data, root_hash
443             FROM {table_prefix}_tree_shards
444             WHERE shard_index = :shard_index"
445        ),
446        named_params![":shard_index": shard_root_addr.index()],
447        |row| Ok((row.get::<_, Vec<u8>>(0)?, row.get::<_, Option<Vec<u8>>>(1)?)),
448    )
449    .optional()
450    .map_err(Error::Query)?
451    .map(|(shard_data, root_hash)| {
452        let shard_tree = read_shard(&mut Cursor::new(shard_data)).map_err(Error::Serialization)?;
453        let located_tree =
454            LocatedPrunableTree::from_parts(shard_root_addr, shard_tree).map_err(|e| {
455                Error::Serialization(io::Error::new(
456                    io::ErrorKind::InvalidData,
457                    format!("Tree contained invalid data at address {e:?}"),
458                ))
459            })?;
460        if let Some(root_hash_data) = root_hash {
461            let root_hash = H::read(Cursor::new(root_hash_data)).map_err(Error::Serialization)?;
462            Ok(located_tree.reannotate_root(Some(Arc::new(root_hash))))
463        } else {
464            Ok(located_tree)
465        }
466    })
467    .transpose()
468}
469
470pub(crate) fn last_shard<H: HashSer>(
471    conn: &rusqlite::Connection,
472    table_prefix: &'static str,
473    shard_root_level: Level,
474) -> Result<Option<LocatedPrunableTree<H>>, Error> {
475    conn.query_row(
476        &format!(
477            "SELECT shard_index, shard_data
478             FROM {table_prefix}_tree_shards
479             ORDER BY shard_index DESC
480             LIMIT 1"
481        ),
482        [],
483        |row| {
484            let shard_index: u64 = row.get(0)?;
485            let shard_data: Vec<u8> = row.get(1)?;
486            Ok((shard_index, shard_data))
487        },
488    )
489    .optional()
490    .map_err(Error::Query)?
491    .map(|(shard_index, shard_data)| {
492        let shard_root = Address::from_parts(shard_root_level, shard_index);
493        let shard_tree = read_shard(&mut Cursor::new(shard_data)).map_err(Error::Serialization)?;
494        LocatedPrunableTree::from_parts(shard_root, shard_tree).map_err(|e| {
495            Error::Serialization(io::Error::new(
496                io::ErrorKind::InvalidData,
497                format!("Tree contained invalid data at address {e:?}"),
498            ))
499        })
500    })
501    .transpose()
502}
503
504/// Returns an error iff the proposed insertion range
505/// for the tree shards would create a discontinuity
506/// in the database.
507#[tracing::instrument(skip(conn))]
508fn check_shard_discontinuity(
509    conn: &rusqlite::Connection,
510    table_prefix: &'static str,
511    proposed_insertion_range: Range<u64>,
512) -> Result<(), Error> {
513    if let Ok((Some(stored_min), Some(stored_max))) = conn
514        .query_row(
515            &format!("SELECT MIN(shard_index), MAX(shard_index) FROM {table_prefix}_tree_shards"),
516            [],
517            |row| {
518                let min = row.get::<_, Option<u64>>(0)?;
519                let max = row.get::<_, Option<u64>>(1)?;
520                Ok((min, max))
521            },
522        )
523        .map_err(Error::Query)
524    {
525        // If the ranges overlap, or are directly adjacent, then we aren't creating a
526        // discontinuity. We can check this by comparing their start-inclusive,
527        // end-exclusive bounds:
528        // - If `cur_start == ins_end` then the proposed insertion range is immediately
529        //   before the current shards. If `cur_start > ins_end` then there is a gap.
530        // - If `ins_start == cur_end` then the proposed insertion range is immediately
531        //   after the current shards. If `ins_start > cur_end` then there is a gap.
532        let (cur_start, cur_end) = (stored_min, stored_max + 1);
533        let (ins_start, ins_end) = (proposed_insertion_range.start, proposed_insertion_range.end);
534        if cur_start > ins_end || ins_start > cur_end {
535            return Err(Error::SubtreeDiscontinuity {
536                attempted_insertion_range: proposed_insertion_range,
537                existing_range: cur_start..cur_end,
538            });
539        }
540    }
541
542    Ok(())
543}
544
545pub(crate) fn put_shard<H: HashSer>(
546    conn: &rusqlite::Transaction<'_>,
547    table_prefix: &'static str,
548    subtree: LocatedPrunableTree<H>,
549) -> Result<(), Error> {
550    let subtree_root_hash = subtree
551        .root()
552        .annotation()
553        .and_then(|ann| {
554            ann.as_ref().map(|rc| {
555                let mut root_hash = vec![];
556                rc.write(&mut root_hash)?;
557                Ok(root_hash)
558            })
559        })
560        .transpose()
561        .map_err(Error::Serialization)?;
562
563    let mut subtree_data = vec![];
564    write_shard(&mut subtree_data, subtree.root()).map_err(Error::Serialization)?;
565
566    let shard_index = subtree.root_addr().index();
567
568    check_shard_discontinuity(conn, table_prefix, shard_index..shard_index + 1)?;
569
570    let mut stmt_put_shard = conn
571        .prepare_cached(&format!(
572            "INSERT INTO {table_prefix}_tree_shards (shard_index, root_hash, shard_data)
573             VALUES (:shard_index, :root_hash, :shard_data)
574             ON CONFLICT (shard_index) DO UPDATE
575             SET root_hash = :root_hash,
576             shard_data = :shard_data"
577        ))
578        .map_err(Error::Query)?;
579
580    stmt_put_shard
581        .execute(named_params![
582            ":shard_index": shard_index,
583            ":root_hash": subtree_root_hash,
584            ":shard_data": subtree_data
585        ])
586        .map_err(Error::Query)?;
587
588    Ok(())
589}
590
591pub(crate) fn get_shard_roots(
592    conn: &rusqlite::Connection,
593    table_prefix: &'static str,
594    shard_root_level: Level,
595) -> Result<Vec<Address>, Error> {
596    let mut stmt = conn
597        .prepare(&format!(
598            "SELECT shard_index FROM {table_prefix}_tree_shards ORDER BY shard_index"
599        ))
600        .map_err(Error::Query)?;
601    let mut rows = stmt.query([]).map_err(Error::Query)?;
602
603    let mut res = vec![];
604    while let Some(row) = rows.next().map_err(Error::Query)? {
605        res.push(Address::from_parts(
606            shard_root_level,
607            row.get(0).map_err(Error::Query)?,
608        ));
609    }
610    Ok(res)
611}
612
613pub(crate) fn truncate_shards(
614    conn: &rusqlite::Connection,
615    table_prefix: &'static str,
616    shard_index: u64,
617) -> Result<(), Error> {
618    conn.execute(
619        &format!("DELETE FROM {table_prefix}_tree_shards WHERE shard_index >= ?"),
620        [shard_index],
621    )
622    .map_err(Error::Query)
623    .map(|_| ())
624}
625
626#[tracing::instrument(skip(conn))]
627pub(crate) fn get_cap<H: HashSer>(
628    conn: &rusqlite::Connection,
629    table_prefix: &'static str,
630) -> Result<PrunableTree<H>, Error> {
631    conn.query_row(
632        &format!("SELECT cap_data FROM {table_prefix}_tree_cap"),
633        [],
634        |row| row.get::<_, Vec<u8>>(0),
635    )
636    .optional()
637    .map_err(Error::Query)?
638    .map_or_else(
639        || Ok(PrunableTree::empty()),
640        |cap_data| read_shard(&mut Cursor::new(cap_data)).map_err(Error::Serialization),
641    )
642}
643
644#[tracing::instrument(skip(conn, cap))]
645pub(crate) fn put_cap<H: HashSer>(
646    conn: &rusqlite::Connection,
647    table_prefix: &'static str,
648    cap: PrunableTree<H>,
649) -> Result<(), Error> {
650    let mut stmt = conn
651        .prepare_cached(&format!(
652            "INSERT INTO {table_prefix}_tree_cap (cap_id, cap_data)
653             VALUES (0, :cap_data)
654             ON CONFLICT (cap_id) DO UPDATE
655             SET cap_data = :cap_data"
656        ))
657        .map_err(Error::Query)?;
658
659    let mut cap_data = vec![];
660    write_shard(&mut cap_data, &cap).map_err(Error::Serialization)?;
661    stmt.execute([cap_data]).map_err(Error::Query)?;
662
663    Ok(())
664}
665
666pub(crate) fn min_checkpoint_id(
667    conn: &rusqlite::Connection,
668    table_prefix: &'static str,
669) -> Result<Option<BlockHeight>, Error> {
670    conn.query_row(
671        &format!("SELECT MIN(checkpoint_id) FROM {table_prefix}_tree_checkpoints"),
672        [],
673        |row| {
674            row.get::<_, Option<u32>>(0)
675                .map(|opt| opt.map(BlockHeight::from))
676        },
677    )
678    .map_err(Error::Query)
679}
680
681pub(crate) fn max_checkpoint_id(
682    conn: &rusqlite::Connection,
683    table_prefix: &'static str,
684) -> Result<Option<BlockHeight>, Error> {
685    conn.query_row(
686        &format!("SELECT MAX(checkpoint_id) FROM {table_prefix}_tree_checkpoints"),
687        [],
688        |row| {
689            row.get::<_, Option<u32>>(0)
690                .map(|opt| opt.map(BlockHeight::from))
691        },
692    )
693    .map_err(Error::Query)
694}
695
696/// Returns the lowest retained checkpoint id that is at or above `floor`, or `None`
697/// if the pool's checkpoint table contains no checkpoint within `[floor, ∞)`.
698pub(crate) fn min_checkpoint_id_at_or_above(
699    conn: &rusqlite::Connection,
700    table_prefix: &'static str,
701    floor: BlockHeight,
702) -> Result<Option<BlockHeight>, Error> {
703    conn.query_row(
704        &format!(
705            "SELECT MIN(checkpoint_id) FROM {table_prefix}_tree_checkpoints
706             WHERE checkpoint_id >= :floor"
707        ),
708        named_params![":floor": u32::from(floor)],
709        |row| {
710            row.get::<_, Option<u32>>(0)
711                .map(|opt| opt.map(BlockHeight::from))
712        },
713    )
714    .map_err(Error::Query)
715}
716
717/// Resets the note commitment tree with the given table prefix to contain only the roots of
718/// subtrees completed at or below `truncation_height`, discarding all of its other state:
719/// scanned shard contents, the cap, and all checkpoints.
720///
721/// This is the truncation outcome for a tree whose retained checkpoints all lie above the
722/// truncation height: no checkpoint exists at or below that height for
723/// `ShardTree::truncate_to_checkpoint` to target, so the tree's scanned contents postdate the
724/// truncation point and must be discarded; the rescan of heights above that point re-creates
725/// them. The roots of subtrees completed at or below the truncation height remain facts
726/// about the retained portion of the chain, however, so they are preserved (in the form
727/// fast sync would deliver them, via [`put_shard_roots`]); discarding them would leave the
728/// wallet unable to construct witnesses spanning those subtrees until they had been
729/// re-downloaded.
730pub(crate) fn truncate_tree_to_subtree_roots<
731    H: Hashable + HashSer + Clone + Eq,
732    const DEPTH: u8,
733    const SHARD_HEIGHT: u8,
734>(
735    conn: &rusqlite::Transaction<'_>,
736    table_prefix: &'static str,
737    truncation_height: BlockHeight,
738) -> Result<(), ShardTreeError<Error>> {
739    // Collect the roots of subtrees completed at or below the truncation height. Subtree end
740    // heights increase with shard index, so these form a prefix of the shard sequence;
741    // re-insertion via `put_shard_roots` requires contiguity from index zero, so collection
742    // stops at the first shard for which no completed root is recorded.
743    let roots = {
744        let mut stmt = conn
745            .prepare(&format!(
746                "SELECT shard_index, subtree_end_height, root_hash
747                 FROM {table_prefix}_tree_shards
748                 WHERE subtree_end_height IS NOT NULL
749                 AND subtree_end_height <= :truncation_height
750                 AND root_hash IS NOT NULL
751                 ORDER BY shard_index"
752            ))
753            .map_err(|e| ShardTreeError::Storage(Error::Query(e)))?;
754
755        let rows = stmt
756            .query_map(
757                named_params![":truncation_height": u32::from(truncation_height)],
758                |row| {
759                    Ok((
760                        row.get::<_, u64>(0)?,
761                        row.get::<_, u32>(1)?,
762                        row.get::<_, Vec<u8>>(2)?,
763                    ))
764                },
765            )
766            .map_err(|e| ShardTreeError::Storage(Error::Query(e)))?;
767
768        let mut roots = vec![];
769        for row in rows {
770            let (shard_index, subtree_end_height, root_hash) =
771                row.map_err(|e| ShardTreeError::Storage(Error::Query(e)))?;
772            if shard_index != u64::try_from(roots.len()).expect("vec length fits in u64") {
773                break;
774            }
775            let root = H::read(Cursor::new(root_hash))
776                .map_err(|e| ShardTreeError::Storage(Error::Serialization(e)))?;
777            roots.push(CommitmentTreeRoot::from_parts(
778                BlockHeight::from(subtree_end_height),
779                root,
780            ));
781        }
782        roots
783    };
784
785    conn.execute_batch(&format!(
786        "DELETE FROM {table_prefix}_tree_checkpoint_marks_removed;
787         DELETE FROM {table_prefix}_tree_checkpoints;
788         DELETE FROM {table_prefix}_tree_shards;
789         DELETE FROM {table_prefix}_tree_cap;"
790    ))
791    .map_err(|e| ShardTreeError::Storage(Error::Query(e)))?;
792
793    put_shard_roots::<H, DEPTH, SHARD_HEIGHT>(conn, table_prefix, 0, &roots)
794}
795
796pub(crate) fn add_checkpoint(
797    conn: &rusqlite::Transaction<'_>,
798    table_prefix: &'static str,
799    checkpoint_id: BlockHeight,
800    checkpoint: Checkpoint,
801) -> Result<(), Error> {
802    let extant_tree_state = conn
803        .query_row(
804            &format!(
805                "SELECT position FROM {table_prefix}_tree_checkpoints WHERE checkpoint_id = :checkpoint_id"
806            ),
807            named_params![":checkpoint_id": u32::from(checkpoint_id),],
808            |row| {
809                row.get::<_, Option<u64>>(0).map(|opt| {
810                    opt.map_or_else(
811                        || TreeState::Empty,
812                        |pos| TreeState::AtPosition(Position::from(pos)),
813                    )
814                })
815            },
816        )
817        .optional()
818        .map_err(Error::Query)?;
819
820    match extant_tree_state {
821        Some(current) => {
822            if current != checkpoint.tree_state() {
823                // If the checkpoint position for a given checkpoint identifier has changed, we treat
824                // this as an error because the wallet should have detected a chain reorg and truncated
825                // the tree.
826                Err(Error::CheckpointConflict {
827                    checkpoint_id,
828                    checkpoint,
829                    extant_tree_state: current,
830                    extant_marks_removed: None,
831                })
832            } else {
833                // if the existing spends are the same, we can skip the insert; if the
834                // existing spends have changed, this is also a conflict.
835                let marks_removed = get_marks_removed(conn, table_prefix, checkpoint_id)?;
836                if &marks_removed == checkpoint.marks_removed() {
837                    Ok(())
838                } else {
839                    Err(Error::CheckpointConflict {
840                        checkpoint_id,
841                        checkpoint,
842                        extant_tree_state: current,
843                        extant_marks_removed: Some(marks_removed),
844                    })
845                }
846            }
847        }
848        None => {
849            let mut stmt_insert_checkpoint = conn
850                .prepare_cached(&format!(
851                    "INSERT INTO {table_prefix}_tree_checkpoints (checkpoint_id, position)
852                     VALUES (:checkpoint_id, :position)"
853                ))
854                .map_err(Error::Query)?;
855
856            stmt_insert_checkpoint
857                .execute(named_params![
858                    ":checkpoint_id": u32::from(checkpoint_id),
859                    ":position": checkpoint.position().map(u64::from)
860                ])
861                .map_err(Error::Query)?;
862
863            let mut stmt_insert_mark_removed = conn
864                .prepare_cached(&format!(
865                    "INSERT INTO {table_prefix}_tree_checkpoint_marks_removed (checkpoint_id, mark_removed_position)
866                     VALUES (:checkpoint_id, :position)"
867                ))
868                .map_err(Error::Query)?;
869
870            for pos in checkpoint.marks_removed() {
871                stmt_insert_mark_removed
872                    .execute(named_params![
873                        ":checkpoint_id": u32::from(checkpoint_id),
874                        ":position": u64::from(*pos)
875                    ])
876                    .map_err(Error::Query)?;
877            }
878
879            Ok(())
880        }
881    }
882}
883
884pub(crate) fn checkpoint_count(
885    conn: &rusqlite::Connection,
886    table_prefix: &'static str,
887) -> Result<usize, Error> {
888    conn.query_row(
889        &format!("SELECT COUNT(*) FROM {table_prefix}_tree_checkpoints"),
890        [],
891        |row| row.get::<_, usize>(0),
892    )
893    .map_err(Error::Query)
894}
895
896fn get_marks_removed(
897    conn: &rusqlite::Connection,
898    table_prefix: &'static str,
899    checkpoint_id: BlockHeight,
900) -> Result<BTreeSet<Position>, Error> {
901    let mut stmt = conn
902        .prepare_cached(&format!(
903            "SELECT mark_removed_position
904            FROM {table_prefix}_tree_checkpoint_marks_removed
905            WHERE checkpoint_id = ?"
906        ))
907        .map_err(Error::Query)?;
908    let mark_removed_rows = stmt
909        .query([u32::from(checkpoint_id)])
910        .map_err(Error::Query)?;
911
912    mark_removed_rows
913        .mapped(|row| row.get::<_, u64>(0).map(Position::from))
914        .collect::<Result<BTreeSet<_>, _>>()
915        .map_err(Error::Query)
916}
917
918pub(crate) fn get_checkpoint(
919    conn: &rusqlite::Connection,
920    table_prefix: &'static str,
921    checkpoint_id: BlockHeight,
922) -> Result<Option<Checkpoint>, Error> {
923    let checkpoint_position = conn
924        .query_row(
925            &format!(
926                "SELECT position
927                 FROM {table_prefix}_tree_checkpoints
928                 WHERE checkpoint_id = ?"
929            ),
930            [u32::from(checkpoint_id)],
931            |row| {
932                row.get::<_, Option<u64>>(0)
933                    .map(|opt| opt.map(Position::from))
934            },
935        )
936        .optional()
937        .map_err(Error::Query)?;
938
939    checkpoint_position
940        .map(|pos_opt| {
941            Ok(Checkpoint::from_parts(
942                pos_opt.map_or(TreeState::Empty, TreeState::AtPosition),
943                get_marks_removed(conn, table_prefix, checkpoint_id)?,
944            ))
945        })
946        .transpose()
947}
948
949pub(crate) fn get_max_checkpointed_height(
950    conn: &rusqlite::Connection,
951    protocol: ShieldedPool,
952    target_height: TargetHeight,
953    min_confirmations: NonZeroU32,
954) -> Result<Option<BlockHeight>, SqliteClientError> {
955    let TableConstants { table_prefix, .. } = table_constants::<SqliteClientError>(protocol)?;
956    let max_checkpoint_height = target_height - u32::from(min_confirmations);
957
958    // We exclude from consideration all checkpoints having heights greater than the maximum
959    // checkpoint height. The checkpoint depth is the number of excluded checkpoints + 1.
960    conn.query_row(
961        &format!(
962            "SELECT checkpoint_id
963             FROM {table_prefix}_tree_checkpoints
964             WHERE checkpoint_id <= :max_checkpoint_height
965             ORDER BY checkpoint_id DESC
966             LIMIT 1",
967        ),
968        named_params![":max_checkpoint_height": u32::from(max_checkpoint_height)],
969        |row| row.get::<_, u32>(0).map(BlockHeight::from),
970    )
971    .optional()
972    .map_err(SqliteClientError::from)
973}
974
975pub(crate) fn get_checkpoint_at_depth(
976    conn: &rusqlite::Connection,
977    table_prefix: &'static str,
978    checkpoint_depth: usize,
979) -> Result<Option<(BlockHeight, Checkpoint)>, rusqlite::Error> {
980    let checkpoint_parts = conn
981        .query_row(
982            &format!(
983                "SELECT checkpoint_id, position
984                FROM {table_prefix}_tree_checkpoints
985                ORDER BY checkpoint_id DESC
986                LIMIT 1
987                OFFSET :offset",
988            ),
989            named_params![":offset": checkpoint_depth],
990            |row| {
991                let checkpoint_id: u32 = row.get(0)?;
992                let position: Option<u64> = row.get(1)?;
993                Ok((
994                    BlockHeight::from(checkpoint_id),
995                    position.map(Position::from),
996                ))
997            },
998        )
999        .optional()?;
1000
1001    checkpoint_parts
1002        .map(|(checkpoint_id, pos_opt)| {
1003            let mut stmt = conn.prepare_cached(&format!(
1004                "SELECT mark_removed_position
1005                    FROM {table_prefix}_tree_checkpoint_marks_removed
1006                    WHERE checkpoint_id = ?"
1007            ))?;
1008            let mark_removed_rows = stmt.query([u32::from(checkpoint_id)])?;
1009
1010            let marks_removed = mark_removed_rows
1011                .mapped(|row| row.get::<_, u64>(0).map(Position::from))
1012                .collect::<Result<BTreeSet<_>, _>>()?;
1013
1014            Ok((
1015                checkpoint_id,
1016                Checkpoint::from_parts(
1017                    pos_opt.map_or(TreeState::Empty, TreeState::AtPosition),
1018                    marks_removed,
1019                ),
1020            ))
1021        })
1022        .transpose()
1023}
1024
1025pub(crate) fn with_checkpoints<F>(
1026    conn: &rusqlite::Transaction<'_>,
1027    table_prefix: &'static str,
1028    limit: usize,
1029    mut callback: F,
1030) -> Result<(), Error>
1031where
1032    F: FnMut(&BlockHeight, &Checkpoint) -> Result<(), Error>,
1033{
1034    let mut stmt_get_checkpoints = conn
1035        .prepare_cached(&format!(
1036            "SELECT checkpoint_id, position
1037            FROM {table_prefix}_tree_checkpoints
1038            ORDER BY position
1039            LIMIT :limit"
1040        ))
1041        .map_err(Error::Query)?;
1042
1043    let mut stmt_get_checkpoint_marks_removed = conn
1044        .prepare_cached(&format!(
1045            "SELECT mark_removed_position
1046            FROM {table_prefix}_tree_checkpoint_marks_removed
1047            WHERE checkpoint_id = :checkpoint_id"
1048        ))
1049        .map_err(Error::Query)?;
1050
1051    let mut rows = stmt_get_checkpoints
1052        .query(named_params![":limit": limit])
1053        .map_err(Error::Query)?;
1054
1055    while let Some(row) = rows.next().map_err(Error::Query)? {
1056        let checkpoint_id = row.get::<_, u32>(0).map_err(Error::Query)?;
1057        let tree_state = row
1058            .get::<_, Option<u64>>(1)
1059            .map(|opt| opt.map_or_else(|| TreeState::Empty, |p| TreeState::AtPosition(p.into())))
1060            .map_err(Error::Query)?;
1061
1062        let mark_removed_rows = stmt_get_checkpoint_marks_removed
1063            .query(named_params![":checkpoint_id": checkpoint_id])
1064            .map_err(Error::Query)?;
1065
1066        let marks_removed = mark_removed_rows
1067            .mapped(|row| row.get::<_, u64>(0).map(Position::from))
1068            .collect::<Result<BTreeSet<_>, _>>()
1069            .map_err(Error::Query)?;
1070
1071        callback(
1072            &BlockHeight::from(checkpoint_id),
1073            &Checkpoint::from_parts(tree_state, marks_removed),
1074        )?
1075    }
1076
1077    Ok(())
1078}
1079
1080pub(crate) fn update_checkpoint_with<F>(
1081    conn: &rusqlite::Transaction<'_>,
1082    table_prefix: &'static str,
1083    checkpoint_id: BlockHeight,
1084    update: F,
1085) -> Result<bool, Error>
1086where
1087    F: Fn(&mut Checkpoint) -> Result<(), Error>,
1088{
1089    if let Some(mut c) = get_checkpoint(conn, table_prefix, checkpoint_id)? {
1090        update(&mut c)?;
1091        remove_checkpoint(conn, table_prefix, checkpoint_id)?;
1092        add_checkpoint(conn, table_prefix, checkpoint_id, c)?;
1093        Ok(true)
1094    } else {
1095        Ok(false)
1096    }
1097}
1098
1099pub(crate) fn remove_checkpoint(
1100    conn: &rusqlite::Transaction<'_>,
1101    table_prefix: &'static str,
1102    checkpoint_id: BlockHeight,
1103) -> Result<(), Error> {
1104    // cascading delete here obviates the need to manually delete from
1105    // `tree_checkpoint_marks_removed`
1106    let mut stmt_delete_checkpoint = conn
1107        .prepare_cached(&format!(
1108            "DELETE FROM {table_prefix}_tree_checkpoints
1109             WHERE checkpoint_id = :checkpoint_id"
1110        ))
1111        .map_err(Error::Query)?;
1112
1113    stmt_delete_checkpoint
1114        .execute(named_params![":checkpoint_id": u32::from(checkpoint_id),])
1115        .map_err(Error::Query)?;
1116
1117    Ok(())
1118}
1119
1120pub(crate) fn add_retained_checkpoint(
1121    conn: &rusqlite::Transaction<'_>,
1122    table_prefix: &'static str,
1123    checkpoint_id: BlockHeight,
1124) -> Result<(), Error> {
1125    conn.prepare_cached(&format!(
1126        "INSERT OR IGNORE INTO {table_prefix}_tree_retained_checkpoints (checkpoint_id)
1127         VALUES (:checkpoint_id)"
1128    ))
1129    .map_err(Error::Query)?
1130    .execute(named_params![":checkpoint_id": u32::from(checkpoint_id)])
1131    .map_err(Error::Query)?;
1132
1133    Ok(())
1134}
1135
1136pub(crate) fn remove_retained_checkpoint(
1137    conn: &rusqlite::Transaction<'_>,
1138    table_prefix: &'static str,
1139    checkpoint_id: BlockHeight,
1140) -> Result<(), Error> {
1141    conn.prepare_cached(&format!(
1142        "DELETE FROM {table_prefix}_tree_retained_checkpoints
1143         WHERE checkpoint_id = :checkpoint_id"
1144    ))
1145    .map_err(Error::Query)?
1146    .execute(named_params![":checkpoint_id": u32::from(checkpoint_id)])
1147    .map_err(Error::Query)?;
1148
1149    Ok(())
1150}
1151
1152pub(crate) fn retained_checkpoints(
1153    conn: &rusqlite::Connection,
1154    table_prefix: &'static str,
1155) -> Result<BTreeSet<BlockHeight>, Error> {
1156    // The retained-checkpoints table is created by a dedicated migration. Tree operations may run
1157    // against a schema that predates that migration (e.g. when a migration test drives the tree at
1158    // an intermediate state); such a wallet simply has no retained checkpoints, so report an empty
1159    // set rather than failing on the missing table.
1160    let table_name = format!("{table_prefix}_tree_retained_checkpoints");
1161    let table_exists = conn
1162        .query_row(
1163            "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = :table_name",
1164            named_params![":table_name": table_name],
1165            |_| Ok(()),
1166        )
1167        .optional()
1168        .map_err(Error::Query)?
1169        .is_some();
1170    if !table_exists {
1171        return Ok(BTreeSet::new());
1172    }
1173
1174    let mut stmt = conn
1175        .prepare_cached(&format!("SELECT checkpoint_id FROM {table_name}"))
1176        .map_err(Error::Query)?;
1177    let rows = stmt.query([]).map_err(Error::Query)?;
1178
1179    rows.mapped(|row| row.get::<_, u32>(0).map(BlockHeight::from))
1180        .collect::<Result<BTreeSet<_>, _>>()
1181        .map_err(Error::Query)
1182}
1183
1184pub(crate) fn truncate_checkpoints_retaining(
1185    conn: &rusqlite::Transaction<'_>,
1186    table_prefix: &'static str,
1187    checkpoint_id: BlockHeight,
1188) -> Result<(), Error> {
1189    // cascading delete here obviates the need to manually delete from
1190    // `<protocol>_tree_checkpoint_marks_removed`
1191    conn.execute(
1192        &format!("DELETE FROM {table_prefix}_tree_checkpoints WHERE checkpoint_id > ?"),
1193        [u32::from(checkpoint_id)],
1194    )
1195    .map_err(Error::Query)?;
1196
1197    // we do however need to manually delete any marks associated with the retained checkpoint
1198    conn.execute(
1199        &format!(
1200            "DELETE FROM {table_prefix}_tree_checkpoint_marks_removed WHERE checkpoint_id = ?"
1201        ),
1202        [u32::from(checkpoint_id)],
1203    )
1204    .map_err(Error::Query)?;
1205
1206    Ok(())
1207}
1208
1209#[tracing::instrument(skip(conn, roots))]
1210pub(crate) fn put_shard_roots<
1211    H: Hashable + HashSer + Clone + Eq,
1212    const DEPTH: u8,
1213    const SHARD_HEIGHT: u8,
1214>(
1215    conn: &rusqlite::Transaction<'_>,
1216    table_prefix: &'static str,
1217    start_index: u64,
1218    roots: &[CommitmentTreeRoot<H>],
1219) -> Result<(), ShardTreeError<Error>> {
1220    if roots.is_empty() {
1221        // nothing to do
1222        return Ok(());
1223    }
1224
1225    // We treat the cap as a tree with `DEPTH - SHARD_HEIGHT` levels, so that we can make a
1226    // batch insertion of root data using `Position::from(start_index)` as the starting position
1227    // and treating the roots as level-0 leaves.
1228    #[derive(Clone, Debug, PartialEq, Eq)]
1229    struct LevelShifter<H, const SHARD_HEIGHT: u8>(H);
1230    impl<H: Hashable, const SHARD_HEIGHT: u8> Hashable for LevelShifter<H, SHARD_HEIGHT> {
1231        fn empty_leaf() -> Self {
1232            Self(H::empty_root(SHARD_HEIGHT.into()))
1233        }
1234
1235        fn combine(level: Level, a: &Self, b: &Self) -> Self {
1236            Self(H::combine(level + SHARD_HEIGHT, &a.0, &b.0))
1237        }
1238
1239        fn empty_root(level: Level) -> Self
1240        where
1241            Self: Sized,
1242        {
1243            Self(H::empty_root(level + SHARD_HEIGHT))
1244        }
1245    }
1246    impl<H: HashSer, const SHARD_HEIGHT: u8> HashSer for LevelShifter<H, SHARD_HEIGHT> {
1247        fn read<R: io::Read>(reader: R) -> io::Result<Self>
1248        where
1249            Self: Sized,
1250        {
1251            H::read(reader).map(Self)
1252        }
1253
1254        fn write<W: io::Write>(&self, writer: W) -> io::Result<()> {
1255            self.0.write(writer)
1256        }
1257    }
1258
1259    let cap = LocatedTree::from_parts(
1260        Address::from_parts((DEPTH - SHARD_HEIGHT).into(), 0),
1261        get_cap::<LevelShifter<H, SHARD_HEIGHT>>(conn, table_prefix)
1262            .map_err(ShardTreeError::Storage)?,
1263    )
1264    .map_err(|e| {
1265        ShardTreeError::Storage(Error::Serialization(io::Error::new(
1266            io::ErrorKind::InvalidData,
1267            format!("Note commitment tree cap was invalid at address {e:?}"),
1268        )))
1269    })?;
1270
1271    let insert_into_cap = tracing::info_span!("insert_into_cap").entered();
1272    let cap_result = cap
1273        .batch_insert::<(), _>(
1274            Position::from(start_index),
1275            roots
1276                .iter()
1277                .map(|r| (LevelShifter(r.root_hash().clone()), Retention::Reference)),
1278        )
1279        .map_err(ShardTreeError::Insert)?
1280        .expect("slice of inserted roots was verified to be nonempty");
1281    drop(insert_into_cap);
1282
1283    put_cap(conn, table_prefix, cap_result.subtree.take_root()).map_err(ShardTreeError::Storage)?;
1284
1285    check_shard_discontinuity(
1286        conn,
1287        table_prefix,
1288        start_index..start_index + (roots.len() as u64),
1289    )
1290    .map_err(ShardTreeError::Storage)?;
1291
1292    // We want to avoid deserializing the subtree just to annotate its root node, so we simply
1293    // cache the downloaded root alongside of any already-persisted subtree. We will update the
1294    // subtree data itself by reannotating the root node of the tree, handling conflicts, at
1295    // the time that we deserialize the tree.
1296    let mut stmt = conn
1297        .prepare_cached(&format!(
1298            "INSERT INTO {table_prefix}_tree_shards (shard_index, subtree_end_height, root_hash, shard_data)
1299            VALUES (:shard_index, :subtree_end_height, :root_hash, :shard_data)
1300            ON CONFLICT (shard_index) DO UPDATE
1301            SET subtree_end_height = :subtree_end_height, root_hash = :root_hash"
1302        ))
1303        .map_err(|e| ShardTreeError::Storage(Error::Query(e)))?;
1304
1305    let put_roots = tracing::info_span!("write_shards").entered();
1306    for (root, i) in roots.iter().zip(0u64..) {
1307        // The `shard_data` value will only be used in the case that no tree already exists.
1308        let mut shard_data: Vec<u8> = vec![];
1309        let tree = PrunableTree::leaf((root.root_hash().clone(), RetentionFlags::EPHEMERAL));
1310        write_shard(&mut shard_data, &tree)
1311            .map_err(|e| ShardTreeError::Storage(Error::Serialization(e)))?;
1312
1313        let mut root_hash_data: Vec<u8> = vec![];
1314        root.root_hash()
1315            .write(&mut root_hash_data)
1316            .map_err(|e| ShardTreeError::Storage(Error::Serialization(e)))?;
1317
1318        stmt.execute(named_params![
1319            ":shard_index": start_index + i,
1320            ":subtree_end_height": u32::from(root.subtree_end_height()),
1321            ":root_hash": root_hash_data,
1322            ":shard_data": shard_data,
1323        ])
1324        .map_err(|e| ShardTreeError::Storage(Error::Query(e)))?;
1325    }
1326    drop(put_roots);
1327
1328    Ok(())
1329}
1330
1331pub(crate) fn check_witnesses(
1332    conn: &rusqlite::Transaction<'_>,
1333    anchor_height: BlockHeight,
1334) -> Result<Vec<Range<BlockHeight>>, SqliteClientError> {
1335    let wallet_birthday = super::wallet_birthday(conn)?.ok_or(SqliteClientError::AccountUnknown)?;
1336    let unspent_sapling_note_meta =
1337        super::sapling::select_unspent_note_meta(conn, wallet_birthday, anchor_height)?;
1338
1339    let mut scan_ranges = vec![];
1340    let mut sapling_incomplete = vec![];
1341    let sapling_tree = sapling_tree(conn)?;
1342    for m in unspent_sapling_note_meta.iter() {
1343        match sapling_tree.witness_at_checkpoint_depth(m.commitment_tree_position(), 0) {
1344            Ok(_) => {}
1345            Err(ShardTreeError::Query(QueryError::TreeIncomplete(mut addrs))) => {
1346                sapling_incomplete.append(&mut addrs);
1347            }
1348            Err(other) => {
1349                return Err(SqliteClientError::CommitmentTree(other));
1350            }
1351        }
1352    }
1353
1354    for addr in sapling_incomplete {
1355        let range = super::get_block_range(conn, ShieldedPool::Sapling, addr)?;
1356        scan_ranges.extend(range);
1357    }
1358
1359    #[cfg(feature = "orchard")]
1360    {
1361        let unspent_orchard_note_meta =
1362            super::orchard::select_unspent_note_meta(conn, wallet_birthday, anchor_height)?;
1363        let mut orchard_incomplete = vec![];
1364        let orchard_tree = orchard_tree(conn)?;
1365        for m in unspent_orchard_note_meta.iter() {
1366            match orchard_tree.witness_at_checkpoint_depth(m.commitment_tree_position(), 0) {
1367                Ok(_) => {}
1368                Err(ShardTreeError::Query(QueryError::TreeIncomplete(mut addrs))) => {
1369                    orchard_incomplete.append(&mut addrs);
1370                }
1371                Err(other) => {
1372                    return Err(SqliteClientError::CommitmentTree(other));
1373                }
1374            }
1375        }
1376
1377        for addr in orchard_incomplete {
1378            let range = super::get_block_range(conn, ShieldedPool::Orchard, addr)?;
1379            scan_ranges.extend(range);
1380        }
1381
1382        let unspent_ironwood_note_meta = super::common::select_unspent_note_meta(
1383            conn,
1384            ShieldedPool::Ironwood,
1385            wallet_birthday,
1386            anchor_height,
1387        )?;
1388        let mut ironwood_incomplete = vec![];
1389        let ironwood_tree = ironwood_tree(conn)?;
1390        for m in unspent_ironwood_note_meta.iter() {
1391            match ironwood_tree.witness_at_checkpoint_depth(m.commitment_tree_position(), 0) {
1392                Ok(_) => {}
1393                Err(ShardTreeError::Query(QueryError::TreeIncomplete(mut addrs))) => {
1394                    ironwood_incomplete.append(&mut addrs);
1395                }
1396                Err(other) => {
1397                    return Err(SqliteClientError::CommitmentTree(other));
1398                }
1399            }
1400        }
1401
1402        for addr in ironwood_incomplete {
1403            let range = super::get_block_range(conn, ShieldedPool::Ironwood, addr)?;
1404            scan_ranges.extend(range);
1405        }
1406    }
1407
1408    Ok(scan_ranges)
1409}
1410
1411/// Generate Orchard Merkle witnesses at a historical height.
1412///
1413/// Loads the wallet's Orchard shard data into an ephemeral in-memory
1414/// [`MemoryShardStore`], inserts the provided frontier at that height as a
1415/// checkpoint, and generates a witness for each of the given note positions.
1416///
1417/// It is assumed that the caller provides the valid frontier at the given height.
1418///
1419/// How it works:
1420/// To construct witnesses at a historical height, we need:
1421/// 1. Authentication path within each note's shard — the scanner marks the
1422///    wallet's notes as MARKED, preventing them and their siblings within a
1423///    shard from being pruned.
1424/// 2. Cap — the upper tree above the shard level.
1425/// 3. Frontier — the right edge at the historical height. It lets ShardTree
1426///    know exactly where the tree ended at that height.
1427///
1428/// The wallet automatically prunes the tree after PRUNING_DEPTH checkpoints.
1429/// These three components are sufficient to reconstruct the tree structure
1430/// needed for witness generation even after pruning has occurred.
1431///
1432/// The wallet DB is strictly read-only. Shard data is read, decoded, and
1433/// inserted into an ephemeral in-memory [`ShardStore`] to avoid tampering with
1434/// the primary wallet DB.
1435///
1436/// Example application: token holder voting. The wallet tree may have advanced past
1437/// the historical height, but we need witnesses anchored at that frontier.
1438///
1439/// # Errors
1440///
1441/// - [`SqliteClientError::CommitmentTree`] if reading the wallet's shard or
1442///   cap data fails, or if the shard data reconstructed from the wallet is
1443///   internally inconsistent at a node the computation requires.
1444/// - [`SqliteClientError::HistoricalFrontierInvalid`] if `frontier_at_height`
1445///   is inconsistent with the shard data reconstructed from the wallet.
1446/// - [`SqliteClientError::HistoricalWitnessUnavailable`] if a witness cannot
1447///   be generated for one of the requested positions at `height` (most
1448///   commonly because the wallet has not yet synced through that height).
1449#[cfg(feature = "orchard")]
1450pub(crate) fn generate_orchard_witnesses_at_historical_height(
1451    conn: &rusqlite::Connection,
1452    note_positions: &[Position],
1453    frontier_at_height: incrementalmerkletree::frontier::NonEmptyFrontier<
1454        orchard::tree::MerkleHashOrchard,
1455    >,
1456    height: BlockHeight,
1457) -> Result<
1458    Vec<
1459        incrementalmerkletree::MerklePath<
1460            orchard::tree::MerkleHashOrchard,
1461            { orchard::NOTE_COMMITMENT_TREE_DEPTH as u8 },
1462        >,
1463    >,
1464    SqliteClientError,
1465> {
1466    generate_orchard_like_witnesses_at_historical_height(
1467        conn,
1468        ORCHARD_TABLES_PREFIX,
1469        note_positions,
1470        frontier_at_height,
1471        height,
1472    )
1473}
1474
1475/// Generates Ironwood Merkle witnesses at a historical height.
1476///
1477/// This is identical to [`generate_orchard_witnesses_at_historical_height`],
1478/// except that it reconstructs witness paths from the Ironwood shard tables.
1479#[cfg(feature = "orchard")]
1480pub(crate) fn generate_ironwood_witnesses_at_historical_height(
1481    conn: &rusqlite::Connection,
1482    note_positions: &[Position],
1483    frontier_at_height: incrementalmerkletree::frontier::NonEmptyFrontier<
1484        orchard::tree::MerkleHashOrchard,
1485    >,
1486    height: BlockHeight,
1487) -> Result<
1488    Vec<
1489        incrementalmerkletree::MerklePath<
1490            orchard::tree::MerkleHashOrchard,
1491            { orchard::NOTE_COMMITMENT_TREE_DEPTH as u8 },
1492        >,
1493    >,
1494    SqliteClientError,
1495> {
1496    generate_orchard_like_witnesses_at_historical_height(
1497        conn,
1498        IRONWOOD_TABLES_PREFIX,
1499        note_positions,
1500        frontier_at_height,
1501        height,
1502    )
1503}
1504
1505#[cfg(feature = "orchard")]
1506fn generate_orchard_like_witnesses_at_historical_height(
1507    conn: &rusqlite::Connection,
1508    table_prefix: &'static str,
1509    note_positions: &[Position],
1510    frontier_at_height: incrementalmerkletree::frontier::NonEmptyFrontier<
1511        orchard::tree::MerkleHashOrchard,
1512    >,
1513    height: BlockHeight,
1514) -> Result<
1515    Vec<
1516        incrementalmerkletree::MerklePath<
1517            orchard::tree::MerkleHashOrchard,
1518            { orchard::NOTE_COMMITMENT_TREE_DEPTH as u8 },
1519        >,
1520    >,
1521    SqliteClientError,
1522> {
1523    // `get_shard_roots` returns addresses ordered by shard index, matching the
1524    // ascending insertion order required by `MemoryShardStore::put_shard`.
1525    // Storage errors flow through `From<ShardTreeError<commitment_tree::Error>>`
1526    // into `SqliteClientError::CommitmentTree`.
1527    let mut store = MemoryShardStore::<orchard::tree::MerkleHashOrchard, BlockHeight>::empty();
1528    let shard_root_level = Level::new(ORCHARD_SHARD_HEIGHT);
1529    let shard_roots =
1530        get_shard_roots(conn, table_prefix, shard_root_level).map_err(ShardTreeError::Storage)?;
1531    for shard_root in shard_roots {
1532        if let Some(shard) =
1533            get_shard::<orchard::tree::MerkleHashOrchard>(conn, table_prefix, shard_root)
1534                .map_err(ShardTreeError::Storage)?
1535        {
1536            store.put_shard(shard).expect("put_shard is infallible");
1537        }
1538    }
1539    let cap = get_cap::<orchard::tree::MerkleHashOrchard>(conn, table_prefix)
1540        .map_err(ShardTreeError::Storage)?;
1541    store.put_cap(cap).expect("put_cap is infallible");
1542
1543    // Only one checkpoint is needed (the historical frontier), but `ShardTree`
1544    // requires a nonzero checkpoint limit.
1545    //
1546    // Pruning-safety invariant: `MemoryShardStore::empty()` starts with zero
1547    // checkpoints in its internal `BTreeMap`, and the only mutations above
1548    // (`put_shard` / `put_cap`) do not touch that map (the `CHECKPOINT`
1549    // retention flags stored *inside* shard leaves are independent of the
1550    // store's `checkpoint_count()`). So when `insert_frontier_nodes` below
1551    // calls `add_checkpoint` exactly once and then `prune_excess_checkpoints`,
1552    // we have `1 > 1 == false` and the freshly inserted historical checkpoint
1553    // is NOT pruned away before `witness_at_checkpoint_id` reads it.
1554    //
1555    // If this code is changed to (a) pre-load wallet checkpoints into the
1556    // in-memory store, (b) add an extra `add_checkpoint` call here, or
1557    // (c) drop `max_checkpoints` below `1`, the historical checkpoint will be
1558    // pruned immediately and witness generation will return
1559    // `HistoricalWitnessUnavailable` at the `witness_at_checkpoint_id` call
1560    // below. The `witnesses_at_historical_height_with_many_wallet_checkpoints`
1561    // test in `mod tests` exists to catch exactly that regression.
1562    let mut tree =
1563        ShardTree::<_, { orchard::NOTE_COMMITMENT_TREE_DEPTH as u8 }, ORCHARD_SHARD_HEIGHT>::new(
1564            store, 1,
1565        );
1566
1567    // Insert the frontier. `Retention::Checkpoint` causes `ShardTree` to
1568    // register a checkpoint at `height` internally; the pruning-safety
1569    // invariant above guarantees it survives the `prune_excess_checkpoints`
1570    // pass that runs at the end of `insert_frontier_nodes`.
1571    //
1572    // `MemoryShardStore::Error` is `Infallible`, so the only variants of
1573    // `ShardTreeError` we can observe are `Insert` (a caller-supplied
1574    // frontier inconsistent with the loaded shards) and `Query` (an
1575    // inconsistency in the loaded shards themselves).
1576    tree.insert_frontier_nodes(
1577        frontier_at_height,
1578        Retention::Checkpoint {
1579            id: height,
1580            marking: Marking::None,
1581        },
1582    )
1583    .map_err(|e| match e {
1584        ShardTreeError::Insert(e) => SqliteClientError::HistoricalFrontierInvalid(e),
1585        ShardTreeError::Query(q) => SqliteClientError::CommitmentTree(ShardTreeError::Query(q)),
1586        ShardTreeError::Storage(inf) => match inf {},
1587    })?;
1588
1589    // Generate a witness per note position. Any `Query` failure (or a `None`
1590    // result) at this stage means the tree reconstructed from the wallet's
1591    // shards does not contain enough information to compute a witness at
1592    // `height`; we surface that as `HistoricalWitnessUnavailable` so the
1593    // caller can either sync further or stop requesting that position.
1594    let mut witnesses = Vec::with_capacity(note_positions.len());
1595    for &pos in note_positions {
1596        let merkle_path = tree
1597            .witness_at_checkpoint_id(pos, &height)
1598            .map_err(|e| match e {
1599                ShardTreeError::Query(_) => SqliteClientError::HistoricalWitnessUnavailable {
1600                    position: pos,
1601                    height,
1602                },
1603                ShardTreeError::Insert(i) => {
1604                    SqliteClientError::CommitmentTree(ShardTreeError::Insert(i))
1605                }
1606                ShardTreeError::Storage(inf) => match inf {},
1607            })?
1608            .ok_or(SqliteClientError::HistoricalWitnessUnavailable {
1609                position: pos,
1610                height,
1611            })?;
1612
1613        witnesses.push(merkle_path);
1614    }
1615
1616    Ok(witnesses)
1617}
1618
1619#[cfg(test)]
1620mod tests {
1621    use tempfile::NamedTempFile;
1622
1623    use incrementalmerkletree::{Marking, Position, Retention};
1624    use incrementalmerkletree_testing::{
1625        check_append, check_checkpoint_rewind, check_remove_mark, check_rewind_remove_mark,
1626        check_root_hashes, check_witness_consistency, check_witnesses,
1627    };
1628    use shardtree::{ShardTree, error::ShardTreeError, store::ShardStore};
1629    use zcash_client_backend::data_api::{
1630        WalletCommitmentTrees,
1631        chain::CommitmentTreeRoot,
1632        testing::{pool::ShieldedPoolTester, sapling::SaplingPoolTester},
1633    };
1634    use zcash_protocol::consensus::{BlockHeight, Network};
1635
1636    use super::SqliteShardStore;
1637    use crate::{
1638        WalletDb,
1639        testing::{
1640            db::{test_clock, test_rng},
1641            pool::ShieldedPoolPersistence,
1642        },
1643        wallet::init::WalletMigrator,
1644    };
1645    // Used only by the orchard-gated `HistoricalWitnessGenerator` type alias below.
1646    use std::collections::BTreeSet;
1647    #[cfg(feature = "orchard")]
1648    use {
1649        crate::error::SqliteClientError, ::orchard::tree::MerkleHashOrchard,
1650        incrementalmerkletree::frontier::Frontier, rand::SeedableRng, rand_chacha::ChaChaRng,
1651        zcash_client_backend::data_api::ORCHARD_SHARD_HEIGHT,
1652    };
1653
1654    fn new_tree<T: ShieldedPoolTester + ShieldedPoolPersistence>(
1655        m: usize,
1656    ) -> ShardTree<SqliteShardStore<rusqlite::Connection, String, 3>, 4, 3> {
1657        let data_file = NamedTempFile::new().unwrap();
1658        let mut db_data = WalletDb::for_path(
1659            data_file.path(),
1660            Network::TestNetwork,
1661            test_clock(),
1662            test_rng(),
1663        )
1664        .unwrap();
1665        data_file.keep().unwrap();
1666
1667        WalletMigrator::new().init_or_migrate(&mut db_data).unwrap();
1668        let store =
1669            SqliteShardStore::<_, String, 3>::from_connection(db_data.conn, T::TABLES_PREFIX)
1670                .unwrap();
1671        ShardTree::new(store, m)
1672    }
1673
1674    fn check_retained_checkpoints(
1675        mut tree: ShardTree<SqliteShardStore<rusqlite::Connection, String, 3>, 4, 3>,
1676    ) {
1677        let h1 = BlockHeight::from(100);
1678        let h2 = BlockHeight::from(200);
1679
1680        assert!(tree.store().retained_checkpoints().unwrap().is_empty());
1681
1682        tree.ensure_retained(h1).unwrap();
1683        tree.ensure_retained(h2).unwrap();
1684        // Retaining an already-retained checkpoint is idempotent.
1685        tree.ensure_retained(h1).unwrap();
1686        assert_eq!(
1687            tree.store().retained_checkpoints().unwrap(),
1688            BTreeSet::from([h1, h2])
1689        );
1690
1691        tree.remove_retained_checkpoint(&h1).unwrap();
1692        assert_eq!(
1693            tree.store().retained_checkpoints().unwrap(),
1694            BTreeSet::from([h2])
1695        );
1696    }
1697
1698    #[cfg(feature = "orchard")]
1699    mod orchard {
1700        use super::new_tree;
1701        use zcash_client_backend::data_api::testing::orchard::OrchardPoolTester;
1702
1703        #[test]
1704        fn append() {
1705            super::check_append(new_tree::<OrchardPoolTester>);
1706        }
1707
1708        #[test]
1709        fn root_hashes() {
1710            super::check_root_hashes(new_tree::<OrchardPoolTester>);
1711        }
1712
1713        #[test]
1714        fn witnesses() {
1715            super::check_witnesses(new_tree::<OrchardPoolTester>);
1716        }
1717
1718        #[test]
1719        fn witness_consistency() {
1720            super::check_witness_consistency(new_tree::<OrchardPoolTester>);
1721        }
1722
1723        #[test]
1724        fn checkpoint_rewind() {
1725            super::check_checkpoint_rewind(new_tree::<OrchardPoolTester>);
1726        }
1727
1728        #[test]
1729        fn remove_mark() {
1730            super::check_remove_mark(new_tree::<OrchardPoolTester>);
1731        }
1732
1733        #[test]
1734        fn rewind_remove_mark() {
1735            super::check_rewind_remove_mark(new_tree::<OrchardPoolTester>);
1736        }
1737
1738        #[test]
1739        fn witnesses_at_historical_height() {
1740            super::witnesses_at_historical_height()
1741        }
1742
1743        #[test]
1744        fn ironwood_witnesses_at_historical_height() {
1745            super::ironwood_witnesses_at_historical_height()
1746        }
1747
1748        #[test]
1749        fn witnesses_at_historical_height_with_many_wallet_checkpoints() {
1750            super::witnesses_at_historical_height_with_many_wallet_checkpoints()
1751        }
1752
1753        #[test]
1754        fn put_shard_roots() {
1755            super::put_shard_roots::<OrchardPoolTester>()
1756        }
1757
1758        #[test]
1759        fn retained_checkpoints() {
1760            super::check_retained_checkpoints(super::new_tree::<OrchardPoolTester>(10));
1761        }
1762    }
1763
1764    #[test]
1765    fn sapling_append() {
1766        check_append(new_tree::<SaplingPoolTester>);
1767    }
1768
1769    #[test]
1770    fn sapling_retained_checkpoints() {
1771        check_retained_checkpoints(new_tree::<SaplingPoolTester>(10));
1772    }
1773
1774    #[test]
1775    fn remove_retained_checkpoints_below() {
1776        let data_file = NamedTempFile::new().unwrap();
1777        let mut db = WalletDb::for_path(
1778            data_file.path(),
1779            Network::TestNetwork,
1780            test_clock(),
1781            test_rng(),
1782        )
1783        .unwrap();
1784        WalletMigrator::new().init_or_migrate(&mut db).unwrap();
1785
1786        db.with_sapling_tree_mut(|tree| {
1787            for h in [100u32, 200, 300] {
1788                tree.ensure_retained(BlockHeight::from(h))?;
1789            }
1790            Ok::<_, ShardTreeError<_>>(())
1791        })
1792        .unwrap();
1793
1794        #[cfg(feature = "orchard")]
1795        {
1796            db.with_orchard_tree_mut(|tree| {
1797                for h in [100u32, 200, 300] {
1798                    tree.ensure_retained(BlockHeight::from(h))?;
1799                }
1800                Ok::<_, ShardTreeError<_>>(())
1801            })
1802            .unwrap();
1803
1804            db.with_ironwood_tree_mut(|tree| {
1805                for h in [100u32, 200, 300] {
1806                    tree.ensure_retained(BlockHeight::from(h))?;
1807                }
1808                Ok::<_, ShardTreeError<_>>(())
1809            })
1810            .unwrap();
1811        }
1812
1813        db.remove_retained_checkpoints_below(BlockHeight::from(250))
1814            .unwrap();
1815
1816        let remaining = db
1817            .with_sapling_tree_mut(|tree| {
1818                tree.store()
1819                    .retained_checkpoints()
1820                    .map_err(ShardTreeError::Storage)
1821            })
1822            .unwrap();
1823        assert_eq!(remaining, BTreeSet::from([BlockHeight::from(300)]));
1824
1825        // The retained checkpoints must be pruned in the Orchard and Ironwood trees as well, not
1826        // just Sapling.
1827        #[cfg(feature = "orchard")]
1828        {
1829            let orchard_remaining = db
1830                .with_orchard_tree_mut(|tree| {
1831                    tree.store()
1832                        .retained_checkpoints()
1833                        .map_err(ShardTreeError::Storage)
1834                })
1835                .unwrap();
1836            assert_eq!(orchard_remaining, BTreeSet::from([BlockHeight::from(300)]));
1837
1838            let ironwood_remaining = db
1839                .with_ironwood_tree_mut(|tree| {
1840                    tree.store()
1841                        .retained_checkpoints()
1842                        .map_err(ShardTreeError::Storage)
1843                })
1844                .unwrap()
1845                .expect("the wallet tracks an Ironwood tree");
1846            assert_eq!(
1847                ironwood_remaining,
1848                BTreeSet::from([BlockHeight::from(300)]),
1849                "retained Ironwood checkpoints below the max height must be released",
1850            );
1851        }
1852    }
1853
1854    #[test]
1855    fn sapling_root_hashes() {
1856        check_root_hashes(new_tree::<SaplingPoolTester>);
1857    }
1858
1859    #[test]
1860    fn sapling_witnesses() {
1861        check_witnesses(new_tree::<SaplingPoolTester>);
1862    }
1863
1864    #[test]
1865    fn sapling_witness_consistency() {
1866        check_witness_consistency(new_tree::<SaplingPoolTester>);
1867    }
1868
1869    #[test]
1870    fn sapling_checkpoint_rewind() {
1871        check_checkpoint_rewind(new_tree::<SaplingPoolTester>);
1872    }
1873
1874    #[test]
1875    fn sapling_remove_mark() {
1876        check_remove_mark(new_tree::<SaplingPoolTester>);
1877    }
1878
1879    #[test]
1880    fn sapling_rewind_remove_mark() {
1881        check_rewind_remove_mark(new_tree::<SaplingPoolTester>);
1882    }
1883
1884    #[test]
1885    fn sapling_put_shard_roots() {
1886        put_shard_roots::<SaplingPoolTester>()
1887    }
1888
1889    fn put_shard_roots<T: ShieldedPoolTester + ShieldedPoolPersistence>() {
1890        let data_file = NamedTempFile::new().unwrap();
1891        let mut db_data = WalletDb::for_path(
1892            data_file.path(),
1893            Network::TestNetwork,
1894            test_clock(),
1895            test_rng(),
1896        )
1897        .unwrap();
1898        data_file.keep().unwrap();
1899
1900        WalletMigrator::new().init_or_migrate(&mut db_data).unwrap();
1901        let tx = db_data.conn.transaction().unwrap();
1902        let store =
1903            SqliteShardStore::<_, String, 3>::from_connection(&tx, T::TABLES_PREFIX).unwrap();
1904
1905        // introduce some roots
1906        let roots = (0u32..4)
1907            .map(|idx| {
1908                CommitmentTreeRoot::from_parts(
1909                    BlockHeight::from((idx + 1) * 3),
1910                    if idx == 3 {
1911                        "abcdefgh".to_string()
1912                    } else {
1913                        idx.to_string()
1914                    },
1915                )
1916            })
1917            .collect::<Vec<_>>();
1918        super::put_shard_roots::<_, 6, 3>(store.conn, T::TABLES_PREFIX, 0, &roots).unwrap();
1919
1920        // simulate discovery of a note
1921        let mut tree = ShardTree::<_, 6, 3>::new(store, 10);
1922        let checkpoint_height = BlockHeight::from(3);
1923        tree.batch_insert(
1924            Position::from(24),
1925            ('a'..='h').map(|c| {
1926                (
1927                    c.to_string(),
1928                    match c {
1929                        'c' => Retention::Marked,
1930                        'h' => Retention::Checkpoint {
1931                            id: checkpoint_height,
1932                            marking: Marking::None,
1933                        },
1934                        _ => Retention::Ephemeral,
1935                    },
1936                )
1937            }),
1938        )
1939        .unwrap();
1940
1941        // construct a witness for the note
1942        let witness = tree
1943            .witness_at_checkpoint_id(Position::from(26), &checkpoint_height)
1944            .unwrap();
1945        assert_eq!(
1946            witness
1947                .expect("an anchor exists at the expected checkpoint height")
1948                .path_elems(),
1949            &[
1950                "d",
1951                "ab",
1952                "efgh",
1953                "2",
1954                "01",
1955                "________________________________"
1956            ]
1957        );
1958    }
1959
1960    /// Test that `generate_orchard_witnesses_at_historical_height` produces valid
1961    /// witnesses when given a frontier extracted from an earlier tree state.
1962    #[cfg(feature = "orchard")]
1963    fn witnesses_at_historical_height() {
1964        witnesses_at_historical_height_for_table(
1965            crate::ORCHARD_TABLES_PREFIX,
1966            super::generate_orchard_witnesses_at_historical_height,
1967        )
1968    }
1969
1970    /// Test that `generate_ironwood_witnesses_at_historical_height` uses the
1971    /// Ironwood shard tables rather than the Orchard shard tables.
1972    #[cfg(feature = "orchard")]
1973    fn ironwood_witnesses_at_historical_height() {
1974        witnesses_at_historical_height_for_table(
1975            crate::IRONWOOD_TABLES_PREFIX,
1976            super::generate_ironwood_witnesses_at_historical_height,
1977        )
1978    }
1979
1980    #[cfg(feature = "orchard")]
1981    type OrchardFrontier =
1982        incrementalmerkletree::frontier::NonEmptyFrontier<::orchard::tree::MerkleHashOrchard>;
1983
1984    #[cfg(feature = "orchard")]
1985    type OrchardMerklePath = incrementalmerkletree::MerklePath<
1986        ::orchard::tree::MerkleHashOrchard,
1987        { ::orchard::NOTE_COMMITMENT_TREE_DEPTH as u8 },
1988    >;
1989
1990    #[cfg(feature = "orchard")]
1991    type HistoricalWitnessGenerator = fn(
1992        &rusqlite::Connection,
1993        &[Position],
1994        OrchardFrontier,
1995        BlockHeight,
1996    ) -> Result<Vec<OrchardMerklePath>, SqliteClientError>;
1997
1998    #[cfg(feature = "orchard")]
1999    fn witnesses_at_historical_height_for_table(
2000        table_prefix: &'static str,
2001        generate_witnesses: HistoricalWitnessGenerator,
2002    ) {
2003        let data_file = NamedTempFile::new().unwrap();
2004        let mut db_data = WalletDb::for_path(
2005            data_file.path(),
2006            Network::TestNetwork,
2007            test_clock(),
2008            test_rng(),
2009        )
2010        .unwrap();
2011        data_file.keep().unwrap();
2012
2013        WalletMigrator::new().init_or_migrate(&mut db_data).unwrap();
2014
2015        let mut rng = ChaChaRng::seed_from_u64(0);
2016
2017        // We build two parallel trees: the wallet's ShardTree (persisted to the DB)
2018        // and a lightweight Frontier that captures the tree state at the historical height.
2019        let mut frontier_tree: Frontier<MerkleHashOrchard, 32> = Frontier::empty();
2020        let historical_height = BlockHeight::from(100);
2021        let note_position = Position::from(2);
2022        let note_leaf;
2023
2024        {
2025            let tx = db_data.conn.transaction().unwrap();
2026            let store =
2027                SqliteShardStore::<_, MerkleHashOrchard, ORCHARD_SHARD_HEIGHT>::from_connection(
2028                    &tx,
2029                    table_prefix,
2030                )
2031                .unwrap();
2032            let mut tree = ShardTree::<
2033                _,
2034                { ::orchard::NOTE_COMMITMENT_TREE_DEPTH as u8 },
2035                ORCHARD_SHARD_HEIGHT,
2036            >::new(store, 100);
2037
2038            let mut leaves = Vec::new();
2039            for _ in 0u64..5 {
2040                leaves.push(MerkleHashOrchard::random(&mut rng));
2041            }
2042            note_leaf = leaves[u64::from(note_position) as usize];
2043
2044            for (i, &leaf) in leaves.iter().enumerate() {
2045                let retention = if i == u64::from(note_position) as usize {
2046                    Retention::Marked
2047                } else {
2048                    Retention::Ephemeral
2049                };
2050                tree.append(leaf, retention).unwrap();
2051                frontier_tree.append(leaf);
2052            }
2053
2054            // Advance the tree past the historical height, simulating the
2055            // wallet continuing to sync afterward.
2056            tree.checkpoint(historical_height).unwrap();
2057            for _ in 0..5 {
2058                tree.append(MerkleHashOrchard::random(&mut rng), Retention::Ephemeral)
2059                    .unwrap();
2060            }
2061            tree.checkpoint(BlockHeight::from(200)).unwrap();
2062
2063            tx.commit().unwrap();
2064        }
2065
2066        let expected_root = frontier_tree.root();
2067        let frontier = frontier_tree.take().expect("frontier is non-empty");
2068
2069        let witnesses =
2070            generate_witnesses(&db_data.conn, &[note_position], frontier, historical_height)
2071                .expect("witness generation should succeed");
2072
2073        assert_eq!(witnesses.len(), 1);
2074        assert_eq!(witnesses[0].root(note_leaf), expected_root);
2075    }
2076
2077    /// Regression test: `generate_orchard_witnesses_at_historical_height` must not
2078    /// lose its freshly inserted historical checkpoint to pruning, even when the
2079    /// wallet has advanced many checkpoints past the historical height (so that
2080    /// the historical checkpoint has long since been pruned from the wallet's own
2081    /// `ShardTree`).
2082    ///
2083    /// Flow:
2084    /// 1. Seed: append 2 leaves to the wallet's `ShardTree` (position 0
2085    ///    `Marked`, position 1 `Ephemeral`), mirror them into a parallel
2086    ///    `Frontier` to capture the ground-truth state, then `checkpoint` at
2087    ///    `historical_height = 10`.
2088    /// 2. Bury: append 249 more `Ephemeral`-only blocks, each with its own
2089    ///    checkpoint. With `WALLET_MAX_CHECKPOINTS = 100` the wallet's pruner
2090    ///    evicts checkpoints 10..=159, including the one at
2091    ///    `historical_height`. The `Marked` retention at position 0 survives
2092    ///    because no surviving checkpoint schedules it for unmarking.
2093    /// 3. Precondition: assert `min_checkpoint_id > historical_height` so the
2094    ///    test fails loudly if a future tweak (e.g. shrinking `TOTAL_BLOCKS`
2095    ///    or growing `WALLET_MAX_CHECKPOINTS`) accidentally leaves the
2096    ///    historical checkpoint alive in the DB.
2097    /// 4. Exercise: call `generate_orchard_witnesses_at_historical_height`
2098    ///    with the captured frontier. Internally it builds a fresh
2099    ///    `MemoryShardStore` (no checkpoints), `ShardTree::new(store, 1)`,
2100    ///    and `insert_frontier_nodes(.., Retention::Checkpoint { id: 10 })`,
2101    ///    which adds exactly one checkpoint and then runs
2102    ///    `prune_excess_checkpoints` (`1 > 1` is false, so no eviction).
2103    /// 5. Verify: witness generation returns `Ok`, and
2104    ///    `witness.root(note_leaf) == frontier_tree.root()` confirms the
2105    ///    reconstructed anchor matches the historical state captured in (1).
2106    ///
2107    /// If anyone later changes the in-memory `max_checkpoints` to `0`, copies
2108    /// wallet checkpoints into the in-memory store, or otherwise inflates
2109    /// `checkpoint_count` before the frontier insertion, the freshly added
2110    /// historical checkpoint would be pruned and `witness_at_checkpoint_id`
2111    /// would return `None`, surfacing the regression as
2112    /// `SqliteClientError::HistoricalWitnessUnavailable`.
2113    #[cfg(feature = "orchard")]
2114    fn witnesses_at_historical_height_with_many_wallet_checkpoints() {
2115        // Wallet tree capacity << number of blocks we sync, so the historical
2116        // checkpoint is forcibly pruned from the wallet's own ShardTree before
2117        // we attempt to generate a witness at that height.
2118        const WALLET_MAX_CHECKPOINTS: usize = 100;
2119        const TOTAL_BLOCKS: u32 = 250;
2120        const LEAVES_PER_BLOCK: usize = 2;
2121
2122        let data_file = NamedTempFile::new().unwrap();
2123        let mut db_data = WalletDb::for_path(
2124            data_file.path(),
2125            Network::TestNetwork,
2126            test_clock(),
2127            test_rng(),
2128        )
2129        .unwrap();
2130        data_file.keep().unwrap();
2131
2132        WalletMigrator::new().init_or_migrate(&mut db_data).unwrap();
2133
2134        let mut rng = ChaChaRng::seed_from_u64(1);
2135
2136        let mut frontier_tree: Frontier<MerkleHashOrchard, 32> = Frontier::empty();
2137        let historical_height = BlockHeight::from(10);
2138        let note_position = Position::from(0);
2139        let note_leaf;
2140
2141        {
2142            let tx = db_data.conn.transaction().unwrap();
2143            let store =
2144                SqliteShardStore::<_, MerkleHashOrchard, ORCHARD_SHARD_HEIGHT>::from_connection(
2145                    &tx, "orchard",
2146                )
2147                .unwrap();
2148            let mut tree = ShardTree::<
2149                _,
2150                { ::orchard::NOTE_COMMITMENT_TREE_DEPTH as u8 },
2151                ORCHARD_SHARD_HEIGHT,
2152            >::new(store, WALLET_MAX_CHECKPOINTS);
2153
2154            // First block: mark the note, snapshot the frontier, then checkpoint
2155            // at `historical_height`.
2156            let note_idx = u64::from(note_position) as usize;
2157            let mut first_leaves = Vec::with_capacity(LEAVES_PER_BLOCK);
2158            for _ in 0..LEAVES_PER_BLOCK {
2159                first_leaves.push(MerkleHashOrchard::random(&mut rng));
2160            }
2161            note_leaf = first_leaves[note_idx];
2162            for (i, &leaf) in first_leaves.iter().enumerate() {
2163                let retention = if i == note_idx {
2164                    Retention::Marked
2165                } else {
2166                    Retention::Ephemeral
2167                };
2168                tree.append(leaf, retention).unwrap();
2169                frontier_tree.append(leaf);
2170            }
2171            tree.checkpoint(historical_height).unwrap();
2172
2173            // Drive the wallet far past the historical height so its pruner
2174            // evicts the historical_height checkpoint from the SQLite store.
2175            for block in 1..TOTAL_BLOCKS {
2176                for _ in 0..LEAVES_PER_BLOCK {
2177                    tree.append(MerkleHashOrchard::random(&mut rng), Retention::Ephemeral)
2178                        .unwrap();
2179                }
2180                tree.checkpoint(historical_height + block).unwrap();
2181            }
2182            tx.commit().unwrap();
2183        }
2184
2185        // Sanity-check the precondition: the historical checkpoint must no
2186        // longer be present in the wallet's checkpoint table.
2187        let min_ckpt = super::min_checkpoint_id(&db_data.conn, "orchard")
2188            .unwrap()
2189            .expect("wallet has checkpoints");
2190        assert!(
2191            min_ckpt > historical_height,
2192            "test precondition: historical checkpoint should have been pruned, \
2193             but min retained checkpoint is {min_ckpt:?} <= {historical_height:?}",
2194        );
2195
2196        let expected_root = frontier_tree.root();
2197        let frontier = frontier_tree.take().expect("frontier is non-empty");
2198
2199        let witnesses = super::generate_orchard_witnesses_at_historical_height(
2200            &db_data.conn,
2201            &[note_position],
2202            frontier,
2203            historical_height,
2204        )
2205        .expect("witness generation should succeed despite deep wallet pruning");
2206
2207        assert_eq!(witnesses.len(), 1);
2208        assert_eq!(witnesses[0].root(note_leaf), expected_root);
2209    }
2210}