1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
// Copyright 2019 Parity Technologies
//
// 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.

//! Generation of compact proofs for Merkle-Patricia tries.

use crate::rstd::{
	boxed::Box, convert::TryInto, marker::PhantomData, ops::Range, vec, vec::Vec,
};

use hash_db::Hasher;

use crate::{
	CError, ChildReference, nibble::LeftNibbleSlice, nibble_ops::NIBBLE_LENGTH, NibbleSlice, node::{NodeHandle, NodeHandlePlan, NodePlan, OwnedNode}, NodeCodec, Recorder,
	Result as TrieResult, Trie, TrieError, TrieHash,
	TrieLayout,
};

struct StackEntry<'a, C: NodeCodec> {
	/// The prefix is the nibble path to the node in the trie.
	prefix: LeftNibbleSlice<'a>,
	node: OwnedNode<Vec<u8>>,
	/// The hash of the node or None if it is referenced inline.
	node_hash: Option<C::HashOut>,
	/// Whether the value should be omitted in the generated proof.
	omit_value: bool,
	/// The next entry in the stack is a child of the preceding entry at this index. For branch
	/// nodes, the index is in [0, NIBBLE_LENGTH] and for extension nodes, the index is in [0, 1].
	child_index: usize,
	/// The child references to use in constructing the proof nodes.
	children: Vec<Option<ChildReference<C::HashOut>>>,
	/// The index into the proof vector that the encoding of this entry should be placed at.
	output_index: Option<usize>,
	_marker: PhantomData<C>,
}

impl<'a, C: NodeCodec> StackEntry<'a, C> {
	fn new(
		prefix: LeftNibbleSlice<'a>,
		node_data: Vec<u8>,
		node_hash: Option<C::HashOut>,
		output_index: Option<usize>,
	) -> TrieResult<Self, C::HashOut, C::Error>
	{
		let node = OwnedNode::new::<C>(node_data)
			.map_err(|err| Box::new(
				TrieError::DecoderError(node_hash.unwrap_or_default(), err)
			))?;
		let children_len = match node.node_plan() {
			NodePlan::Empty | NodePlan::Leaf { .. } => 0,
			NodePlan::Extension { .. } => 1,
			NodePlan::Branch { .. } | NodePlan::NibbledBranch { .. } => NIBBLE_LENGTH,
		};
		Ok(StackEntry {
			prefix,
			node,
			node_hash,
			omit_value: false,
			child_index: 0,
			children: vec![None; children_len],
			output_index,
			_marker: PhantomData::default(),
		})
	}

	/// Encode this entry to an encoded trie node with data properly omitted.
	fn encode_node(mut self) -> TrieResult<Vec<u8>, C::HashOut, C::Error> {
		let node_data = self.node.data();
		Ok(match self.node.node_plan() {
			NodePlan::Empty => node_data.to_vec(),
			NodePlan::Leaf { .. } if !self.omit_value => node_data.to_vec(),
			NodePlan::Leaf { partial, value: _ } => {
				let partial = partial.build(node_data);
				C::leaf_node(partial.right(), &[])
			}
			NodePlan::Extension { .. } if self.child_index == 0 => node_data.to_vec(),
			NodePlan::Extension { partial: partial_plan, child: _ } => {
				let partial = partial_plan.build(node_data);
				let child = self.children[0]
					.expect(
						"for extension nodes, children[0] is guaranteed to be Some when \
						child_index > 0; \
						the branch guard guarantees that child_index > 0"
					);
				C::extension_node(
					partial.right_iter(),
					partial.len(),
					child
				)
			}
			NodePlan::Branch { value, children } => {
				Self::complete_branch_children(
					node_data,
					children,
					self.child_index,
					&mut self.children
				)?;
				C::branch_node(
					self.children.into_iter(),
					value_with_omission(node_data, value, self.omit_value)
				)
			},
			NodePlan::NibbledBranch { partial: partial_plan, value, children } => {
				let partial = partial_plan.build(node_data);
				Self::complete_branch_children(
					node_data,
					children,
					self.child_index,
					&mut self.children
				)?;
				C::branch_node_nibbled(
					partial.right_iter(),
					partial.len(),
					self.children.into_iter(),
					value_with_omission(node_data, value, self.omit_value)
				)
			},
		})
	}

