jam_std_common/
bundle.rs

1use codec::Encode;
2use jam_types::{
3	Hash, SegmentTreeRoot, WorkPackage, WorkPackageHash, PROVEN_PER_SEGMENT, SEGMENT_LEN,
4};
5
6use crate::{cd_merkle_proof, cd_merkle_root, hash_raw, hash_raw_concat, CdMerkleProof};
7
8// Bundled import data.
9pub struct ImportData {
10	// Imported segment content.
11	pub segment: Vec<u8>,
12	// Imported segment proof.
13	pub proof: CdMerkleProof,
14}
15
16/// Construct an encoded Work Package Bundle from its parts
17///
18/// This function assumes the number of imports and the length of the extrinsics match the Work
19/// Package and does not explicitly check that fact.
20///
21/// Returns encoded package hash and the encoded bundle.
22///
23/// # Arguments
24///
25/// - `package`: work package.
26/// - `chunked_extrinsics`: this iterator goes over arbitrarily chunked extrinsics. All of the
27///   chunks are concatenated to form the final extrinsic data.
28/// - `imports`: one element for each work package item. Each element contains a vector of import
29///   data for that item.
30pub fn build_encoded_bundle(
31	package: &WorkPackage,
32	chunked_extrinsics: impl IntoIterator<Item = impl AsRef<[u8]>>,
33	imports: &[Vec<ImportData>],
34) -> (WorkPackageHash, Vec<u8>) {
35	let mut encoded = Vec::new();
36	package.encode_to(&mut encoded);
37	let wp_hash = hash_raw(&encoded);
38	for chunk in chunked_extrinsics {
39		encoded.extend_from_slice(chunk.as_ref());
40	}
41	for item_imports in imports {
42		for import in item_imports {
43			encoded.extend_from_slice(&import.segment);
44		}
45	}
46	for item_imports in imports {
47		for import in item_imports {
48			import.proof.encode_to(&mut encoded);
49		}
50	}
51	(wp_hash.into(), encoded)
52}
53
54/// Generate import proofs for exported segments.
55/// Returns a full import proof for each segment and a tree root.
56pub fn import_proofs(
57	segments: &[impl AsRef<[u8; SEGMENT_LEN]>],
58) -> (Vec<CdMerkleProof>, SegmentTreeRoot) {
59	let (page_proofs, root) = import_proof_data(segments);
60	let proofs = page_proofs
61		.into_iter()
62		.enumerate()
63		.flat_map(|(i, (prefix, hashes))| {
64			let padded_len =
65				if i > 0 { PROVEN_PER_SEGMENT } else { hashes.len().next_power_of_two() };
66			(0..hashes.len()).map(move |index| {
67				let mut proof = prefix.to_vec();
68				proof.extend(cd_merkle_proof(&hashes, padded_len, index).0);
69				proof
70			})
71		})
72		.collect::<Vec<_>>();
73	(proofs, root)
74}
75
76/// Prepare import proof data for a set of segments.
77/// This function splits the segment into `PROVEN_PER_SEGMENT` pages and for each page returns
78/// a page prefix proof along with the vector of `PROVEN_PER_SEGMENT` hashes that form the page.
79pub fn import_proof_data(
80	segments: &[impl AsRef<[u8; SEGMENT_LEN]>],
81) -> (impl IntoIterator<Item = (CdMerkleProof, Vec<Hash>)>, SegmentTreeRoot) {
82	// Construct the aligned pages.
83	let pages = segments
84		.chunks(PROVEN_PER_SEGMENT)
85		.map(|i| {
86			i.iter()
87				.map(|s| hash_raw_concat([b"leaf", s.as_ref().as_slice()]))
88				.collect::<Vec<_>>()
89		})
90		.collect::<Vec<_>>();
91	// Figure out the roots of each page's subtree.
92	let subtree_roots = pages
93		.iter()
94		.map(|s| {
95			let padded_len =
96				if pages.len() > 1 { PROVEN_PER_SEGMENT } else { s.len().next_power_of_two() };
97			cd_merkle_root(s, padded_len)
98		})
99		.collect::<Vec<_>>();
100
101	// Bottom row of the super-tree is each of these page-roots. We ensure it is a Po2:
102	let mut tree = vec![];
103	let root = if !segments.is_empty() {
104		tree = vec![subtree_roots];
105		let leafs = pages.len().next_power_of_two();
106		if tree[0].len() < leafs {
107			let empty_root = cd_merkle_root(&[], PROVEN_PER_SEGMENT);
108			tree[0].resize(leafs, empty_root);
109		}
110		// Rest of the super-tree can be derived from bottom row:
111		tree.reserve(leafs.trailing_zeros() as usize + 1);
112		loop {
113			let last = tree.last().expect(
114				"Initially the tree has exactly one node, \
115				we don't remove the nodes in the loop",
116			);
117			if last.len() <= 1 {
118				break;
119			}
120			let next = last.chunks(2).map(|c| cd_merkle_root(c, 2)).collect::<Vec<_>>();
121			tree.push(next);
122		}
123
124		tree.pop().expect(
125			"Initially the tree has exactly one node, \
126			we don't remove the nodes in the loop",
127		)[0]
128		.into()
129	} else {
130		Default::default()
131	};
132
133	// Construct the proof pages from the super-tree and subtree contents.
134	let pages = pages.into_iter().enumerate().map(move |(i, page)| {
135		let proof: CdMerkleProof = tree
136			.iter()
137			.enumerate()
138			.map(|(level, row)| row[(i >> level) ^ 1])
139			.rev()
140			.collect();
141		(proof, page)
142	});
143	(pages, root)
144}
145
146#[test]
147fn import_proof_data_empty_hash_when_empty() {
148	let empty: [jam_types::SegmentBytes; 0] = [];
149	let (_, root) = import_proof_data(&empty);
150	assert_eq!(root, Default::default())
151}
jam_std_common/bundle.rs

jam_std_common/
bundle.rs
