use shardtree::ShardTree;
use shardtree::error::ShardTreeError;
use shardtree::store::memory::MemoryShardStore;
use std::convert::Infallible;
use incrementalmerkletree::Hashable;
use zcash_protocol::consensus::BlockHeight;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CheckpointAppendOutcome {
Appended,
NotAboveNewest,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RollbackOutcome {
RolledBack,
NoSuchCheckpoint,
}
pub trait ShardTreeExt {
fn append_checkpoint(
&mut self,
checkpoint_id: BlockHeight,
) -> Result<CheckpointAppendOutcome, ShardTreeError<Infallible>>;
fn rollback_to_checkpoint(
&mut self,
checkpoint_id: BlockHeight,
) -> Result<RollbackOutcome, ShardTreeError<Infallible>>;
}
impl<H, const DEPTH: u8, const SHARD_HEIGHT: u8> ShardTreeExt
for ShardTree<MemoryShardStore<H, BlockHeight>, DEPTH, SHARD_HEIGHT>
where
H: Hashable + Clone + PartialEq,
{
fn append_checkpoint(
&mut self,
checkpoint_id: BlockHeight,
) -> Result<CheckpointAppendOutcome, ShardTreeError<Infallible>> {
Ok(if self.checkpoint(checkpoint_id)? {
CheckpointAppendOutcome::Appended
} else {
CheckpointAppendOutcome::NotAboveNewest
})
}
fn rollback_to_checkpoint(
&mut self,
checkpoint_id: BlockHeight,
) -> Result<RollbackOutcome, ShardTreeError<Infallible>> {
Ok(if self.truncate_to_checkpoint(&checkpoint_id)? {
RollbackOutcome::RolledBack
} else {
RollbackOutcome::NoSuchCheckpoint
})
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn raw_boolean_calls_are_confined_to_this_module() {
let src_root = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("src");
let mut rust_sources = Vec::new();
collect_rust_sources(&src_root, &mut rust_sources);
assert!(!rust_sources.is_empty());
let raw_patterns = [
".checkpoint(",
".truncate_to_checkpoint(",
".truncate_to_checkpoint_depth(",
];
for path in rust_sources {
if path.ends_with("shardtree_ext.rs") {
continue;
}
let source = std::fs::read_to_string(&path).unwrap();
for pattern in raw_patterns {
assert!(
!source.contains(pattern),
"{} calls a raw shardtree boolean operation ({pattern}); \
route it through crate::shardtree_ext instead",
path.display(),
);
}
}
}
fn collect_rust_sources(dir: &std::path::Path, out: &mut Vec<std::path::PathBuf>) {
for entry in std::fs::read_dir(dir).unwrap() {
let path = entry.unwrap().path();
if path.is_dir() {
collect_rust_sources(&path, out);
} else if path.extension().is_some_and(|extension| extension == "rs") {
out.push(path);
}
}
}
#[test]
fn outcomes_match_store_behavior() {
let mut tree: ShardTree<
MemoryShardStore<orchard::tree::MerkleHashOrchard, BlockHeight>,
{ orchard::NOTE_COMMITMENT_TREE_DEPTH as u8 },
{ crate::witness::SHARD_HEIGHT },
> = ShardTree::new(MemoryShardStore::empty(), 10);
let five = BlockHeight::from_u32(5);
assert_eq!(
tree.append_checkpoint(five).unwrap(),
CheckpointAppendOutcome::Appended
);
assert_eq!(
tree.append_checkpoint(five).unwrap(),
CheckpointAppendOutcome::NotAboveNewest
);
assert_eq!(
tree.rollback_to_checkpoint(five).unwrap(),
RollbackOutcome::RolledBack
);
assert_eq!(
tree.rollback_to_checkpoint(BlockHeight::from_u32(4))
.unwrap(),
RollbackOutcome::NoSuchCheckpoint
);
}
}