[PIBD] PMMR Desegmenter Structure (Pt. 1) (#3667)

* initial commit of WIP pibd explorations

* correct calling for obtaining and validating first segment

* update test to properly iterate through each segment of the test pmmrs, validating each segment as it goes

* updated test to fully segment and validate PMMRs from compacted and uncompacted sample data. Also contains method of running test againt live chain data

* remove logger change

* change test file name

* change test file name

* change directory reference in test for CI

* add initial (experimental) structure for PIBD desegmenting

* move bitmap desegmentation logic into desegmenter

* added txhashset methods to apply pibd segments (note this only works for fully unpruned trees atm)

* change last_pos to mmr_size

* fix to pmmr::peaks call

* don't verify POW when copying headers

* prepare for commit of work thus far'

* update test paths

* few updates based on early review

chain/src/chain.rs

@@ -27,7 +27,7 @@ use crate::error::{Error, ErrorKind};
 use crate::pipe;
 use crate::store;
 use crate::txhashset;
-use crate::txhashset::{PMMRHandle, Segmenter, TxHashSet};
+use crate::txhashset::{Desegmenter, PMMRHandle, Segmenter, TxHashSet};
 use crate::types::{
 	BlockStatus, ChainAdapter, CommitPos, NoStatus, Options, Tip, TxHashsetWriteStatus,
 };
@@ -153,6 +153,7 @@ pub struct Chain {
 	txhashset: Arc<RwLock<txhashset::TxHashSet>>,
 	header_pmmr: Arc<RwLock<txhashset::PMMRHandle<BlockHeader>>>,
 	pibd_segmenter: Arc<RwLock<Option<Segmenter>>>,
+	pibd_desegmenter: Arc<RwLock<Option<Desegmenter>>>,
 	// POW verification function
 	pow_verifier: fn(&BlockHeader) -> Result<(), pow::Error>,
 	denylist: Arc<RwLock<Vec<Hash>>>,
@@ -202,6 +203,7 @@ impl Chain {
 			txhashset: Arc::new(RwLock::new(txhashset)),
 			header_pmmr: Arc::new(RwLock::new(header_pmmr)),
 			pibd_segmenter: Arc::new(RwLock::new(None)),
+			pibd_desegmenter: Arc::new(RwLock::new(None)),
 			pow_verifier,
 			denylist: Arc::new(RwLock::new(vec![])),
 			archive_mode,
@@ -863,6 +865,43 @@ impl Chain {
 		))
 	}
 
+	/// instantiate desegmenter (in same lazy fashion as segmenter, though this should not be as
+	/// expensive an operation)
+	pub fn desegmenter(&self, archive_header: &BlockHeader) -> Result<Desegmenter, Error> {
+		// Use our cached desegmenter if we have one and the associated header matches.
+		if let Some(d) = self.pibd_desegmenter.read().as_ref() {
+			if d.header() == archive_header {
+				return Ok(d.clone());
+			}
+		}
+		// If no desegmenter or headers don't match init
+		// TODO: (Check whether we can do this.. we *should* be able to modify this as the desegmenter
+		// is in flight and we cross a horizon boundary, but needs more thinking)
+		let desegmenter = self.init_desegmenter(archive_header)?;
+		let mut cache = self.pibd_desegmenter.write();
+		*cache = Some(desegmenter.clone());
+
+		return Ok(desegmenter);
+	}
+
+	/// initialize a desegmenter, which is capable of extending the hashset by appending
+	/// PIBD segments of the three PMMR trees + Bitmap PMMR
+	/// header should be the same header as selected for the txhashset.zip archive
+	fn init_desegmenter(&self, header: &BlockHeader) -> Result<Desegmenter, Error> {
+		debug!(
+			"init_desegmenter: initializing new desegmenter for {} at {}",
+			header.hash(),
+			header.height
+		);
+
+		Ok(Desegmenter::new(
+			self.txhashset(),
+			self.header_pmmr.clone(),
+			header.clone(),
+			self.store.clone(),
+		))
+	}
+
 	/// To support the ability to download the txhashset from multiple peers in parallel,
 	/// the peers must all agree on the exact binary representation of the txhashset.
 	/// This means compacting and rewinding to the exact same header.

chain/src/txhashset.rs

@@ -16,12 +16,14 @@
 //! kernel) more conveniently and transactionally.
 
 mod bitmap_accumulator;
+mod desegmenter;
 mod rewindable_kernel_view;
 mod segmenter;
 mod txhashset;
 mod utxo_view;
 
 pub use self::bitmap_accumulator::*;
+pub use self::desegmenter::*;
 pub use self::rewindable_kernel_view::*;
 pub use self::segmenter::*;
 pub use self::txhashset::*;

chain/src/txhashset/bitmap_accumulator.rs

@@ -21,7 +21,7 @@ use croaring::Bitmap;
 
 use crate::core::core::hash::{DefaultHashable, Hash};
 use crate::core::core::pmmr::segment::{Segment, SegmentIdentifier, SegmentProof};
-use crate::core::core::pmmr::{self, ReadablePMMR, ReadonlyPMMR, VecBackend, PMMR};
+use crate::core::core::pmmr::{self, Backend, ReadablePMMR, ReadonlyPMMR, VecBackend, PMMR};
 use crate::core::ser::{self, PMMRable, Readable, Reader, Writeable, Writer};
 use crate::error::{Error, ErrorKind};
 use enum_primitive::FromPrimitive;
@@ -190,6 +190,17 @@ impl BitmapAccumulator {
 	pub fn readonly_pmmr(&self) -> ReadonlyPMMR<BitmapChunk, VecBackend<BitmapChunk>> {
 		ReadonlyPMMR::at(&self.backend, self.backend.size())
 	}
+
+	/// Return a raw in-memory bitmap of this accumulator
+	pub fn as_bitmap(&self) -> Result<Bitmap, Error> {
+		let mut bitmap = Bitmap::create();
+		for (chunk_count, chunk_index) in self.backend.leaf_idx_iter(0).enumerate() {
+			//TODO: Unwrap
+			let chunk = self.backend.get_data(chunk_index).unwrap();
+			bitmap.add_many(&chunk.set_iter(chunk_count * 1024).collect::<Vec<u32>>());
+		}
+		Ok(bitmap)
+	}
 }
 
 /// A bitmap "chunk" representing 1024 contiguous bits of the overall bitmap.

chain/src/txhashset/desegmenter.rs

@@ -0,0 +1,243 @@
+// Copyright 2021 The Grin Developers
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Manages the reconsitution of a txhashset from segments produced by the
+//! segmenter
+
+use std::sync::Arc;
+
+use crate::core::core::hash::Hash;
+use crate::core::core::pmmr;
+use crate::core::core::{BlockHeader, OutputIdentifier, Segment, TxKernel};
+use crate::error::Error;
+use crate::txhashset::{BitmapAccumulator, BitmapChunk, TxHashSet};
+use crate::util::secp::pedersen::RangeProof;
+use crate::util::RwLock;
+
+use crate::store;
+use crate::txhashset;
+
+use croaring::Bitmap;
+
+/// Desegmenter for rebuilding a txhashset from PIBD segments
+#[derive(Clone)]
+pub struct Desegmenter {
+	txhashset: Arc<RwLock<TxHashSet>>,
+	header_pmmr: Arc<RwLock<txhashset::PMMRHandle<BlockHeader>>>,
+	archive_header: BlockHeader,
+	store: Arc<store::ChainStore>,
+
+	bitmap_accumulator: BitmapAccumulator,
+	bitmap_segments: Vec<Segment<BitmapChunk>>,
+	output_segments: Vec<Segment<OutputIdentifier>>,
+	rangeproof_segments: Vec<Segment<RangeProof>>,
+	kernel_segments: Vec<Segment<TxKernel>>,
+
+	bitmap_mmr_leaf_count: u64,
+	bitmap_mmr_size: u64,
+	// In-memory 'raw' bitmap corresponding to contents of bitmap accumulator
+	bitmap_cache: Option<Bitmap>,
+}
+
+impl Desegmenter {
+	/// Create a new segmenter based on the provided txhashset and the specified block header
+	pub fn new(
+		txhashset: Arc<RwLock<TxHashSet>>,
+		header_pmmr: Arc<RwLock<txhashset::PMMRHandle<BlockHeader>>>,
+		archive_header: BlockHeader,
+		store: Arc<store::ChainStore>,
+	) -> Desegmenter {
+		let mut retval = Desegmenter {
+			txhashset,
+			header_pmmr,
+			archive_header,
+			store,
+			bitmap_accumulator: BitmapAccumulator::new(),
+			bitmap_segments: vec![],
+			output_segments: vec![],
+			rangeproof_segments: vec![],
+			kernel_segments: vec![],
+
+			bitmap_mmr_leaf_count: 0,
+			bitmap_mmr_size: 0,
+
+			bitmap_cache: None,
+		};
+		retval.calc_bitmap_mmr_sizes();
+		retval
+	}
+
+	/// Return reference to the header used for validation
+	pub fn header(&self) -> &BlockHeader {
+		&self.archive_header
+	}
+	/// Return size of bitmap mmr
+	pub fn expected_bitmap_mmr_size(&self) -> u64 {
+		self.bitmap_mmr_size
+	}
+
+	/// 'Finalize' the bitmap accumulator, storing an in-memory copy of the bitmap for
+	/// use in further validation and setting the accumulator on the underlying txhashset
+	/// TODO: Could be called automatically when we have the calculated number of
+	/// required segments for the archive header
+	/// TODO: Accumulator will likely need to be stored locally to deal with server
+	/// being shut down and restarted
+	pub fn finalize_bitmap(&mut self) -> Result<(), Error> {
+		debug!(
+			"pibd_desgmenter: caching bitmap - accumulator root: {}",
+			self.bitmap_accumulator.root()
+		);
+		self.bitmap_cache = Some(self.bitmap_accumulator.as_bitmap()?);
+
+		// Set the txhashset's bitmap accumulator
+		let mut header_pmmr = self.header_pmmr.write();
+		let mut txhashset = self.txhashset.write();
+		let mut batch = self.store.batch()?;
+		txhashset::extending(
+			&mut header_pmmr,
+			&mut txhashset,
+			&mut batch,
+			|ext, _batch| {
+				let extension = &mut ext.extension;
+				// TODO: Unwrap
+				extension.set_bitmap_accumulator(self.bitmap_accumulator.clone());
+				Ok(())
+			},
+		)?;
+		Ok(())
+	}
+
+	// Calculate and store number of leaves and positions in the bitmap mmr given the number of
+	// outputs specified in the header. Should be called whenever the header changes
+	fn calc_bitmap_mmr_sizes(&mut self) {
+		// Number of leaves (BitmapChunks)
+		self.bitmap_mmr_leaf_count =
+			(pmmr::n_leaves(self.archive_header.output_mmr_size) + 1023) / 1024;
+		debug!(
+			"pibd_desgmenter - expected number of leaves in bitmap MMR: {}",
+			self.bitmap_mmr_leaf_count
+		);
+		// Total size of Bitmap PMMR
+		self.bitmap_mmr_size = pmmr::peaks(self.bitmap_mmr_leaf_count)
+			.last()
+			.unwrap_or(&pmmr::insertion_to_pmmr_index(self.bitmap_mmr_leaf_count))
+			.clone();
+		debug!(
+			"pibd_desgmenter - expected size of bitmap MMR: {}",
+			self.bitmap_mmr_size
+		);
+	}
+
+	/// Adds and validates a bitmap chunk
+	/// TODO: Still experimenting, this expects chunks received to be in order
+	pub fn add_bitmap_segment(
+		&mut self,
+		segment: Segment<BitmapChunk>,
+		output_root_hash: Hash,
+	) -> Result<(), Error> {
+		debug!("pibd_desegmenter: add bitmap segment");
+		segment.validate_with(
+			self.bitmap_mmr_size, // Last MMR pos at the height being validated, in this case of the bitmap root
+			None,
+			self.archive_header.output_root, // Output root we're checking for
+			self.archive_header.output_mmr_size,
+			output_root_hash, // Other root
+			true,
+		)?;
+		// All okay, add leaves to bitmap accumulator
+		let (_sid, _hash_pos, _hashes, _leaf_pos, leaf_data, _proof) = segment.parts();
+		for chunk in leaf_data.into_iter() {
+			self.bitmap_accumulator.append_chunk(chunk)?;
+		}
+		Ok(())
+	}
+
+	/// Adds a output segment
+	/// TODO: Still experimenting, expects chunks received to be in order
+	pub fn add_output_segment(&self, segment: Segment<OutputIdentifier>) -> Result<(), Error> {
+		debug!("pibd_desegmenter: add output segment");
+		segment.validate_with(
+			self.archive_header.output_mmr_size, // Last MMR pos at the height being validated
+			self.bitmap_cache.as_ref(),
+			self.archive_header.output_root, // Output root we're checking for
+			self.archive_header.output_mmr_size,
+			self.bitmap_accumulator.root(), // Other root
+			false,
+		)?;
+		let mut header_pmmr = self.header_pmmr.write();
+		let mut txhashset = self.txhashset.write();
+		let mut batch = self.store.batch()?;
+		txhashset::extending(
+			&mut header_pmmr,
+			&mut txhashset,
+			&mut batch,
+			|ext, _batch| {
+				let extension = &mut ext.extension;
+				extension.apply_output_segment(segment)?;
+				Ok(())
+			},
+		)?;
+		Ok(())
+	}
+
+	/// Adds a Rangeproof segment
+	/// TODO: Still experimenting, expects chunks received to be in order
+	pub fn add_rangeproof_segment(&self, segment: Segment<RangeProof>) -> Result<(), Error> {
+		debug!("pibd_desegmenter: add rangeproof segment");
+		segment.validate(
+			self.archive_header.output_mmr_size, // Last MMR pos at the height being validated
+			self.bitmap_cache.as_ref(),
+			self.archive_header.range_proof_root, // Range proof root we're checking for
+		)?;
+		let mut header_pmmr = self.header_pmmr.write();
+		let mut txhashset = self.txhashset.write();
+		let mut batch = self.store.batch()?;
+		txhashset::extending(
+			&mut header_pmmr,
+			&mut txhashset,
+			&mut batch,
+			|ext, _batch| {
+				let extension = &mut ext.extension;
+				extension.apply_rangeproof_segment(segment)?;
+				Ok(())
+			},
+		)?;
+		Ok(())
+	}
+
+	/// Adds a Kernel segment
+	/// TODO: Still experimenting, expects chunks received to be in order
+	pub fn add_kernel_segment(&self, segment: Segment<TxKernel>) -> Result<(), Error> {
+		debug!("pibd_desegmenter: add kernel segment");
+		segment.validate(
+			self.archive_header.kernel_mmr_size, // Last MMR pos at the height being validated
+			None,
+			self.archive_header.kernel_root, // Kernel root we're checking for
+		)?;
+		let mut header_pmmr = self.header_pmmr.write();
+		let mut txhashset = self.txhashset.write();
+		let mut batch = self.store.batch()?;
+		txhashset::extending(
+			&mut header_pmmr,
+			&mut txhashset,
+			&mut batch,
+			|ext, _batch| {
+				let extension = &mut ext.extension;
+				extension.apply_kernel_segment(segment)?;
+				Ok(())
+			},
+		)?;
+		Ok(())
+	}
+}

chain/src/txhashset/txhashset.rs

@@ -22,7 +22,9 @@ use crate::core::core::merkle_proof::MerkleProof;
 use crate::core::core::pmmr::{
 	self, Backend, ReadablePMMR, ReadonlyPMMR, RewindablePMMR, VecBackend, PMMR,
 };
-use crate::core::core::{Block, BlockHeader, KernelFeatures, Output, OutputIdentifier, TxKernel};
+use crate::core::core::{
+	Block, BlockHeader, KernelFeatures, Output, OutputIdentifier, Segment, TxKernel,
+};
 use crate::core::global;
 use crate::core::ser::{PMMRable, ProtocolVersion};
 use crate::error::{Error, ErrorKind};
@@ -1177,6 +1179,11 @@ impl<'a> Extension<'a> {
 		)
 	}
 
