fix(pepper-sync): retain ZIP 318 anchor checkpoints

pepper-sync/src/sync.rs

@@ -446,6 +446,8 @@ where
         consensus_parameters,
     )?;
 
+    repin_anchor_checkpoints(consensus_parameters, &mut *wallet.write().await)?;
+
     let ufvks = wallet
         .read()
         .await
@@ -1726,8 +1728,9 @@ where
     H: incrementalmerkletree::Hashable + Clone + PartialEq,
 {
     let mut truncation_height = None;
+    let checkpoint_count = tree.store().checkpoint_count().expect("infallible");
     tree.store()
-        .for_each_checkpoint((MAX_REORG_ALLOWANCE + 1) as usize, |height, _| {
+        .for_each_checkpoint(checkpoint_count, |height, _| {
             if truncation_height.is_some() {
                 return Ok(());
             }
@@ -2030,6 +2033,7 @@ where
             fetch_request_sender,
             scan_range,
             highest_scanned_height,
+            witness::anchor_retention_policy(consensus_parameters),
             sapling_located_trees,
             orchard_located_trees,
             ironwood_located_trees,
@@ -2309,6 +2313,36 @@ where
     Ok(())
 }
 
+/// Re-derives the pinned anchor-checkpoint set of each shard tree from the retention policy in
+/// force for `consensus_parameters`, as of the wallet's newest scanned block.
+fn repin_anchor_checkpoints<W>(
+    consensus_parameters: &impl consensus::Parameters,
+    wallet: &mut W,
+) -> Result<(), SyncError<W::Error>>
+where
+    W: SyncWallet + SyncShardTrees,
+{
+    let Some(policy) = witness::anchor_retention_policy(consensus_parameters) else {
+        return Ok(());
+    };
+    let Some(highest_scanned_height) = wallet
+        .get_sync_state()
+        .map_err(SyncError::WalletError)?
+        .highest_scanned_height()
+    else {
+        return Ok(());
+    };
+    let window = witness::anchor_retention_window(&policy, highest_scanned_height);
+    let shard_trees = wallet
+        .get_shard_trees_mut()
+        .map_err(SyncError::WalletError)?;
+    witness::repin_anchor_checkpoints(&policy, &window, shard_trees.sapling.store_mut());
+    witness::repin_anchor_checkpoints(&policy, &window, shard_trees.orchard.store_mut());
+    witness::repin_anchor_checkpoints(&policy, &window, shard_trees.ironwood.store_mut());
+
+    Ok(())
+}
+
 async fn add_initial_frontier<W>(
     consensus_parameters: &impl consensus::Parameters,
     fetch_request_sender: mpsc::UnboundedSender<FetchRequest>,

pepper-sync/src/wallet/serialization.rs

@@ -1341,9 +1341,21 @@ impl ShardTrees {
                 Ok(())
             })
             .expect("Infallible");
-        if checkpoints.len() > MAX_REORG_ALLOWANCE as usize {
-            let keep_from = checkpoints.len() - MAX_REORG_ALLOWANCE as usize;
-            checkpoints.drain(..keep_from);
+        let retained = store.retained_checkpoints().expect("Infallible");
+        let non_retained_count = checkpoints
+            .iter()
+            .filter(|(checkpoint_id, _)| !retained.contains(checkpoint_id))
+            .count();
+        if non_retained_count > MAX_REORG_ALLOWANCE as usize {
+            let mut excess = non_retained_count - MAX_REORG_ALLOWANCE as usize;
+            checkpoints.retain(|(checkpoint_id, _)| {
+                if retained.contains(checkpoint_id) || excess == 0 {
+                    true
+                } else {
+                    excess -= 1;
+                    false
+                }
+            });
         }
         write_with_error_handling!(write_checkpoints, checkpoints);
 
@@ -1529,4 +1541,59 @@ mod tests {
                 .is_none()
         );
     }
+
+    /// A pinned anchor checkpoint survives the newest-100 write cap, while the pinned set
+    /// itself is not persisted.
+    #[test]
+    fn shardtree_roundtrip_keeps_pinned_anchor_checkpoints() {
+        let mut shard_trees = ShardTrees::new();
+
+        for height in 1..=150 {
+            let height = BlockHeight::from_u32(height);
+            shard_trees
+                .sapling
+                .store_mut()
+                .add_checkpoint(
+                    height,
+                    Checkpoint::from_parts(TreeState::Empty, BTreeSet::new()),
+                )
+                .expect("infallible");
+        }
+        shard_trees
+            .sapling
+            .store_mut()
+            .add_retained_checkpoint(BlockHeight::from_u32(24))
+            .expect("infallible");
+
+        let mut bytes = Vec::new();
+        shard_trees.write(&mut bytes).expect("write should succeed");
+        let roundtripped = ShardTrees::read(bytes.as_slice()).expect("read should succeed");
+
+        let sapling_store = roundtripped.sapling.store();
+        assert_eq!(sapling_store.checkpoint_count().expect("infallible"), 101);
+        assert!(
+            sapling_store
+                .get_checkpoint(&BlockHeight::from_u32(24))
+                .expect("infallible")
+                .is_some()
+        );
+        assert!(
+            sapling_store
+                .get_checkpoint(&BlockHeight::from_u32(50))
+                .expect("infallible")
+                .is_none()
+        );
+        assert!(
+            sapling_store
+                .get_checkpoint(&BlockHeight::from_u32(51))
+                .expect("infallible")
+                .is_some()
+        );
+        assert!(
+            sapling_store
+                .retained_checkpoints()
+                .expect("infallible")
+                .is_empty()
+        );
+    }
 }

pepper-sync/src/wallet/traits.rs

@@ -10,6 +10,7 @@ use orchard::tree::MerkleHashOrchard;
 use shardtree::ShardTree;
 use shardtree::store::memory::MemoryShardStore;
 use shardtree::store::{Checkpoint, ShardStore, TreeState};
+use zcash_client_backend::data_api::anchor_retention::AnchorRetention;
 use zcash_keys::keys::UnifiedFullViewingKey;
 use zcash_primitives::transaction::TxId;
 use zcash_protocol::consensus::BlockHeight;
@@ -252,11 +253,13 @@ pub trait SyncShardTrees: SyncWallet {
     /// Update wallet shard trees with new shard tree data.
     ///
     /// `highest_scanned_height` is the height of the highest scanned block in the wallet not including the `scan_range` we are updating.
+    #[allow(clippy::too_many_arguments)]
     fn update_shard_trees(
         &mut self,
         fetch_request_sender: mpsc::UnboundedSender<FetchRequest>,
         scan_range: &ScanRange,
         highest_scanned_height: BlockHeight,
+        anchor_retention: Option<AnchorRetention>,
         sapling_located_trees: Vec<LocatedTreeData<sapling_crypto::Node>>,
         orchard_located_trees: Vec<LocatedTreeData<MerkleHashOrchard>>,
         ironwood_located_trees: Vec<LocatedTreeData<MerkleHashOrchard>>,
@@ -336,6 +339,70 @@ pub trait SyncShardTrees: SyncWallet {
                 .await?;
             }
 
+            if let Some(retention) = &anchor_retention {
+                let as_of = std::cmp::max(highest_scanned_height, scan_range.block_range().end - 1);
+                let window = witness::anchor_retention_window(retention, as_of);
+                witness::repin_anchor_checkpoints(
+                    retention,
+                    &window,
+                    shard_trees.sapling.store_mut(),
+                );
+                witness::repin_anchor_checkpoints(
+                    retention,
+                    &window,
+                    shard_trees.orchard.store_mut(),
+                );
+                witness::repin_anchor_checkpoints(
+                    retention,
+                    &window,
+                    shard_trees.ironwood.store_mut(),
+                );
+
+                let start = std::cmp::max(*window.start(), scan_range.block_range().start);
+                let end = std::cmp::min(*window.end(), scan_range.block_range().end - 1);
+                for boundary in retention.retained_in_range(start..=end) {
+                    if checkpoint_range.contains(&boundary) {
+                        continue;
+                    }
+                    add_checkpoint::<
+                        Sapling,
+                        sapling_crypto::Node,
+                        { sapling_crypto::NOTE_COMMITMENT_TREE_DEPTH },
+                        { witness::SHARD_HEIGHT },
+                    >(
+                        fetch_request_sender.clone(),
+                        boundary,
+                        &sapling_located_trees,
+                        &mut shard_trees.sapling,
+                    )
+                    .await?;
+                    add_checkpoint::<
+                        Orchard,
+                        MerkleHashOrchard,
+                        { orchard::NOTE_COMMITMENT_TREE_DEPTH as u8 },
+                        { witness::SHARD_HEIGHT },
+                    >(
+                        fetch_request_sender.clone(),
+                        boundary,
+                        &orchard_located_trees,
+                        &mut shard_trees.orchard,
+                    )
+                    .await?;
+                    add_checkpoint::<
+                        Ironwood,
+                        MerkleHashOrchard,
+                        { orchard::NOTE_COMMITMENT_TREE_DEPTH as u8 },
+                        { witness::SHARD_HEIGHT },
+                    >(
+                        fetch_request_sender.clone(),
+                        boundary,
+                        &ironwood_located_trees,
+                        &mut shard_trees.ironwood,
+                    )
+                    .await?;
+                }
+            }
+
             for tree in sapling_located_trees {
                 shard_trees
                     .sapling

pepper-sync/src/witness.rs

@@ -9,6 +9,7 @@ use incrementalmerkletree::{
 use orchard::tree::MerkleHashOrchard;
 use sapling_crypto::Node;
 use shardtree::LocatedPrunableTree;
+use zcash_client_backend::data_api::anchor_retention::{AnchorRetention, AnchorRetentionInterval};
 use zcash_primitives::{block::BlockHash, merkle_tree::read_commitment_tree};
 use zcash_protocol::consensus::BlockHeight;
 use zingo_netutils::lightwallet_protocol::TreeState;
@@ -371,3 +372,116 @@ pub(crate) fn get_ironwood_tree(
         )
     }
 }
+
+/// How many grid intervals behind the wallet's newest scanned block anchor-boundary checkpoints
+/// stay pinned: the ZIP 318 anchor age cap plus two intervals of slack.
+pub(crate) const ANCHOR_RETENTION_INTERVALS: u32 = zcash_protocol::zip318::ANCHOR_AGE_CAP + 2;
+
+/// The anchor-checkpoint retention policy in force for `consensus_parameters` — the ZIP 318
+/// grid from the Ironwood pool's activation — or `None` when that upgrade never activates.
+pub(crate) fn anchor_retention_policy(
+    consensus_parameters: &impl zcash_protocol::consensus::Parameters,
+) -> Option<AnchorRetention> {
+    crate::wallet::PoolActivation::of(consensus_parameters, zcash_protocol::ShieldedPool::Ironwood)
+        .map(|activation| {
+            AnchorRetention::new(activation.height(), AnchorRetentionInterval::default())
+        })
+}
+
+/// The height range whose grid-boundary checkpoints must currently be pinned, given `as_of`,
+/// the wallet's newest scanned block.
+pub(crate) fn anchor_retention_window(
+    policy: &AnchorRetention,
+    as_of: BlockHeight,
+) -> std::ops::RangeInclusive<BlockHeight> {
+    let span = policy
+        .intervals()
+        .iter()
+        .map(|interval| interval.block_count().get())
+        .max()
+        .unwrap_or(0)
+        .saturating_mul(ANCHOR_RETENTION_INTERVALS);
+    as_of.saturating_sub(span)..=as_of
+}
+
+/// Re-derives one shard store's pinned-checkpoint set from `policy`, pinning every grid
+/// boundary in `window` and releasing every pin outside it.
+pub(crate) fn repin_anchor_checkpoints<S>(
+    policy: &AnchorRetention,
+    window: &std::ops::RangeInclusive<BlockHeight>,
+    store: &mut S,
+) where
+    S: ShardStore<CheckpointId = BlockHeight, Error = std::convert::Infallible>,
+{
+    for id in store.retained_checkpoints().expect("infallible") {
+        if !window.contains(&id) || !policy.retains(id) {
+            store.remove_retained_checkpoint(&id).expect("infallible");
+        }
+    }
+    for boundary in policy.retained_in_range(window.clone()) {
+        store.add_retained_checkpoint(boundary).expect("infallible");
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::shardtree_ext::ShardTreeExt;
+    use crate::wallet::empty_shard_tree;
+
+    /// Pinned grid boundaries survive checkpoint pruning until re-deriving the pins releases
+    /// them.
+    #[test]
+    fn pinned_boundaries_survive_pruning_until_released() {
+        let policy =
+            AnchorRetention::new(BlockHeight::from_u32(1), AnchorRetentionInterval::default());
+        let mut tree = empty_shard_tree::<
+            sapling_crypto::Node,
+            { sapling_crypto::NOTE_COMMITMENT_TREE_DEPTH },
+            SHARD_HEIGHT,
+        >();
+
+        for height in 1..=600u32 {
+            let height = BlockHeight::from_u32(height);
+            if policy.retains(height) {
+                tree.store_mut()
+                    .add_retained_checkpoint(height)
+                    .expect("infallible");
+            }
+            tree.append_checkpoint(height).expect("infallible");
+        }
+
+        let checkpoint_exists = |store: &shardtree::store::memory::MemoryShardStore<
+            sapling_crypto::Node,
+            BlockHeight,
+        >,
+                                 height: u32| {
+            store
+                .get_checkpoint(&BlockHeight::from_u32(height))
+                .expect("infallible")
+                .is_some()
+        };
+        for boundary in [144, 288, 432, 576] {
+            assert!(
+                checkpoint_exists(tree.store(), boundary),
+                "boundary {boundary}"
+            );
+        }
+        for pruned in [143, 300, 431] {
+            assert!(
+                !checkpoint_exists(tree.store(), pruned),
+                "non-boundary {pruned}"
+            );
+        }
+
+        let window = anchor_retention_window(&policy, BlockHeight::from_u32(1200));
+        assert_eq!(*window.start(), BlockHeight::from_u32(336));
+        repin_anchor_checkpoints(&policy, &window, tree.store_mut());
+        tree.append_checkpoint(BlockHeight::from_u32(601))
+            .expect("infallible");
+        assert!(!checkpoint_exists(tree.store(), 144));
+        assert!(!checkpoint_exists(tree.store(), 288));
+        assert!(checkpoint_exists(tree.store(), 432));
+        assert!(checkpoint_exists(tree.store(), 576));
+    }
+}

zingolib/src/wallet/send.rs

@@ -99,6 +99,20 @@ impl LightWallet {
     /// Whether this wallet can materialize `protocol`'s note commitment tree root, and witnesses to it, as of `height`.
     pub(crate) fn anchor_is_computable(&self, protocol: ShieldedPool, height: BlockHeight) -> bool {
         self.shards_are_scanned(protocol, None, height)
+            && self.checkpoint_is_retained(protocol, height)
+    }
+
+    /// Whether `protocol`'s shard tree retains a checkpoint at `height`.
+    fn checkpoint_is_retained(&self, protocol: ShieldedPool, height: BlockHeight) -> bool {
+        use shardtree::store::ShardStore;
+
+        match protocol {
+            ShieldedPool::Sapling => self.shard_trees.sapling.store().get_checkpoint(&height),
+            ShieldedPool::Orchard => self.shard_trees.orchard.store().get_checkpoint(&height),
+            ShieldedPool::Ironwood => self.shard_trees.ironwood.store().get_checkpoint(&height),
+        }
+        .expect("memory shard store is infallible")
+        .is_some()
     }
 
     /// Whether every shard carrying `protocol` notes between `note_height` (the scan floor when absent) and `anchor_height` is scanned.