use spine::{Hasher, frontier_for_size, nary_mr};
use crate::consistency::ProofStep;
use crate::error::{Error, Result};
use crate::mountain::{bag_path, bag_peaks, mountain_skeleton};
use crate::schedule::reduction_count;
pub trait NodeReader {
type Error;
fn get_node(
&self,
alg_id: u64,
left: u64,
height: u32,
) -> impl std::future::Future<Output = ReadResult<Option<Vec<u8>>, Self>> + Send;
fn get_leaf(
&self,
index: u64,
) -> impl std::future::Future<Output = ReadResult<Vec<u8>, Self>> + Send;
}
pub type ReadResult<T, R> = std::result::Result<T, <R as NodeReader>::Error>;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LogKind {
Flat,
Subtree,
}
impl LogKind {
#[must_use]
pub fn to_byte(self) -> u8 {
match self {
Self::Flat => 0,
Self::Subtree => 1,
}
}
#[must_use]
pub fn from_byte(b: u8) -> Option<Self> {
match b {
0 => Some(Self::Flat),
1 => Some(Self::Subtree),
_ => None,
}
}
}
#[derive(Debug)]
pub struct AlgView {
pub hasher: Box<dyn Hasher>,
pub epochs: Vec<(u64, u64)>,
pub frontier: Vec<Vec<u8>>,
pub frontier_coords: Vec<(u64, u32)>,
}
impl Clone for AlgView {
fn clone(&self) -> Self {
Self {
hasher: self.hasher.clone_box(),
epochs: self.epochs.clone(),
frontier: self.frontier.clone(),
frontier_coords: self.frontier_coords.clone(),
}
}
}
impl AlgView {
#[must_use]
pub fn is_active(&self) -> bool {
self.epochs.last().is_some_and(|&(_, end)| end == u64::MAX)
}
#[must_use]
pub fn is_active_at(&self, i: u64) -> bool {
self.epochs
.iter()
.any(|&(start, end)| start <= i && i < end)
}
#[must_use]
pub fn tree_size(&self, global_size: u64) -> u64 {
if self.is_active() {
global_size
} else {
self.epochs.last().map_or(0, |&(_, end)| end)
}
}
#[must_use]
pub fn first_activation(&self) -> u64 {
self.epochs.first().map_or(0, |&(start, _)| start)
}
#[must_use]
pub fn active_range(&self, lo: u64, hi: u64) -> bool {
self.epochs
.iter()
.any(|&(start, end)| start < hi && end > lo)
}
#[must_use]
pub fn fully_active(&self, lo: u64, hi: u64) -> bool {
self.epochs
.iter()
.any(|&(start, end)| start <= lo && hi <= end)
}
#[must_use]
pub fn effective_size_at(&self, size: u64) -> u64 {
if size <= self.first_activation() {
size
} else if self.is_active_at(size - 1) {
size
} else {
self.epochs
.iter()
.filter(|&&(_, end)| end < size)
.map(|&(_, end)| end)
.max()
.unwrap_or(0)
}
}
#[must_use]
pub fn epochs_at(&self, size: u64) -> Option<Vec<(u64, u64)>> {
let mut out = Vec::new();
for &(start, end) in &self.epochs {
if start > size {
break;
}
if end > size {
out.push((start, u64::MAX));
break;
}
out.push((start, end));
}
if out.is_empty() { None } else { Some(out) }
}
}
#[must_use]
pub fn compute_root(view: &AlgView, k: usize) -> Vec<u8> {
bag_peaks(view.hasher.as_ref(), &view.frontier, k as u64)
}
pub fn carry<E>(
view: &mut AlgView,
alg_id: u64,
digest: Vec<u8>,
count: u64,
arity: u64,
out_nodes: &mut Vec<(u64, u64, u32, Vec<u8>)>,
) -> Result<(), E> {
view.frontier.push(digest);
view.frontier_coords.push((count, 0));
let merges = reduction_count(count, arity);
for _ in 0..merges {
let mut children = Vec::with_capacity(arity as usize);
let mut coords = Vec::with_capacity(arity as usize);
for _ in 0..arity as usize {
children.push(view.frontier.pop().ok_or_else(|| Error::Corrupted {
alg_id,
reason: "frontier stack underflow during reduction".to_string(),
})?);
coords.push(view.frontier_coords.pop().ok_or_else(|| Error::Corrupted {
alg_id,
reason: "frontier_coords stack underflow during reduction".to_string(),
})?);
}
children.reverse();
coords.reverse();
let child_refs: Vec<&[u8]> = children.iter().map(|c| c.as_slice()).collect();
let parent = nary_mr(view.hasher.as_ref(), &child_refs);
let parent_left_index = coords[0].0;
let parent_height = coords[0].1 + 1;
if parent != view.hasher.null() {
out_nodes.push((alg_id, parent_left_index, parent_height, parent.clone()));
}
view.frontier.push(parent);
view.frontier_coords
.push((parent_left_index, parent_height));
}
Ok(())
}
pub async fn get_node_hash<R: NodeReader>(
reader: &R,
view: &AlgView,
alg_id: u64,
left: u64,
height: u32,
arity: u64,
) -> Result<Vec<u8>, R::Error> {
let cap = match arity.checked_pow(height) {
Some(c) => c,
None => return Ok(view.hasher.null()),
};
let limit = match left.checked_add(cap) {
Some(val) => val,
None => return Ok(view.hasher.null()),
};
if !view.active_range(left, limit) {
return Ok(view.hasher.null());
}
if let Some(hash) = reader
.get_node(alg_id, left, height)
.await
.map_err(Error::Read)?
{
let expected = view.hasher.null().len();
if hash.len() != expected {
return Err(Error::Corrupted {
alg_id,
reason: format!(
"node at left {left} height {height} has wrong digest length: expected \
{expected}, got {}",
hash.len()
),
});
}
Ok(hash)
} else {
Err(Error::Corrupted {
alg_id,
reason: format!(
"missing internal node for algorithm {alg_id} at left {left} height {height}"
),
})
}
}
pub async fn inclusion_proof<R: NodeReader>(
reader: &R,
view: &AlgView,
alg_id: u64,
index: u64,
tree_size: u64,
global_size: u64,
arity: u64,
) -> Result<Option<crate::consistency::InclusionProof>, R::Error> {
let max_size = view.tree_size(global_size);
if tree_size == 0 || index >= tree_size || tree_size > max_size {
return Ok(None);
}
let k = arity;
let coords = frontier_for_size(tree_size, k);
let mut target_f_idx = None;
for (f_idx, &(left, height)) in coords.iter().enumerate() {
let cap = k.pow(height);
if index >= left && index < left + cap {
target_f_idx = Some((f_idx, left, height));
break;
}
}
let (f_idx, left, height) = match target_f_idx {
Some(val) => val,
None => return Ok(None),
};
let mut path = Vec::new();
log_level_bisection_path_to_height(
reader, view, alg_id, left, height, index, 0, arity, &mut path,
)
.await?;
path.reverse();
let mut peaks = Vec::with_capacity(coords.len());
for &(l, h) in &coords {
let hash = get_node_hash(reader, view, alg_id, l, h, arity).await?;
peaks.push(hash);
}
let bag = bag_path(&peaks, f_idx, view.hasher.as_ref(), arity);
path.extend(bag);
debug_assert!(
mountain_skeleton(k, tree_size, index).is_some_and(|skeleton| {
skeleton.len() == path.len()
&& path.iter().zip(skeleton.iter()).all(|(step, shape)| {
step.position == shape.position && step.siblings.len() == shape.sibling_count
})
}),
"generated inclusion proof must match the canonical mountain skeleton"
);
Ok(Some(crate::consistency::InclusionProof { path }))
}
pub async fn leaf_proof<R: NodeReader>(
reader: &R,
view: &AlgView,
alg_id: u64,
index: u64,
tree_size: u64,
global_size: u64,
arity: u64,
) -> Result<Option<spine::LeafProof>, R::Error> {
let Some(proof) =
inclusion_proof(reader, view, alg_id, index, tree_size, global_size, arity).await?
else {
return Ok(None);
};
let leaf_hash = get_node_hash(reader, view, alg_id, index, 0, arity).await?;
Ok(Some(spine::LeafProof::new(
leaf_hash, index, tree_size, arity, proof.path,
)))
}
#[allow(clippy::too_many_arguments)]
pub async fn consistency_proof<R: NodeReader>(
reader: &R,
view: &AlgView,
alg_id: u64,
old_size: u64,
new_size: u64,
global_size: u64,
arity: u64,
) -> Result<Option<crate::consistency::ConsistencyProof>, R::Error> {
let max_size = view.tree_size(global_size);
if old_size == 0 || old_size >= new_size || new_size > max_size {
return Ok(None);
}
let k = arity;
let old_coords = frontier_for_size(old_size, k);
let &(boundary_left, boundary_height) = old_coords.last().ok_or_else(|| Error::Corrupted {
alg_id,
reason: "empty old_coords for non-zero old_size".to_string(),
})?;
let start_hash =
get_node_hash(reader, view, alg_id, boundary_left, boundary_height, arity).await?;
let new_coords = frontier_for_size(new_size, k);
let mut target_new_f_idx = None;
for (f_idx, &(new_left, new_height)) in new_coords.iter().enumerate() {
let cap = k.pow(new_height);
if boundary_left >= new_left && boundary_left < new_left + cap {
target_new_f_idx = Some((f_idx, new_left, new_height));
break;
}
}
let (f_idx, left, height) = match target_new_f_idx {
Some(val) => val,
None => return Ok(None),
};
if height < boundary_height {
return Ok(None);
}
let mut peak_path = Vec::new();
log_level_bisection_path_to_height(
reader,
view,
alg_id,
left,
height,
boundary_left,
boundary_height,
arity,
&mut peak_path,
)
.await?;
peak_path.reverse();
let mut new_peaks = Vec::with_capacity(new_coords.len());
for &(l, h) in &new_coords {
let hash = get_node_hash(reader, view, alg_id, l, h, arity).await?;
new_peaks.push(hash);
}
debug_assert!(
mountain_skeleton(k, new_size, boundary_left).is_some_and(|skeleton| {
let (bh, nh) = (boundary_height as usize, height as usize);
skeleton.len() >= nh
&& skeleton[bh..nh].len() == peak_path.len()
&& peak_path
.iter()
.zip(&skeleton[bh..nh])
.all(|(step, shape)| {
step.position == shape.position
&& step.siblings.len() == shape.sibling_count
})
}),
"generated consistency peak_path must match the canonical mountain skeleton"
);
Ok(Some(crate::consistency::ConsistencyProof {
boundary_hash: start_hash,
peak_path,
new_peaks,
split_index: f_idx,
}))
}
#[allow(clippy::too_many_arguments)]
async fn log_level_bisection_path_to_height<R: NodeReader>(
reader: &R,
view: &AlgView,
alg_id: u64,
left_index: u64,
height: u32,
target_index: u64,
target_height: u32,
arity: u64,
path: &mut Vec<ProofStep>,
) -> Result<(), R::Error> {
let mut curr_left = left_index;
let mut curr_height = height;
let k = arity;
while curr_height > target_height {
let child_capacity = k.pow(curr_height - 1);
let child_idx = (target_index - curr_left) / child_capacity;
let mut siblings = Vec::with_capacity(arity as usize - 1);
for j in 0..arity as usize {
let j_u64 = j as u64;
if j_u64 == child_idx {
continue;
}
let c_left = curr_left + j_u64 * child_capacity;
let hash = get_node_hash(reader, view, alg_id, c_left, curr_height - 1, arity).await?;
siblings.push(hash);
}
path.push(ProofStep {
siblings,
position: child_idx as usize,
});
curr_left += child_idx * child_capacity;
curr_height -= 1;
}
Ok(())
}
pub async fn peak_at<R: NodeReader>(
reader: &R,
view: &AlgView,
alg_id: u64,
left: u64,
height: u32,
arity: u64,
) -> Result<Vec<u8>, R::Error> {
get_node_hash(reader, view, alg_id, left, height, arity).await
}
pub async fn root_for_at<R: NodeReader>(
reader: &R,
view: &AlgView,
alg_id: u64,
size: u64,
arity: u64,
) -> Result<Vec<u8>, R::Error> {
let alg_size = view.effective_size_at(size);
if alg_size == 0 {
return Ok(view.hasher.empty());
}
let coords = frontier_for_size(alg_size, arity);
let mut frontier = Vec::with_capacity(coords.len());
for &(left, height) in &coords {
let hash = get_node_hash(reader, view, alg_id, left, height, arity).await?;
frontier.push(hash);
}
Ok(bag_peaks(view.hasher.as_ref(), &frontier, arity))
}
pub async fn frontier_peaks<R: NodeReader>(
reader: &R,
view: &AlgView,
alg_id: u64,
size: u64,
arity: u64,
) -> Result<Vec<Vec<u8>>, R::Error> {
let coords = frontier_for_size(size, arity);
let mut peaks = Vec::with_capacity(coords.len());
for &(left, height) in &coords {
peaks.push(peak_at(reader, view, alg_id, left, height, arity).await?);
}
Ok(peaks)
}
pub fn validate_epochs<E>(alg_id: u64, epochs: &[(u64, u64)], global_size: u64) -> Result<(), E> {
if epochs.is_empty() {
return Err(Error::Corrupted {
alg_id,
reason: "epoch sequence is empty".to_string(),
});
}
let mut last_end = 0;
for (i, &(start, end)) in epochs.iter().enumerate() {
if start > end {
return Err(Error::Corrupted {
alg_id,
reason: format!("epoch start {start} exceeds end {end}"),
});
}
if start < last_end {
return Err(Error::Corrupted {
alg_id,
reason: format!("epoch start {start} is less than prior end {last_end}"),
});
}
if end != u64::MAX && end > global_size {
return Err(Error::Corrupted {
alg_id,
reason: format!("epoch end {end} exceeds global size {global_size}"),
});
}
if end == u64::MAX && i != epochs.len() - 1 {
return Err(Error::Corrupted {
alg_id,
reason: "open epoch (end = u64::MAX) is not the final entry".to_string(),
});
}
last_end = end;
}
if let Some(&(start, end)) = epochs.last() {
if end == u64::MAX && start > global_size {
return Err(Error::Corrupted {
alg_id,
reason: format!("active epoch start {start} exceeds global size {global_size}"),
});
}
}
Ok(())
}
pub async fn reconstruct_view<R: NodeReader>(
reader: &R,
alg_id: u64,
hasher: Box<dyn Hasher>,
epochs: &[(u64, u64)],
global_size: u64,
arity: u64,
) -> Result<AlgView, R::Error> {
validate_epochs(alg_id, epochs, global_size)?;
let is_active = epochs.last().is_some_and(|&(_, end)| end == u64::MAX);
let tree_size = if is_active {
global_size
} else {
epochs.last().map_or(0, |&(_, end)| end)
};
let mut view = AlgView {
hasher,
epochs: epochs.to_vec(),
frontier: Vec::new(),
frontier_coords: Vec::new(),
};
if tree_size == 0 {
return Ok(view);
}
let coords = frontier_for_size(tree_size, arity);
let mut frontier = Vec::with_capacity(coords.len());
for &(left, height) in &coords {
let cap = arity.pow(height);
let hash = if !view.active_range(left, left + cap) {
view.hasher.null()
} else {
reader
.get_node(alg_id, left, height)
.await
.map_err(Error::Read)?
.ok_or_else(|| Error::Corrupted {
alg_id,
reason: format!(
"missing frontier node for algorithm {alg_id} at left {left} height \
{height}"
),
})?
};
frontier.push(hash);
}
view.frontier = frontier;
view.frontier_coords = coords;
Ok(view)
}
#[allow(clippy::type_complexity)]
#[allow(clippy::too_many_arguments)]
pub fn reconstruct_subtree_root<'a, R>(
reader: &'a R,
alg_id: u64,
view: &'a AlgView,
lo: u64,
hi: u64,
k: u64,
kind: LogKind,
store_mixed: bool,
) -> std::pin::Pin<
Box<
dyn std::future::Future<Output = Result<(Vec<u8>, Vec<(u64, u32, Vec<u8>)>), R::Error>>
+ Send
+ 'a,
>,
>
where
R: NodeReader + Sync + 'a,
{
Box::pin(async move {
let size = hi - lo;
if size == 0 {
return Ok((view.hasher.empty(), Vec::new()));
}
if size == 1 {
if view.is_active_at(lo) {
if kind == LogKind::Subtree {
if let Some(hash) = reader.get_node(alg_id, lo, 0).await.map_err(Error::Read)? {
let expected = view.hasher.null().len();
if hash.len() != expected {
return Err(Error::Corrupted {
alg_id,
reason: format!(
"subtree node at left {lo} height 0 has wrong digest length: \
expected {expected}, got {}",
hash.len()
),
});
}
return Ok((hash, Vec::new()));
} else {
return Ok((view.hasher.null(), Vec::new()));
}
} else {
let data = reader.get_leaf(lo).await.map_err(Error::Read)?;
return Ok((view.hasher.leaf(&data), Vec::new()));
}
}
return Ok((view.hasher.null(), Vec::new()));
}
if !view.active_range(lo, hi) {
return Ok((view.hasher.null(), Vec::new()));
}
let is_power_of_k = {
let mut temp = size;
while temp % k == 0 {
temp /= k;
}
temp == 1
};
if is_power_of_k {
let height = {
let mut h = 0;
let mut temp = size;
while temp > 1 {
temp /= k;
h += 1;
}
h as u32
};
if view.fully_active(lo, hi) {
if let Some(hash) = reader
.get_node(alg_id, lo, height)
.await
.map_err(Error::Read)?
{
let expected = view.hasher.null().len();
if hash.len() != expected {
return Err(Error::Corrupted {
alg_id,
reason: format!(
"node at left {lo} height {height} has wrong digest length: \
expected {expected}, got {}",
hash.len()
),
});
}
return Ok((hash, Vec::new()));
}
}
let child_size = size / k;
let mut child_hashes = Vec::with_capacity(k as usize);
let mut mixed_nodes = Vec::new();
for j in 0..k {
let c_lo = lo + j * child_size;
let c_hi = lo + (j + 1) * child_size;
let (child_hash, child_mixed) = reconstruct_subtree_root(
reader,
alg_id,
view,
c_lo,
c_hi,
k,
kind,
store_mixed,
)
.await?;
child_hashes.push(child_hash);
mixed_nodes.extend(child_mixed);
}
let child_refs: Vec<&[u8]> = child_hashes.iter().map(|c| c.as_slice()).collect();
let hash = nary_mr(view.hasher.as_ref(), &child_refs);
if store_mixed {
mixed_nodes.push((lo, height, hash.clone()));
}
Ok((hash, mixed_nodes))
} else {
let coords = frontier_for_size(size, k);
let mut component_hashes = Vec::with_capacity(coords.len());
let mut mixed_nodes = Vec::new();
for &(part_left, part_height) in &coords {
let cap = k.pow(part_height);
let c_lo = lo + part_left;
let c_hi = c_lo + cap;
let (part_root, part_mixed) = reconstruct_subtree_root(
reader,
alg_id,
view,
c_lo,
c_hi,
k,
kind,
store_mixed,
)
.await?;
component_hashes.push(part_root);
mixed_nodes.extend(part_mixed);
}
let root = bag_peaks(view.hasher.as_ref(), &component_hashes, k);
Ok((root, mixed_nodes))
}
})
}
#[cfg(test)]
mod tests {
use sha2::Digest as _;
use super::*;
#[derive(Debug)]
struct TestHasher;
impl Hasher for TestHasher {
fn leaf(&self, data: &[u8]) -> Vec<u8> {
sha2::Sha256::digest(data).to_vec()
}
fn node(&self, children: &[&[u8]]) -> Vec<u8> {
let mut h = sha2::Sha256::new();
for c in children {
h.update(c);
}
h.finalize().to_vec()
}
fn empty(&self) -> Vec<u8> {
sha2::Sha256::digest(b"").to_vec()
}
fn hash(&self, data: &[u8]) -> Vec<u8> {
sha2::Sha256::digest(data).to_vec()
}
fn clone_box(&self) -> Box<dyn Hasher> {
Box::new(TestHasher)
}
}
#[test]
fn epochs_at_snapshot_clamping() {
let view = AlgView {
hasher: Box::new(TestHasher),
epochs: vec![(2, 5), (7, 12), (15, u64::MAX)],
frontier: Vec::new(),
frontier_coords: Vec::new(),
};
assert_eq!(view.epochs_at(1), None);
assert_eq!(view.epochs_at(3), Some(vec![(2, u64::MAX)]));
assert_eq!(view.epochs_at(5), Some(vec![(2, 5)]));
assert_eq!(view.epochs_at(6), Some(vec![(2, 5)]));
assert_eq!(view.epochs_at(7), Some(vec![(2, 5), (7, u64::MAX)]));
assert_eq!(
view.epochs_at(20),
Some(vec![(2, 5), (7, 12), (15, u64::MAX)])
);
}
}