+	/// Sets the bitmap accumulator (as received during PIBD sync)
+	pub fn set_bitmap_accumulator(&mut self, accumulator: BitmapAccumulator) {
+		self.bitmap_accumulator = accumulator;
+	}
+
 	// Prune output and rangeproof PMMRs based on provided pos.
 	// Input is not valid if we cannot prune successfully.
 	fn apply_input(&mut self, commit: Commitment, pos: CommitPos) -> Result<(), Error> {
@@ -1232,6 +1239,33 @@ impl<'a> Extension<'a> {
 		Ok(1 + output_pos)
 	}
 
+	/// Apply an output segment to the output PMMR. must be called in order
+	/// TODO: Not complete
+	pub fn apply_output_segment(
+		&mut self,
+		segment: Segment<OutputIdentifier>,
+	) -> Result<(), Error> {
+		let (_sid, _hash_pos, _hashes, _leaf_pos, leaf_data, _proof) = segment.parts();
+		for output_identifier in leaf_data {
+			self.output_pmmr
+				.push(&output_identifier)
+				.map_err(&ErrorKind::TxHashSetErr)?;
+		}
+		Ok(())
+	}
+
+	/// Apply a rangeproof segment to the output PMMR. must be called in order
+	/// TODO: Not complete
+	pub fn apply_rangeproof_segment(&mut self, segment: Segment<RangeProof>) -> Result<(), Error> {
+		let (_sid, _hash_pos, _hashes, _leaf_pos, leaf_data, _proof) = segment.parts();
+		for proof in leaf_data {
+			self.rproof_pmmr
+				.push(&proof)
+				.map_err(&ErrorKind::TxHashSetErr)?;
+		}
+		Ok(())
+	}
+
 	/// Apply kernels to the kernel MMR.
 	/// Validate any NRD relative height locks via the "recent" kernel index.
 	/// Note: This is used for both block processing and tx validation.
@@ -1251,6 +1285,18 @@ impl<'a> Extension<'a> {
 		Ok(())
 	}
 
+	/// Apply a kernel segment to the output PMMR. must be called in order
+	/// TODO: Not complete
+	pub fn apply_kernel_segment(&mut self, segment: Segment<TxKernel>) -> Result<(), Error> {
+		let (_sid, _hash_pos, _hashes, _leaf_pos, leaf_data, _proof) = segment.parts();
+		for kernel in leaf_data {
+			self.kernel_pmmr
+				.push(&kernel)
+				.map_err(&ErrorKind::TxHashSetErr)?;
+		}
+		Ok(())
+	}
+
 	/// Push kernel onto MMR (hash and data files).
 	fn apply_kernel(&mut self, kernel: &TxKernel) -> Result<u64, Error> {
 		let pos = self