	/// Populate the remaining references in `children` with references copied from
	/// `child_handles`.
	///
	/// Preconditions:
	/// - children has size NIBBLE_LENGTH.
	fn complete_branch_children(
		node_data: &[u8],
		child_handles: &[Option<NodeHandlePlan>; NIBBLE_LENGTH],
		child_index: usize,
		children: &mut [Option<ChildReference<C::HashOut>>],
	) -> TrieResult<(), C::HashOut, C::Error>
	{
		for i in child_index..NIBBLE_LENGTH {
			children[i] = child_handles[i]
				.as_ref()
				.map(|child_plan|
					child_plan
						.build(node_data)
						.try_into()
						.map_err(|hash| Box::new(
							TrieError::InvalidHash(C::HashOut::default(), hash)
						))
				)
				.transpose()?;
		}
		Ok(())
	}

	/// Sets the reference for the child at index `child_index`. If the child is hash-referenced in
	/// the trie, the proof node reference will be an omitted child. If the child is
	/// inline-referenced in the trie, the proof node reference will also be inline.
	fn set_child(&mut self, encoded_child: &[u8]) {
		let child_ref = match self.node.node_plan() {
			NodePlan::Empty | NodePlan::Leaf { .. } => panic!(
				"empty and leaf nodes have no children; \
				thus they are never descended into; \
				thus set_child will not be called on an entry with one of these types"
			),
			NodePlan::Extension { child, .. } => {
				assert_eq!(
					self.child_index, 0,
					"extension nodes only have one child; \
					set_child is called when the only child is popped from the stack; \
					child_index is 0 before child is pushed to the stack; qed"
				);
				Some(Self::replacement_child_ref(encoded_child, child))
			}
			NodePlan::Branch { children, .. } | NodePlan::NibbledBranch { children, .. } => {
				assert!(
					self.child_index < NIBBLE_LENGTH,
					"extension nodes have at most NIBBLE_LENGTH children; \
					set_child is called when the only child is popped from the stack; \
					child_index is <NIBBLE_LENGTH before child is pushed to the stack; qed"
				);
				children[self.child_index]
					.as_ref()
					.map(|child| Self::replacement_child_ref(encoded_child, child))
			}
		};
		self.children[self.child_index] = child_ref;
		self.child_index += 1;
	}

	/// Build a proof node child reference. If the child is hash-referenced in the trie, the proof
	/// node reference will be an omitted child. If the child is inline-referenced in the trie, the
	/// proof node reference will also be inline.
	fn replacement_child_ref(encoded_child: &[u8], child: &NodeHandlePlan)
							 -> ChildReference<C::HashOut>
	{
		match child {
			NodeHandlePlan::Hash(_) => ChildReference::Inline(C::HashOut::default(), 0),
			NodeHandlePlan::Inline(_) => {
				let mut hash = C::HashOut::default();
				assert!(
					encoded_child.len() <= hash.as_ref().len(),
					"the encoding of the raw inline node is checked to be at most the hash length
					before descending; \
					the encoding of the proof node is always smaller than the raw node as data is \
					only stripped"
				);
				&mut hash.as_mut()[..encoded_child.len()].copy_from_slice(encoded_child);
				ChildReference::Inline(hash, encoded_child.len())
			}
		}
	}
}

