use incrementalmerkletree::Hashable;
use shardtree::ShardTree;
use shardtree::store::ShardStore;
use shardtree::store::memory::MemoryShardStore;
use zcash_protocol::consensus::{self, BlockHeight};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct WalletTruncationState {
pub birthday: BlockHeight,
pub highest_scanned_height: BlockHeight,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TreeTruncationFacts {
pub checkpoint_at_target: Option<BlockHeight>,
pub newest_checkpoint: Option<BlockHeight>,
}
pub(crate) fn tree_facts<H, const DEPTH: u8, const SHARD_HEIGHT: u8>(
tree: &ShardTree<MemoryShardStore<H, BlockHeight>, DEPTH, SHARD_HEIGHT>,
truncate_height: BlockHeight,
) -> TreeTruncationFacts
where
H: Hashable + Clone + PartialEq,
{
TreeTruncationFacts {
checkpoint_at_target: tree
.store()
.get_checkpoint(&truncate_height)
.expect("infallible")
.map(|_| truncate_height),
newest_checkpoint: tree.store().max_checkpoint_id().expect("infallible"),
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PoolTruncation {
ToCheckpoint {
checkpoint: BlockHeight,
},
Untouched,
RequiresRescan {
newest_checkpoint: BlockHeight,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TruncationPlan {
NoOp,
ClearAll,
Truncate {
height: BlockHeight,
},
}
#[must_use]
pub fn plan_truncation(
wallet_state: WalletTruncationState,
truncate_height: BlockHeight,
) -> TruncationPlan {
if truncate_height == consensus::H0 || truncate_height < wallet_state.birthday {
return TruncationPlan::ClearAll;
}
if truncate_height > wallet_state.highest_scanned_height {
return TruncationPlan::NoOp;
}
TruncationPlan::Truncate {
height: truncate_height,
}
}
pub(crate) fn plan_pool_truncation(
facts: TreeTruncationFacts,
truncate_height: BlockHeight,
) -> PoolTruncation {
if let Some(checkpoint) = facts.checkpoint_at_target {
return PoolTruncation::ToCheckpoint { checkpoint };
}
match facts.newest_checkpoint {
Some(newest_checkpoint) if newest_checkpoint > truncate_height => {
PoolTruncation::RequiresRescan { newest_checkpoint }
}
_ => PoolTruncation::Untouched,
}
}
pub(crate) fn plan_unscanned_state_rollback(
facts: TreeTruncationFacts,
highest_scanned_height: BlockHeight,
) -> Option<BlockHeight> {
facts.checkpoint_at_target.filter(|_| {
facts
.newest_checkpoint
.is_some_and(|newest_checkpoint| newest_checkpoint > highest_scanned_height)
})
}
#[cfg(test)]
mod test {
use super::*;
use crate::shardtree_ext::{CheckpointAppendOutcome, ShardTreeExt};
use crate::wallet::ShardTrees;
fn wallet_state() -> WalletTruncationState {
WalletTruncationState {
birthday: BlockHeight::from_u32(6),
highest_scanned_height: BlockHeight::from_u32(10),
}
}
fn facts(at_target: Option<u32>, newest_checkpoint: Option<u32>) -> TreeTruncationFacts {
TreeTruncationFacts {
checkpoint_at_target: at_target.map(BlockHeight::from_u32),
newest_checkpoint: newest_checkpoint.map(BlockHeight::from_u32),
}
}
#[test]
fn migrated_empty_tree_is_untouched() {
assert_eq!(
plan_truncation(wallet_state(), BlockHeight::from_u32(8)),
TruncationPlan::Truncate {
height: BlockHeight::from_u32(8),
}
);
assert_eq!(
plan_pool_truncation(facts(Some(8), Some(10)), BlockHeight::from_u32(8)),
PoolTruncation::ToCheckpoint {
checkpoint: BlockHeight::from_u32(8),
}
);
assert_eq!(
plan_pool_truncation(facts(None, Some(0)), BlockHeight::from_u32(8)),
PoolTruncation::Untouched
);
}
#[test]
fn checkpointless_tree_is_untouched() {
assert_eq!(
plan_pool_truncation(facts(None, None), BlockHeight::from_u32(8)),
PoolTruncation::Untouched
);
}
#[test]
fn pruned_checkpoint_requires_rescan() {
assert_eq!(
plan_pool_truncation(facts(None, Some(10)), BlockHeight::from_u32(8)),
PoolTruncation::RequiresRescan {
newest_checkpoint: BlockHeight::from_u32(10),
}
);
}
#[test]
fn checkpoint_at_target_rolls_back() {
assert_eq!(
plan_pool_truncation(facts(Some(8), Some(10)), BlockHeight::from_u32(8)),
PoolTruncation::ToCheckpoint {
checkpoint: BlockHeight::from_u32(8),
}
);
}
#[test]
fn state_above_highest_scanned_rolls_back() {
assert_eq!(
plan_unscanned_state_rollback(facts(Some(10), Some(11)), BlockHeight::from_u32(10)),
Some(BlockHeight::from_u32(10))
);
}
#[test]
fn tree_at_highest_scanned_is_left_as_it_is() {
assert_eq!(
plan_unscanned_state_rollback(facts(Some(10), Some(10)), BlockHeight::from_u32(10)),
None
);
assert_eq!(
plan_unscanned_state_rollback(facts(None, Some(0)), BlockHeight::from_u32(10)),
None
);
}
#[test]
fn tree_without_a_checkpoint_at_highest_scanned_is_left_as_it_is() {
assert_eq!(
plan_unscanned_state_rollback(facts(None, Some(12)), BlockHeight::from_u32(10)),
None
);
}
#[test]
fn target_above_highest_scanned_is_a_no_op() {
let plan = plan_truncation(wallet_state(), BlockHeight::from_u32(11));
assert_eq!(plan, TruncationPlan::NoOp);
}
#[test]
fn target_below_birthday_clears_all() {
let plan = plan_truncation(wallet_state(), BlockHeight::from_u32(5));
assert_eq!(plan, TruncationPlan::ClearAll);
}
#[test]
fn height_zero_clears_all() {
let plan = plan_truncation(
WalletTruncationState {
birthday: consensus::H0,
highest_scanned_height: BlockHeight::from_u32(10),
},
consensus::H0,
);
assert_eq!(plan, TruncationPlan::ClearAll);
}
#[test]
fn target_at_birthday_truncates() {
let plan = plan_truncation(wallet_state(), BlockHeight::from_u32(6));
assert!(matches!(plan, TruncationPlan::Truncate { height }
if height == BlockHeight::from_u32(6)));
}
#[test]
fn exhaustive_sweep_matches_the_reference_model() {
const MAX: u32 = 6;
let heights = || (0..=MAX).map(BlockHeight::from_u32);
let newest_choices = || std::iter::once(None).chain(heights().map(Some));
for birthday in heights() {
for highest_scanned_height in heights() {
for target in heights() {
let plan = plan_truncation(
WalletTruncationState {
birthday,
highest_scanned_height,
},
target,
);
if target == consensus::H0 || target < birthday {
assert_eq!(plan, TruncationPlan::ClearAll);
continue;
}
if target > highest_scanned_height {
assert_eq!(plan, TruncationPlan::NoOp);
continue;
}
assert_eq!(plan, TruncationPlan::Truncate { height: target });
for checkpoint_at_target in [None, Some(target)] {
for newest_checkpoint in newest_choices() {
let outcome = plan_pool_truncation(
TreeTruncationFacts {
checkpoint_at_target,
newest_checkpoint,
},
target,
);
let expected = if let Some(checkpoint) = checkpoint_at_target {
PoolTruncation::ToCheckpoint { checkpoint }
} else {
match newest_checkpoint {
Some(newest_checkpoint) if newest_checkpoint > target => {
PoolTruncation::RequiresRescan { newest_checkpoint }
}
_ => PoolTruncation::Untouched,
}
};
assert_eq!(outcome, expected);
if newest_checkpoint.is_none_or(|newest| newest <= target) {
assert!(!matches!(outcome, PoolTruncation::RequiresRescan { .. }));
}
}
}
}
}
}
}
#[test]
fn gather_reports_store_ground_truth() {
let mut shard_trees = ShardTrees::new();
for height in [3u32, 5] {
assert_eq!(
shard_trees
.orchard
.append_checkpoint(BlockHeight::from_u32(height))
.unwrap(),
CheckpointAppendOutcome::Appended
);
}
assert_eq!(
tree_facts(&shard_trees.orchard, BlockHeight::from_u32(5)),
TreeTruncationFacts {
checkpoint_at_target: Some(BlockHeight::from_u32(5)),
newest_checkpoint: Some(BlockHeight::from_u32(5)),
}
);
assert_eq!(
tree_facts(&shard_trees.sapling, BlockHeight::from_u32(5)),
TreeTruncationFacts {
checkpoint_at_target: None,
newest_checkpoint: Some(BlockHeight::from_u32(0)),
}
);
assert_eq!(
tree_facts(&shard_trees.ironwood, BlockHeight::from_u32(5)),
tree_facts(&shard_trees.sapling, BlockHeight::from_u32(5)),
);
assert_eq!(
tree_facts(&shard_trees.orchard, BlockHeight::from_u32(4)),
TreeTruncationFacts {
checkpoint_at_target: None,
newest_checkpoint: Some(BlockHeight::from_u32(5)),
}
);
}
}
#[cfg(test)]
mod properties {
use std::collections::BTreeSet;
use proptest::prelude::*;
use super::*;
fn facts_for_set(checkpoints: &BTreeSet<u32>, target: u32) -> TreeTruncationFacts {
TreeTruncationFacts {
checkpoint_at_target: checkpoints
.contains(&target)
.then(|| BlockHeight::from_u32(target)),
newest_checkpoint: checkpoints
.last()
.map(|&height| BlockHeight::from_u32(height)),
}
}
fn checkpoint_sets() -> impl Strategy<Value = BTreeSet<u32>> {
proptest::collection::btree_set(0u32..=1_000, 0..16)
}
proptest! {
#[test]
fn pool_outcome_agrees_with_the_set_model(
checkpoints in checkpoint_sets(),
target in 0u32..=1_000,
) {
let outcome = plan_pool_truncation(
facts_for_set(&checkpoints, target),
BlockHeight::from_u32(target),
);
let expected = if checkpoints.contains(&target) {
PoolTruncation::ToCheckpoint {
checkpoint: BlockHeight::from_u32(target),
}
} else {
match checkpoints.last() {
Some(&newest) if newest > target => PoolTruncation::RequiresRescan {
newest_checkpoint: BlockHeight::from_u32(newest),
},
_ => PoolTruncation::Untouched,
}
};
prop_assert_eq!(outcome, expected);
}
#[test]
fn nothing_above_target_is_never_condemned(
candidates in checkpoint_sets(),
target in 0u32..=1_000,
) {
let checkpoints: BTreeSet<u32> = candidates
.into_iter()
.filter(|&height| height <= target)
.collect();
let outcome = plan_pool_truncation(
facts_for_set(&checkpoints, target),
BlockHeight::from_u32(target),
);
prop_assert!(
!matches!(outcome, PoolTruncation::RequiresRescan { .. }),
"a tree with nothing above target {} was condemned to rescan: {:?}",
target,
outcome,
);
}
#[test]
fn untouched_is_upward_closed(
checkpoints in checkpoint_sets(),
target in 0u32..1_000,
) {
let outcome_at = |t: u32| {
plan_pool_truncation(facts_for_set(&checkpoints, t), BlockHeight::from_u32(t))
};
if matches!(outcome_at(target), PoolTruncation::Untouched) {
prop_assert!(matches!(outcome_at(target + 1), PoolTruncation::Untouched));
}
}
#[test]
fn planner_is_total_over_its_regions(
birthday in 0u32..=1_000,
highest_scanned_height in 0u32..=1_000,
target in 0u32..=1_000,
) {
let plan = plan_truncation(
WalletTruncationState {
birthday: BlockHeight::from_u32(birthday),
highest_scanned_height: BlockHeight::from_u32(highest_scanned_height),
},
BlockHeight::from_u32(target),
);
if target == 0 || target < birthday {
prop_assert_eq!(plan, TruncationPlan::ClearAll);
} else if target > highest_scanned_height {
prop_assert_eq!(plan, TruncationPlan::NoOp);
} else {
prop_assert!(
matches!(plan, TruncationPlan::Truncate { height }
if height == BlockHeight::from_u32(target)),
"expected Truncate at height {}, got {:?}",
target,
plan,
);
}
}
}
}