diff --git a/chain/src/txhashset/segmenter.rs b/chain/src/txhashset/segmenter.rs index ed6a2aee22..345a25829b 100644 --- a/chain/src/txhashset/segmenter.rs +++ b/chain/src/txhashset/segmenter.rs @@ -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, @@ -93,7 +93,7 @@ impl Segmenter { ) -> Result<(Segment, 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", @@ -113,9 +113,10 @@ impl Segmenter { id: SegmentIdentifier, ) -> Result<(Segment, Hash), Error> { let now = Instant::now(); + let bitmap = self.bitmap_snapshot.as_bitmap().ok(); 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", @@ -132,9 +133,10 @@ impl Segmenter { /// Create a rangeproof segment. pub fn rangeproof_segment(&self, id: SegmentIdentifier) -> Result, Error> { let now = Instant::now(); + let bitmap = self.bitmap_snapshot.as_bitmap().ok(); 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, diff --git a/chain/tests/bitmap_segment.rs b/chain/tests/bitmap_segment.rs index 960625baf7..baa64346bb 100644 --- a/chain/tests/bitmap_segment.rs +++ b/chain/tests/bitmap_segment.rs @@ -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()); diff --git a/core/src/core/pmmr/segment.rs b/core/src/core/pmmr/segment.rs index 8df8b8f989..1794c67768 100644 --- a/core/src/core/pmmr/segment.rs +++ b/core/src/core/pmmr/segment.rs @@ -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( segment_id: SegmentIdentifier, pmmr: &ReadonlyPMMR<'_, U, B>, - prunable: bool, + bitmap: Option<&Bitmap>, ) -> Result where U: PMMRable, @@ -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 { + if let Some(hash) = pmmr.get_from_file(pos0) { + segment.hashes.push(hash); + segment.hash_pos.push(pos0); + } + } } } } @@ -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 +} diff --git a/core/tests/segment.rs b/core/tests/segment.rs index cb81402fac..bc000055e1 100644 --- a/core/tests/segment.rs +++ b/core/tests/segment.rs @@ -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 diff --git a/store/tests/segment.rs b/store/tests/segment.rs index a24c1b8c5c..49ec24f3d6 100644 --- a/store/tests/segment.rs +++ b/store/tests/segment.rs @@ -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() @@ -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() @@ -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() @@ -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() @@ -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) @@ -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 @@ -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); @@ -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!( @@ -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 @@ -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 @@ -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!( @@ -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 @@ -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(); @@ -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!( @@ -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::::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);