Make the sparse PIBD segment optimization (#3820) safe to apply

chain/src/txhashset/desegmenter.rs

@@ -486,9 +486,13 @@ impl Desegmenter {
 				self.bitmap_mmr_size,
 				self.default_bitmap_segment_height,
 			);
-			// Advance iterator to next expected segment
+			// Advance iterator to next expected segment. `local_pmmr_size` is
+			// a count (positions 0..size-1 are present locally), so a segment
+			// whose last position equals it still contains new data and must
+			// be requested: with a strict `>` a single-leaf bitmap MMR (last
+			// position 0, local size 0) is never requested and sync stalls.
 			while let Some(id) = identifier_iter.next() {
-				if id.segment_pos_range(self.bitmap_mmr_size).1 > local_pmmr_size {
+				if id.segment_pos_range(self.bitmap_mmr_size).1 >= local_pmmr_size {
 					if !self.has_bitmap_segment_with_id(id) {
 						return_vec.push(SegmentTypeIdentifier::new(SegmentType::Bitmap, id));
 						if return_vec.len() >= max_elements {
@@ -672,18 +676,23 @@ impl Desegmenter {
 			"pibd_desegmenter - expected number of leaves in bitmap MMR: {}",
 			self.bitmap_mmr_leaf_count
 		);
-		// Total size of Bitmap PMMR
+		// Total size of Bitmap PMMR. `unwrap_or` evaluates its argument
+		// eagerly, so the fallback used to run - and panic on an empty peak
+		// set - whenever the bitmap MMR had a single leaf (output MMR with
+		// <= 1024 leaves). Evaluate it lazily and degrade to 0 instead of
+		// panicking when even the reduced peak set is empty.
 		self.bitmap_mmr_size =
 			1 + pmmr::peaks(pmmr::insertion_to_pmmr_index(self.bitmap_mmr_leaf_count))
 				.last()
-				.unwrap_or(
-					&(pmmr::peaks(pmmr::insertion_to_pmmr_index(
-						self.bitmap_mmr_leaf_count - 1,
+				.copied()
+				.unwrap_or_else(|| {
+					pmmr::peaks(pmmr::insertion_to_pmmr_index(
+						self.bitmap_mmr_leaf_count.saturating_sub(1),
 					))
 					.last()
-					.unwrap()),
-				)
-				.clone();
+					.copied()
+					.unwrap_or(0)
+				});
 
 		trace!(
 			"pibd_desegmenter - expected size of bitmap MMR: {}",

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<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",
@@ -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();
 		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<Segment<RangeProof>, 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,

chain/src/txhashset/txhashset.rs

@@ -1426,6 +1426,22 @@ impl<'a> Extension<'a> {
 								.rewind(0, &Bitmap::new())
 								.map_err(&Error::TxHashSetErr)?;
 						}
+						// A pruned-subtree hash must cover a region that is
+						// entirely absent locally. If any position under
+						// pos0 was already applied (leftmost < current
+						// size), collapsing here would silently discard
+						// already-applied leaves/hashes and desynchronize
+						// the prune list from the physical hash file -
+						// reject instead of corrupting local state.
+						let leftmost = pmmr::bintree_leftmost(pos0);
+						if leftmost < self.output_pmmr.size {
+							return Err(Error::InvalidSegment(format!(
+								"output segment hash at pos {} would overwrite {} already-applied position(s) starting at {} (sparse/dense segment mismatch)",
+								pos0,
+								self.output_pmmr.size - leftmost,
+								leftmost
+							)));
+						}
 						self.output_pmmr
 							.push_pruned_subtree(hashes[idx], pos0)
 							.map_err(&Error::TxHashSetErr)?;
@@ -1470,6 +1486,18 @@ impl<'a> Extension<'a> {
 								.rewind(0, &Bitmap::new())
 								.map_err(&Error::TxHashSetErr)?;
 						}
+						// See apply_output_segment: reject a collapse that
+						// would overwrite already-applied positions instead
+						// of silently corrupting the prune list.
+						let leftmost = pmmr::bintree_leftmost(pos0);
+						if leftmost < self.rproof_pmmr.size {
+							return Err(Error::InvalidSegment(format!(
+								"rangeproof segment hash at pos {} would overwrite {} already-applied position(s) starting at {} (sparse/dense segment mismatch)",
+								pos0,
+								self.rproof_pmmr.size - leftmost,
+								leftmost
+							)));
+						}
 						self.rproof_pmmr
 							.push_pruned_subtree(hashes[idx], pos0)
 							.map_err(&Error::TxHashSetErr)?;

chain/tests/bitmap_segment.rs

@@ -65,7 +65,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());

core/src/core/pmmr/segment.rs

@@ -340,10 +340,23 @@ where
 	T: Readable + Writeable + Debug,
 {
 	/// Generate a segment from a PMMR
+	/// If bitmap is provided, only hashes at pruning boundaries are included.
+	///
+	/// Hash boundaries are computed via the same bottom-up postorder walk the
+	/// receiver-side reconstruction (`root`) performs, tracking which
+	/// positions are already derivable ("known") from the selected leaves,
+	/// instead of an independent bitmap-only heuristic. A position becomes
+	/// known if both children are known; if exactly one child is known, the
+	/// *other* child's hash is required and emitted (mirroring the
+	/// `(None, Some)` / `(Some, None)` lookups in `root`), after which the
+	/// parent is known too. This guarantees the produced segment is exactly
+	/// what `root` and `TxHashSet::apply_*_segment` expect: no position is
+	/// ever described twice (as both leaf data and part of a collapsed
+	/// ancestor hash), and no hash needed for reconstruction is missing.
 	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>,
@@ -356,25 +369,103 @@ where
 			return Err(SegmentError::NonExistent);
 		}
 
-		// Fill leaf data and hashes
 		let (segment_first_pos, segment_last_pos) = segment.segment_pos_range(mmr_size);
+
+		// Pass 1: select leaf data - present on disk, and (for prunable MMRs)
+		// required by the bitmap or its sibling, or the final uneven leaf:
+		// the exact set the receiver-side `root` consumes.
 		for pos0 in segment_first_pos..=segment_last_pos {
-			if pmmr::is_leaf(pos0) {
+			if !pmmr::is_leaf(pos0) {
+				continue;
+			}
+			let include = match bitmap {
+				None => true,
+				Some(bm) => {
+					let idx_1 = pmmr::n_leaves(pos0 + 1) - 1;
+					let idx_2 = if pmmr::is_left_sibling(pos0) {
+						idx_1 + 1
+					} else {
+						idx_1 - 1
+					};
+					bm.contains(idx_1 as u32) || bm.contains(idx_2 as u32) || pos0 == mmr_size - 1
+				}
+			};
+			if include {
 				if let Some(data) = pmmr.get_data_from_file(pos0) {
 					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);
+		}
+
+		// Pass 2: bottom-up walk mirroring `root`, deciding exactly which
+		// interior hashes are needed to reconstruct the root from the leaf
+		// set chosen above.
+		if bitmap.is_some() {
+			let leaf_set: std::collections::HashSet<u64> =
+				segment.leaf_pos.iter().copied().collect();
+			let mut needed = Vec::<(u64, Hash)>::new();
+			// (pos0, known) - `known` means the receiver can derive this
+			// position's hash from what the segment ships below it.
+			let mut stack =
+				Vec::<(u64, bool)>::with_capacity(2 * (segment.identifier.height as usize) + 1);
+			for pos0 in segment_first_pos..=segment_last_pos {
+				let height = pmmr::bintree_postorder_height(pos0);
+				let is_known = if height == 0 {
+					leaf_set.contains(&pos0)
+				} else {
+					let (_, right_known) = stack.pop().unwrap();
+					let (_, left_known) = stack.pop().unwrap();
+					match (left_known, right_known) {
+						(true, true) => true,
+						(false, false) => false,
+						(known_l, _known_r) => {
+							let left_child_pos = 1 + pos0 - (1 << height);
+							let right_child_pos = pos0;
+							let need_pos = if known_l {
+								right_child_pos - 1
+							} else {
+								left_child_pos - 1
+							};
+							if let Some(hash) = pmmr.get_from_file(need_pos) {
+								needed.push((need_pos, hash));
+								true
+							} else {
+								// The needed child hash is unavailable (its
+								// subtree is compacted beyond this level):
+								// leave the parent unknown so the requirement
+								// propagates upward, or ultimately resolves
+								// via the whole-segment fallback below.
+								false
+							}
+						}
+					}
+				};
+				stack.push((pos0, is_known));
+			}
+			// Any subtree root still unknown has no parent inside the
+			// segment to demand it: the segment root of a fully pruned full
+			// segment, or a fully pruned peak in the final (partial)
+			// segment. The receiver looks these up directly (`root` bags
+			// pruned in-segment peaks; `first_unpruned_parent` probes the
+			// segment root before climbing), so ship them if we have them.
+			for (pos0, known) in stack {
+				if !known {
+					if let Some(hash) = pmmr.get_from_file(pos0) {
+						needed.push((pos0, hash));
+					}
 				}
 			}
+			// Serialization requires strictly increasing hash positions and
+			// the walk can demand a left child after a deeper right-subtree
+			// sibling, so sort before finalizing.
+			needed.sort_unstable_by_key(|&(pos0, _)| pos0);
+			for (pos0, hash) in needed {
+				segment.hash_pos.push(pos0);
+				segment.hashes.push(hash);
+			}
 		}
 
 		let mut start_pos = None;

core/tests/segment.rs

@@ -41,7 +41,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