jam-std-common 0.1.28

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

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 imports and the length of the extrinsics match the Work
/// Package and does not explicitly check that fact.
///
/// Returns encoded package hash and the encoded bundle.
///
/// # Arguments
///
/// - `package`: work package.
/// - `chunked_extrinsics`: this iterator goes over arbitrarily chunked extrinsics. All of the
///   chunks are concatenated to form the final extrinsic data.
/// - `imports`: one element for each work package item. Each element contains a vector of import
///   data for that item.
pub fn build_encoded_bundle(
	package: &WorkPackage,
	chunked_extrinsics: impl IntoIterator<Item = impl AsRef<[u8]>>,
	imports: &[Vec<ImportData>],
) -> (WorkPackageHash, Vec<u8>) {
	let mut encoded = Vec::new();
	package.encode_to(&mut encoded);
	let wp_hash = hash_raw(&encoded);
	for chunk in chunked_extrinsics {
		encoded.extend_from_slice(chunk.as_ref());
	}
	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: &[impl AsRef<[u8; SEGMENT_LEN]>],
) -> (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: &[impl AsRef<[u8; SEGMENT_LEN]>],
) -> (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().as_slice()]))
				.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 empty: [jam_types::SegmentBytes; 0] = [];
	let (_, root) = import_proof_data(&empty);
	assert_eq!(root, Default::default())
}