jam-std-common 0.1.26

use codec::Encode;
use jam_types::{Hash, Segment, SegmentTreeRoot, WorkPackage, WorkPackageHash, PROVEN_PER_SEGMENT};

use crate::{cd_merkle_proof, cd_merkle_root, hash_raw, hash_raw_concat, CdMerkleProof};

// Bundled import data.
pub struct ImportData {
	// Imported segment content.
	pub segment: Vec<u8>,
	// Imported segment proof.
	pub proof: CdMerkleProof,
}

/// Construct an encoded Work Package Bundle from its parts
///
/// This function assumes the number of extrinsics and imports matches the Work Package and does not
/// explicitly check that fact. `package` - Work Package,
/// `extrinsics` - One element for each Work Package item. Each element contains a vector of
/// extrinsic data for that item. `imports` - One element for each Work Package item. Each element
/// contains a vector of import data for that item.
/// Returns encoded package hash and the econded bundle.
pub fn build_encoded_bundle(
	package: &WorkPackage,
	extrinsics: &[Vec<Vec<u8>>],
	imports: &[Vec<ImportData>],
) -> (WorkPackageHash, Vec<u8>) {
	let mut encoded = Vec::new();
	package.encode_to(&mut encoded);
	let wp_hash = hash_raw(&encoded);
	for item_xts in extrinsics {
		for xt in item_xts {
			encoded.extend_from_slice(xt);
		}
	}
	for item_imports in imports {
		for import in item_imports {
			encoded.extend_from_slice(&import.segment);
		}
	}
	for item_imports in imports {
		for import in item_imports {
			import.proof.encode_to(&mut encoded);
		}
	}
	(wp_hash.into(), encoded)
}

/// Generate import proofs for exported segments.
/// Returns a full import proof for each segment and a tree root.
pub fn import_proofs(segments: &[Segment]) -> (Vec<CdMerkleProof>, SegmentTreeRoot) {
	let (page_proofs, root) = import_proof_data(segments);
	let proofs = page_proofs
		.into_iter()
		.enumerate()
		.flat_map(|(i, (prefix, hashes))| {
			let padded_len =
				if i > 0 { PROVEN_PER_SEGMENT } else { hashes.len().next_power_of_two() };
			(0..hashes.len()).map(move |index| {
				let mut proof = prefix.to_vec();
				proof.extend(cd_merkle_proof(&hashes, padded_len, index).0);
				proof
			})
		})
		.collect::<Vec<_>>();
	(proofs, root)
}

/// Prepare import proof data for a set of segments.
/// This function splits the segment into `PROVEN_PER_SEGMENT` pages and for each page returns
/// a page prefix proof along with the vector of `PROVEN_PER_SEGMENT` hashes that form the page.
pub fn import_proof_data(
	segments: &[Segment],
) -> (impl IntoIterator<Item = (CdMerkleProof, Vec<Hash>)>, SegmentTreeRoot) {
	// Construct the aligned pages.
	let pages = segments
		.chunks(PROVEN_PER_SEGMENT)
		.map(|i| i.iter().map(|s| hash_raw_concat([b"leaf", s.as_ref()])).collect::<Vec<_>>())
		.collect::<Vec<_>>();
	// Figure out the roots of each page's subtree.
	let subtree_roots = pages
		.iter()
		.map(|s| {
			let padded_len =
				if pages.len() > 1 { PROVEN_PER_SEGMENT } else { s.len().next_power_of_two() };
			cd_merkle_root(s, padded_len)
		})
		.collect::<Vec<_>>();

	// Bottom row of the super-tree is each of these page-roots. We ensure it is a Po2:
	let mut tree = vec![];
	let root = if !segments.is_empty() {
		tree = vec![subtree_roots];
		let leafs = pages.len().next_power_of_two();
		if tree[0].len() < leafs {
			let empty_root = cd_merkle_root(&[], PROVEN_PER_SEGMENT);
			tree[0].resize(leafs, empty_root);
		}
		// Rest of the super-tree can be derived from bottom row:
		tree.reserve(leafs.trailing_zeros() as usize + 1);
		loop {
			let last = tree.last().expect(
				"Initially the tree has exactly one node, \
				we don't remove the nodes in the loop",
			);
			if last.len() <= 1 {
				break;
			}
			let next = last.chunks(2).map(|c| cd_merkle_root(c, 2)).collect::<Vec<_>>();
			tree.push(next);
		}

		tree.pop().expect(
			"Initially the tree has exactly one node, \
			we don't remove the nodes in the loop",
		)[0]
		.into()
	} else {
		Default::default()
	};

	// Construct the proof pages from the super-tree and subtree contents.
	let pages = pages.into_iter().enumerate().map(move |(i, page)| {
		let proof: CdMerkleProof = tree
			.iter()
			.enumerate()
			.map(|(level, row)| row[(i >> level) ^ 1])
			.rev()
			.collect();
		(proof, page)
	});
	(pages, root)
}

#[test]
fn import_proof_data_empty_hash_when_empty() {
	let (_, root) = import_proof_data(&[]);
	assert_eq!(root, Default::default())
}