/// Generate a compact proof for key-value pairs in a trie given a set of keys.
///
/// Assumes inline nodes have only inline children.
pub fn generate_proof<'a, T, L, I, K>(trie: &T, keys: I)
									  -> TrieResult<Vec<Vec<u8>>, TrieHash<L>, CError<L>>
	where
		T: Trie<L>,
		L: TrieLayout,
		I: IntoIterator<Item=&'a K>,
		K: 'a + AsRef<[u8]>
{
	// Sort and deduplicate keys.
	let mut keys = keys.into_iter()
		.map(|key| key.as_ref())
		.collect::<Vec<_>>();
	keys.sort();
	keys.dedup();

	// The stack of nodes through a path in the trie. Each entry is a child node of the preceding
	// entry.
	let mut stack = <Vec<StackEntry<L::Codec>>>::new();

	// The mutated trie nodes comprising the final proof.
	let mut proof_nodes = Vec::new();

	for key_bytes in keys {
		let key = LeftNibbleSlice::new(key_bytes);

		// Unwind the stack until the new entry is a child of the last entry on the stack.
		unwind_stack(&mut stack, &mut proof_nodes, Some(&key))?;

		// Perform the trie lookup for the next key, recording the sequence of nodes traversed.
		let mut recorder = Recorder::new();
		let expected_value = trie.get_with(key_bytes, &mut recorder)?;
		let mut recorded_nodes = recorder.drain().into_iter().peekable();

		// Skip over recorded nodes already on the stack. Their indexes into the respective vector
		// (either `stack` or `recorded_nodes`) match under the assumption that inline nodes have
		// only inline children.
		{
			let mut stack_iter = stack.iter().peekable();
			while let (Some(next_record), Some(next_entry)) =
			(recorded_nodes.peek(), stack_iter.peek())
				{
					if next_entry.node_hash != Some(next_record.hash) {
						break;
					}
					recorded_nodes.next();
					stack_iter.next();
				}
		}

		loop {
			let step = match stack.last_mut() {
				Some(entry) => match_key_to_node::<L::Codec>(
					entry.node.data(),
					entry.node.node_plan(),
					&mut entry.omit_value,
					&mut entry.child_index,
					&mut entry.children,
					&key,
					entry.prefix.len(),
				)?,
				// If stack is empty, descend into the root node.
				None => Step::Descend {
					child_prefix_len: 0,
					child: NodeHandle::Hash(trie.root().as_ref()),
				},
			};

			match step {
				Step::Descend { child_prefix_len, child } => {
					let child_prefix = key.truncate(child_prefix_len);
					let child_entry = match child {
						NodeHandle::Hash(hash) => {
							let child_record = recorded_nodes.next()
								.expect(
									"this function's trie traversal logic mirrors that of Lookup; \
									thus the sequence of traversed nodes must be the same; \
									so the next child node must have been recorded and must have \
									the expected hash"
								);
							// Proof for `assert_eq` is in the `expect` proof above.
							assert_eq!(child_record.hash.as_ref(), hash);

							let output_index = proof_nodes.len();
							// Insert a placeholder into output which will be replaced when this
							// new entry is popped from the stack.
							proof_nodes.push(Vec::new());
							StackEntry::new(
								child_prefix,
								child_record.data,
								Some(child_record.hash),
								Some(output_index),
							)?
						}
						NodeHandle::Inline(data) => {
							if data.len() > L::Hash::LENGTH {
								return Err(Box::new(
									TrieError::InvalidHash(<TrieHash<L>>::default(), data.to_vec())
								));
							}
							StackEntry::new(
								child_prefix,
								data.to_vec(),
								None,
								None,
							)?
						}
					};
					stack.push(child_entry);
				}
				Step::FoundValue(value) => {
					assert_eq!(
						value,
						expected_value.as_ref().map(|v| v.as_ref()),
						"expected_value is found using `trie_db::Lookup`; \
						value is found by traversing the same nodes recorded during the lookup \
						using the same logic; \
						thus the values found must be equal"
					);
					assert!(
						recorded_nodes.next().is_none(),
						"the recorded nodes are only recorded on the lookup path to the current \
						key; \
						recorded nodes is the minimal sequence of trie nodes on the lookup path; \
						the value was found by traversing recorded nodes, so there must be none \
						remaining"
					);
					break;
				}
			}
		}
	}

	unwind_stack(&mut stack, &mut proof_nodes, None)?;
	Ok(proof_nodes)
}

enum Step<'a> {
	Descend {
		child_prefix_len: usize,
		child: NodeHandle<'a>,
	},
	FoundValue(Option<&'a [u8]>),
}

