Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions chain/src/txhashset/segmenter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ impl Segmenter {
let now = Instant::now();
let txhashset = self.txhashset.read();
let kernel_pmmr = txhashset.kernel_pmmr_at(&self.header);
let segment = Segment::from_pmmr(id, &kernel_pmmr, false)?;
let segment = Segment::from_pmmr(id, &kernel_pmmr, None)?;
debug!(
"kernel_segment: id: ({}, {}), leaves: {}, hashes: {}, proof hashes: {}, took {}ms",
segment.id().height,
Expand Down Expand Up @@ -93,7 +93,7 @@ impl Segmenter {
) -> Result<(Segment<BitmapChunk>, Hash), Error> {
let now = Instant::now();
let bitmap_pmmr = self.bitmap_snapshot.readonly_pmmr();
let segment = Segment::from_pmmr(id, &bitmap_pmmr, false)?;
let segment = Segment::from_pmmr(id, &bitmap_pmmr, None)?;
let output_root = self.output_root()?;
debug!(
"bitmap_segment: id: ({}, {}), leaves: {}, hashes: {}, proof hashes: {}, took {}ms",
Expand All @@ -113,9 +113,10 @@ impl Segmenter {
id: SegmentIdentifier,
) -> Result<(Segment<OutputIdentifier>, Hash), Error> {
let now = Instant::now();
let bitmap = self.bitmap_snapshot.as_bitmap().ok();
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we avoid .ok() here? If building the bitmap fails, I think this should return an error

let txhashset = self.txhashset.read();
let output_pmmr = txhashset.output_pmmr_at(&self.header);
let segment = Segment::from_pmmr(id, &output_pmmr, true)?;
let segment = Segment::from_pmmr(id, &output_pmmr, bitmap.as_ref())?;
let bitmap_root = self.bitmap_root()?;
debug!(
"output_segment: id: ({}, {}), leaves: {}, hashes: {}, proof hashes: {}, took {}ms",
Expand All @@ -132,9 +133,10 @@ impl Segmenter {
/// Create a rangeproof segment.
pub fn rangeproof_segment(&self, id: SegmentIdentifier) -> Result<Segment<RangeProof>, Error> {
let now = Instant::now();
let bitmap = self.bitmap_snapshot.as_bitmap().ok();
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same here, I think bitmap errors should be returned instead of being converted to None.

let txhashset = self.txhashset.read();
let pmmr = txhashset.rangeproof_pmmr_at(&self.header);
let segment = Segment::from_pmmr(id, &pmmr, true)?;
let segment = Segment::from_pmmr(id, &pmmr, bitmap.as_ref())?;
debug!(
"rangeproof_segment: id: ({}, {}), leaves: {}, hashes: {}, proof hashes: {}, took {}ms",
segment.id().height,
Expand Down
2 changes: 1 addition & 1 deletion chain/tests/bitmap_segment.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ fn test_roundtrip(entries: usize) {
.unwrap();

let mmr = accumulator.readonly_pmmr();
let segment = Segment::from_pmmr(identifier, &mmr, false).unwrap();
let segment = Segment::from_pmmr(identifier, &mmr, None).unwrap();

// Convert to `BitmapSegment`
let bms = BitmapSegment::from(segment.clone());
Expand Down
51 changes: 44 additions & 7 deletions core/src/core/pmmr/segment.rs
Original file line number Diff line number Diff line change
Expand Up @@ -307,10 +307,11 @@ where
T: Readable + Writeable + Debug,
{
/// Generate a segment from a PMMR
/// If bitmap is provided, only hashes at pruning boundaries are included.
pub fn from_pmmr<U, B>(
segment_id: SegmentIdentifier,
pmmr: &ReadonlyPMMR<'_, U, B>,
prunable: bool,
bitmap: Option<&Bitmap>,
) -> Result<Self, SegmentError>
where
U: PMMRable<E = T>,
Expand All @@ -331,15 +332,40 @@ where
segment.leaf_data.push(data);
segment.leaf_pos.push(pos0);
continue;
} else if !prunable {
} else if bitmap.is_none() {
return Err(SegmentError::MissingLeaf(pos0));
}
}
// TODO: optimize, no need to send every intermediary hash
if prunable {
if let Some(hash) = pmmr.get_from_file(pos0) {
segment.hashes.push(hash);
segment.hash_pos.push(pos0);
if let Some(bm) = bitmap {
// Only include hash if this subtree is fully pruned
// AND the sibling subtree is NOT fully pruned (pruning boundary)
if subtree_fully_pruned(pos0, bm, mmr_size) {
// Find sibling position
let height = pmmr::bintree_postorder_height(pos0);
let subtree_size = (1u64 << (height + 1)) - 1;
let sibling_pos0 = if pmmr::is_left_sibling(pos0) {
// Right sibling is at pos0 + size of this subtree
Some(pos0 + subtree_size)
} else {
// Left sibling: go back by sibling's subtree size
pos0.checked_sub(subtree_size)
};
// Need hash if sibling exists in segment and is not fully pruned
let need_hash = match sibling_pos0 {
Some(sib) if sib >= segment_first_pos && sib <= segment_last_pos => {
!subtree_fully_pruned(sib, bm, mmr_size)
}
_ => {
// Sibling outside segment or underflow - may need hash for proof
true
}
};
if need_hash {
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If need_hash is true, silently skipping a missing hash can produce a segment that later fails validation with MissingHash or Invalid Root. Should this return an error instead?

like:

let hash = pmmr
    .get_from_file(pos0)
    .ok_or_else(|| SegmentError::MissingHash(pos0))?;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This also matches the failure pattern you observed, @ardocrat.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll check this out

if let Some(hash) = pmmr.get_from_file(pos0) {
segment.hashes.push(hash);
segment.hash_pos.push(pos0);
}
}
}
}
}
Expand Down Expand Up @@ -842,3 +868,14 @@ impl Writeable for SegmentProof {
Ok(())
}
}

/// Check if a subtree rooted at pos0 is fully pruned (no unspent leaves in bitmap)
fn subtree_fully_pruned(pos0: u64, bitmap: &Bitmap, mmr_size: u64) -> bool {
let leftmost = pmmr::bintree_leftmost(pos0);
let rightmost = pmmr::bintree_rightmost(pos0);
let n_leaves = pmmr::n_leaves(mmr_size);
let start_leaf = pmmr::n_leaves(leftmost + 1).saturating_sub(1);
let end_leaf = min(pmmr::n_leaves(rightmost + 1), n_leaves);
// If any leaf in range is in bitmap (unspent), subtree is not fully pruned
bitmap.range_cardinality(start_leaf as u32..end_leaf as u32) == 0
}
2 changes: 1 addition & 1 deletion core/tests/segment.rs
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ fn test_unprunable_size(height: u8, n_leaves: u32) {

for idx in 0..n_segments {
let id = SegmentIdentifier { height, idx };
let segment = Segment::from_pmmr(id, &mmr, false).unwrap();
let segment = Segment::from_pmmr(id, &mmr, None).unwrap();
println!(
"\n\n>>>>>>> N_LEAVES = {}, LAST_POS = {}, SEGMENT = {}:\n{:#?}",
n_leaves, last_pos, idx, segment
Expand Down
35 changes: 18 additions & 17 deletions store/tests/segment.rs
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ fn prunable_mmr() {

// Validate a segment before any pruning
let mmr = ReadonlyPMMR::at(&mut ba, last_pos);
let segment = Segment::from_pmmr(id, &mmr, true).unwrap();
let segment = Segment::from_pmmr(id, &mmr, Some(&bitmap)).unwrap();
assert_eq!(
segment.root(last_pos, Some(&bitmap)).unwrap().unwrap(),
mmr.get_hash(29).unwrap()
Expand All @@ -70,7 +70,7 @@ fn prunable_mmr() {

// Validate
let mmr = ReadonlyPMMR::at(&mut ba, last_pos);
let segment = Segment::from_pmmr(id, &mmr, true).unwrap();
let segment = Segment::from_pmmr(id, &mmr, Some(&bitmap)).unwrap();
assert_eq!(
segment.root(last_pos, Some(&bitmap)).unwrap().unwrap(),
mmr.get_hash(29).unwrap()
Expand All @@ -86,7 +86,7 @@ fn prunable_mmr() {

// Validate
let mmr = ReadonlyPMMR::at(&mut ba, last_pos);
let segment = Segment::from_pmmr(id, &mmr, true).unwrap();
let segment = Segment::from_pmmr(id, &mmr, Some(&bitmap)).unwrap();
assert_eq!(
segment.root(last_pos, Some(&bitmap)).unwrap().unwrap(),
mmr.get_hash(29).unwrap()
Expand All @@ -102,7 +102,7 @@ fn prunable_mmr() {

// Validate
let mmr = ReadonlyPMMR::at(&mut ba, last_pos);
let segment = Segment::from_pmmr(id, &mmr, true).unwrap();
let segment = Segment::from_pmmr(id, &mmr, Some(&bitmap)).unwrap();
assert_eq!(
segment.root(last_pos, Some(&bitmap)).unwrap().unwrap(),
mmr.get_hash(29).unwrap()
Expand All @@ -117,13 +117,13 @@ fn prunable_mmr() {
ba.sync().unwrap();

let mmr = ReadonlyPMMR::at(&mut ba, last_pos);
assert!(Segment::from_pmmr(id, &mmr, true).is_ok());
assert!(Segment::from_pmmr(id, &mmr, Some(&bitmap)).is_ok());

// Final segment is not full, test it before pruning
let id = SegmentIdentifier { height: 3, idx: 9 };

let mmr = ReadonlyPMMR::at(&mut ba, last_pos);
let segment = Segment::from_pmmr(id, &mmr, true).unwrap();
let segment = Segment::from_pmmr(id, &mmr, Some(&bitmap)).unwrap();
segment.validate(last_pos, Some(&bitmap), root).unwrap();

// Prune second and third to last leaves (a full peak in the MMR)
Expand All @@ -135,7 +135,7 @@ fn prunable_mmr() {

// Validate
let mmr = ReadonlyPMMR::at(&mut ba, last_pos);
let segment = Segment::from_pmmr(id, &mmr, true).unwrap();
let segment = Segment::from_pmmr(id, &mmr, Some(&bitmap)).unwrap();
segment.validate(last_pos, Some(&bitmap), root).unwrap();

// Prune final element
Expand All @@ -147,7 +147,7 @@ fn prunable_mmr() {

// Validate
let mmr = ReadonlyPMMR::at(&mut ba, last_pos);
let segment = Segment::from_pmmr(id, &mmr, true).unwrap();
let segment = Segment::from_pmmr(id, &mmr, Some(&bitmap)).unwrap();
segment.validate(last_pos, Some(&bitmap), root).unwrap();

std::mem::drop(ba);
Expand Down Expand Up @@ -185,7 +185,7 @@ fn pruned_segment() {
// Validate the empty segment 1
let id = SegmentIdentifier { height: 2, idx: 1 };
let mmr = ReadonlyPMMR::at(&mut ba, last_pos);
let segment = Segment::from_pmmr(id, &mmr, true).unwrap();
let segment = Segment::from_pmmr(id, &mmr, Some(&bitmap)).unwrap();
assert_eq!(segment.leaf_iter().count(), 0);
assert_eq!(segment.hash_iter().count(), 1);
assert_eq!(
Expand All @@ -206,7 +206,7 @@ fn pruned_segment() {

// Validate the empty segment 1 again
let mmr = ReadonlyPMMR::at(&mut ba, last_pos);
let segment = Segment::from_pmmr(id, &mmr, true).unwrap();
let segment = Segment::from_pmmr(id, &mmr, Some(&bitmap)).unwrap();
assert_eq!(segment.leaf_iter().count(), 0);
assert_eq!(segment.hash_iter().count(), 1);
// Since both 7 and 14 are now pruned, the first unpruned hash will be at 15
Expand All @@ -228,7 +228,7 @@ fn pruned_segment() {

// Validate the empty segment 1 again
let mmr = ReadonlyPMMR::at(&mut ba, last_pos);
let segment = Segment::from_pmmr(id, &mmr, true).unwrap();
let segment = Segment::from_pmmr(id, &mmr, Some(&bitmap)).unwrap();
assert_eq!(segment.leaf_iter().count(), 0);
assert_eq!(segment.hash_iter().count(), 1);
// Since both 15 and 30 are now pruned, the first unpruned hash will be at 31: the mmr root
Expand Down Expand Up @@ -260,7 +260,7 @@ fn pruned_segment() {
// Validate segment 4
let id = SegmentIdentifier { height: 2, idx: 4 };
let mmr = ReadonlyPMMR::at(&mut ba, last_pos);
let segment = Segment::from_pmmr(id, &mmr, true).unwrap();
let segment = Segment::from_pmmr(id, &mmr, Some(&bitmap)).unwrap();
assert_eq!(segment.leaf_iter().count(), 0);
assert_eq!(segment.hash_iter().count(), 1);
assert_eq!(
Expand All @@ -275,7 +275,7 @@ fn pruned_segment() {
// Segment 5 has 2 peaks
let id = SegmentIdentifier { height: 2, idx: 5 };
let mmr = ReadonlyPMMR::at(&mut ba, last_pos);
let segment = Segment::from_pmmr(id, &mmr, true).unwrap();
let segment = Segment::from_pmmr(id, &mmr, Some(&bitmap)).unwrap();
assert_eq!(segment.leaf_iter().count(), 3);
assert_eq!(
segment
Expand All @@ -297,7 +297,7 @@ fn pruned_segment() {

// Segment 5 should be unchanged
let mmr = ReadonlyPMMR::at(&mut ba, last_pos);
let segment = Segment::from_pmmr(id, &mmr, true).unwrap();
let segment = Segment::from_pmmr(id, &mmr, Some(&bitmap)).unwrap();
assert_eq!(segment, prev_segment);
segment.validate(last_pos, Some(&bitmap), root).unwrap();

Expand All @@ -310,7 +310,7 @@ fn pruned_segment() {

// Validate segment 5 again
let mmr = ReadonlyPMMR::at(&mut ba, last_pos);
let segment = Segment::from_pmmr(id, &mmr, true).unwrap();
let segment = Segment::from_pmmr(id, &mmr, Some(&bitmap)).unwrap();
assert_eq!(segment.leaf_iter().count(), 1);
assert_eq!(segment.hash_iter().count(), 1);
assert_eq!(
Expand Down Expand Up @@ -354,14 +354,15 @@ fn ser_round_trip() {

let mmr = ReadonlyPMMR::at(&ba, last_pos);
let id = SegmentIdentifier { height: 3, idx: 0 };
let segment = Segment::from_pmmr(id, &mmr, true).unwrap();
let segment = Segment::from_pmmr(id, &mmr, Some(&bitmap)).unwrap();

let mut cursor = Cursor::new(Vec::<u8>::new());
let mut writer = BinWriter::new(&mut cursor, ProtocolVersion(1));
Writeable::write(&segment, &mut writer).unwrap();
// With optimization, only 1 hash is needed (at pruning boundary) instead of 7
assert_eq!(
cursor.position(),
(9) + (8 + 7 * (8 + 32)) + (8 + 6 * (8 + 16)) + (8 + 2 * 32)
(9) + (8 + 1 * (8 + 32)) + (8 + 6 * (8 + 16)) + (8 + 2 * 32)
);
cursor.set_position(0);

Expand Down