Simple Merkle Tree

merkle_tree.rs

//! This is minimal Merkle tree implementation with proof checking

use std::{
	collections::hash_map::DefaultHasher,
	hash::{Hash, Hasher},
};

fn main() {
	let s = "";
	let h = calculate_merkle_root("Trust me, bro!");
	let result = generate_proof("Trust me, bro! German", 3);
	let result2 = generate_proof("I need to find an index of an element in a vector of strings", 8);
	println!("{:x}", h);
	println!("{:x}", result.0);
	println!("{}", validate_proof(result2.0, "element", result2.1));
}

/// We'll use Rust's built-in hashing which returns a u64 type.
/// This alias just helps us understand when we're treating the number as a hash
type HashValue = u64;

/// Helper function that makes the hashing interface easier to understand.
fn hash<T: Hash>(t: &T) -> HashValue {
	let mut s = DefaultHasher::new();
	t.hash(&mut s);
	s.finish()
}

/// Given a vector of data blocks this function adds padding blocks to the end
/// until the length is a power of two which is needed for Merkle trees.
fn pad_base_layer(blocks: &mut Vec<&str>) {
	let n = blocks.len().next_power_of_two() - blocks.len();
	for _ in 0..n {
		blocks.push(" ");
	}
}

/// Helper function to combine two hashes and compute the hash of the combination.
/// This will be useful when building the intermediate nodes in the Merkle tree.
///
/// There are many correct ways to do this, but the easiest one and the one that I recommend
/// is to convert the hashes to strings and concatenate them.
fn concatenate_hash_values(left: HashValue, right: HashValue) -> HashValue {
	let string = format!("{:x}{:x}", left, right);
	hash(&string)
}

/// Calculates the Merkle root of a sentence. We consider each word in the sentence to
/// be one block. Words are separated by one or more spaces.
///
/// Example:
/// Sentence: "Trust me, bro!"
/// "Trust", "me," "bro!"
/// Notice that the punctuation like the comma and exclamation point are included in the words
/// but the spaces are not.
fn calculate_merkle_root(sentence: &str) -> HashValue {
	let mut array: Vec<&str> = sentence
		.split_ascii_whitespace()
		.filter(|el| !el.contains(' ') && !el.is_empty())
		.collect();
	pad_base_layer(&mut array);
	let array: Vec<u64> = array.iter().map(hash).collect();
	let mut hashes = calculate_layer(array);
	while hashes.len() != 1 {
		hashes = calculate_layer(hashes);
	}
	hashes[0]
}

fn calculate_layer(arr: Vec<u64>) -> Vec<u64> {
	if arr.len() == 1 {
		return arr;
	}
	let mut hashes: Vec<u64> = vec![];
	let mut i = 0;
	while hashes.len() != arr.len() / 2 {
		let left = arr[i];
		let right = arr[i + 1];
		let result: u64 = concatenate_hash_values(left, right);
		hashes.push(result);
		i += 2;
	}
	hashes
}

/// A representation of a sibling node along the Merkle path from the data
/// to the root. It is necessary to specify which side the sibling is on
/// so that the hash values can be combined in the same order.
enum SiblingNode {
	Left(HashValue),
	Right(HashValue),
}

/// Generates a Merkle proof that one particular word is contained
/// in the given sentence. You provide the sentence and the index of the word
/// which you want a proof.
///
/// Panics if the index is beyond the length of the word.
///
/// Example: I want to prove that the word "Trust" is in the sentence "Trust me, bro!"
/// So I call generate_proof("Trust me, bro!", 0)
/// And I get back the merkle root and list of intermediate nodes from which the
/// root can be reconstructed.
fn generate_proof(sentence: &str, index: usize) -> (HashValue, Vec<SiblingNode>) {
	let mut array: Vec<&str> = sentence
		.split_ascii_whitespace()
		.filter(|el| !el.contains(' ') && !el.is_empty())
		.collect();
	pad_base_layer(&mut array);
	let mut array: Vec<u64> = array.iter().map(hash).collect();
	if index > array.len() - 1 {
		panic!("index out of bounds");
	}
	let mut siblings: Vec<SiblingNode> = vec![];
	let mut current_sibling: u64;
	let mut current_hash: u64;

	let mut index = index;
	while array.len() != 1 {
		if index % 2 == 0 {
			current_sibling = array[index + 1];
			siblings.push(SiblingNode::Right(current_sibling));
			current_hash = concatenate_hash_values(array[index], current_sibling);
		} else {
			current_sibling = array[index - 1];
			siblings.push(SiblingNode::Left(current_sibling));
			current_hash = concatenate_hash_values(current_sibling, array[index]);
		}
		array = calculate_layer(array);
		index = array.iter().position(|x| *x == current_hash).unwrap();
	}

	(array[0], siblings)
}

/// Checks whether the given word is contained in a sentence, without knowing the whole sentence.
/// Rather we only know the merkle root of the sentence and a proof.
fn validate_proof(root: HashValue, word: &str, proof: Vec<SiblingNode>) -> bool {
	let inter_hash = hash(&word);
	let my_root: u64 = proof.iter().fold(inter_hash, |hash, sibling| match sibling {
		SiblingNode::Left(value) => concatenate_hash_values(*value, hash),
		SiblingNode::Right(value) => concatenate_hash_values(hash, *value),
	});
	my_root == root
}

#[test]
fn there_is_some_hash() {
	let h = calculate_merkle_root("Trust me, bro!");
	assert!(h > 0);
}

#[test]
fn padding() {
	let mut arr_4: Vec<&str> = "Trust me, bro!".split_ascii_whitespace().collect();
	let mut arr_8: Vec<&str> = "Trust me, bro bro bro!".split_ascii_whitespace().collect();
	pad_base_layer(&mut arr_4);
	pad_base_layer(&mut arr_8);
	assert_eq!(
		arr_4.len(),
		4,
		"You are not padding correctly. Expected 4 elements, found {}",
		arr_4.len()
	);
	assert_eq!(
		arr_8.len(),
		8,
		"You are not padding correctly. Expected 8 elements, found {}",
		arr_8.len()
	);
}

#[test]
fn different_hashes() {
	let h = calculate_merkle_root("Trust me, bro!");
	let h2 = calculate_merkle_root("Trust me, bro bro!");
	assert_ne!(h, h2, "Hashes for 3 word string is the same as for 4 word string");
}

#[test]
fn check_root_of_proof() {
	let h = calculate_merkle_root("Trust me, bro!");
  let result = generate_proof("Trust me, bro!", 1);
  assert_eq!(h, result.0, "Merkle root from generated proof is not the same as the original merkle root");
}

#[test]
fn check_different_roots_of_proof() {
	let h = calculate_merkle_root("Trust me, German!");
  let result = generate_proof("Trust me, bro!", 1);
  assert_ne!(h, result.0, "Merkle root from generated proof is the same as the merkle root for different sentence");
}

#[test]
fn check_merkle_proofs() {
  let proof = generate_proof("Trust me, bro!", 1);
  let result = validate_proof(proof.0, "me,", proof.1);
  assert!(result, "Your proof is incorrect. The word is in the sentence, but not identified by your solution");
}

#[test]
fn check_merkle_proofs_incorrect() {
  let proof = generate_proof("Trust me, bro!", 1);
  let result = validate_proof(proof.0, "hello", proof.1);
  assert!(!result, "Your proof is incorrect. The word is NOT in the sentence, but you validated it to be in");
}

No tests yet

Amazing!