jam_std_common/
bundle.rs

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

jam_std_common/
bundle.rs