/// Determine the next algorithmic step to take by matching the current key against the current top
/// entry on the stack.
fn match_key_to_node<'a, C: NodeCodec>(
	node_data: &'a [u8],
	node_plan: &NodePlan,
	omit_value: &mut bool,
	child_index: &mut usize,
	children: &mut [Option<ChildReference<C::HashOut>>],
	key: &LeftNibbleSlice,
	prefix_len: usize,
) -> TrieResult<Step<'a>, C::HashOut, C::Error>
{
	Ok(match node_plan {
		NodePlan::Empty => Step::FoundValue(None),
		NodePlan::Leaf { partial: partial_plan, value: value_range } => {
			let partial = partial_plan.build(node_data);
			if key.contains(&partial, prefix_len) &&
				key.len() == prefix_len + partial.len()
			{
				*omit_value = true;
				Step::FoundValue(Some(&node_data[value_range.clone()]))
			} else {
				Step::FoundValue(None)
			}
		}
		NodePlan::Extension { partial: partial_plan, child: child_plan } => {
			let partial = partial_plan.build(node_data);
			if key.contains(&partial, prefix_len) {
				assert_eq!(*child_index, 0);
				let child_prefix_len = prefix_len + partial.len();
				let child = child_plan.build(&node_data);
				Step::Descend { child_prefix_len, child }
			} else {
				Step::FoundValue(None)
			}
		}
		NodePlan::Branch { value, children: child_handles } =>
			match_key_to_branch_node::<C>(
				node_data,
				value,
				&child_handles,
				omit_value,
				child_index,
				children,
				key,
				prefix_len,
				NibbleSlice::new(&[]),
			)?,
		NodePlan::NibbledBranch { partial: partial_plan, value, children: child_handles } =>
			match_key_to_branch_node::<C>(
				node_data,
				value,
				&child_handles,
				omit_value,
				child_index,
				children,
				key,
				prefix_len,
				partial_plan.build(node_data),
			)?,
	})
}

fn match_key_to_branch_node<'a, 'b, C: NodeCodec>(
	node_data: &'a [u8],
	value_range: &'b Option<Range<usize>>,
	child_handles: &'b [Option<NodeHandlePlan>; NIBBLE_LENGTH],
	omit_value: &mut bool,
	child_index: &mut usize,
	children: &mut [Option<ChildReference<C::HashOut>>],
	key: &'b LeftNibbleSlice<'b>,
	prefix_len: usize,
	partial: NibbleSlice<'b>,
) -> TrieResult<Step<'a>, C::HashOut, C::Error>
{
	if !key.contains(&partial, prefix_len) {
		return Ok(Step::FoundValue(None));
	}

	if key.len() == prefix_len + partial.len() {
		*omit_value = true;
		let value = value_range.clone().map(|range| &node_data[range]);
		return Ok(Step::FoundValue(value));
	}

	let new_index = key.at(prefix_len + partial.len())
		.expect(
			"key contains partial key after entry key offset; \
			thus key len is greater than equal to entry key len plus partial key len; \
			also they are unequal due to else condition;
			qed"
		)
		as usize;
	assert!(*child_index <= new_index);
	while *child_index < new_index {
		children[*child_index] = child_handles[*child_index]
			.as_ref()
			.map(|child_plan|
				child_plan
					.build(node_data)
					.try_into()
					.map_err(|hash| Box::new(
						TrieError::InvalidHash(C::HashOut::default(), hash)
					))
			)
			.transpose()?;
		*child_index += 1;
	}
	if let Some(child_plan) = &child_handles[*child_index] {
		Ok(Step::Descend {
			child_prefix_len: prefix_len + partial.len() + 1,
			child: child_plan.build(node_data),
		})
	} else {
		Ok(Step::FoundValue(None))
	}
}

fn value_with_omission<'a>(
	node_data: &'a [u8],
	value_range: &Option<Range<usize>>,
	omit: bool
) -> Option<&'a [u8]>
{
	if omit {
		None
	} else {
		value_range.clone().map(|range| &node_data[range])
	}
}

/// Unwind the stack until the given key is prefixed by the entry at the top of the stack. If the
/// key is None, unwind the stack completely. As entries are popped from the stack, they are
/// encoded into proof nodes and added to the finalized proof.
fn unwind_stack<C: NodeCodec>(
	stack: &mut Vec<StackEntry<C>>,
	proof_nodes: &mut Vec<Vec<u8>>,
	maybe_key: Option<&LeftNibbleSlice>,
) -> TrieResult<(), C::HashOut, C::Error>
{
	while let Some(entry) = stack.pop() {
		match maybe_key {
			Some(key) if key.starts_with(&entry.prefix) => {
				// Stop if the key lies below this entry in the trie.
				stack.push(entry);
				break;
			}
			_ => {
				// Pop and finalize node from the stack.
				let index = entry.output_index;
				let encoded = entry.encode_node()?;
				if let Some(parent_entry) = stack.last_mut() {
					parent_entry.set_child(&encoded);
				}
				if let Some(index) = index {
					proof_nodes[index] = encoded;
				}
			}
		}
	}
	Ok(())
}