Skip to content

Instantly share code, notes, and snippets.

@sug0
Created September 12, 2022 07:42
Show Gist options
  • Save sug0/744b5a0e4a6efd45a604e9376d2e10bd to your computer and use it in GitHub Desktop.
Save sug0/744b5a0e4a6efd45a604e9376d2e10bd to your computer and use it in GitHub Desktop.
Simple Merkle tree in Rust
#![feature(iter_array_chunks)]
#![feature(hasher_prefixfree_extras)]
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
#[derive(Clone, Debug)]
struct Tree<T> {
root: Option<Node<T>>,
}
#[derive(Clone, Debug, Hash)]
enum Node<T> {
Leaf(Box<Leaf<T>>),
Branch(Box<Branch<T>>),
}
#[derive(Clone, Debug)]
struct Leaf<T> {
hash: u64,
attrib: T,
}
#[derive(Clone, Debug)]
struct Branch<T> {
hash: u64,
left: Node<T>,
right: Option<Node<T>>,
}
fn main() {
Tree::merkelize([
"ganda cena mano",
"de broa velho",
"aew mermau",
"parabens aew velho",
"caraca",
"mlk",
"meu deus",
"nossa senhora",
"vai embora do americo",
])
.display_dot();
}
impl<T> Tree<T> {
fn merkelize<A>(attributes: A) -> Self
where
A: IntoIterator<Item = T>,
T: Hash,
{
let mut nodes: Vec<_> = attributes
.into_iter()
.map(|a| Node::Leaf(Box::new(Leaf::new(a))))
.collect();
if nodes.is_empty() {
return Self { root: None };
}
while nodes.len() != 1 {
nodes = Node::reparent(nodes);
}
Self { root: nodes.pop() }
}
fn display_dot(&self) {
println!("digraph BST {{");
if let Some(node) = &self.root {
let mut nulls = 0;
node.display_dot_inner(&mut nulls);
}
println!("}}");
}
}
impl<T> Node<T> {
fn display_dot_inner(&self, nulls: &mut usize) {
match self {
Node::Leaf(leaf) => {
let hash = leaf.hash;
let id = *nulls;
*nulls += 1;
println!(" {hash} -> null{id};");
println!(" null{id} [shape=point];");
}
Node::Branch(branch) => {
let h1 = branch.hash;
let h2 = branch.left.get_hash();
println!(" {h1} -> {h2};");
branch.left.display_dot_inner(nulls);
if let Some(right) = &branch.right {
let h2 = right.get_hash();
println!(" {h1} -> {h2};");
right.display_dot_inner(nulls);
} else {
let id = *nulls;
*nulls += 1;
println!(" {h1} -> null{id};");
println!(" null{id} [shape=point];");
}
}
}
}
fn get_hash(&self) -> u64 {
match self {
Node::Leaf(leaf) => leaf.hash,
Node::Branch(branch) => branch.hash,
}
}
fn reparent(nodes: Vec<Node<T>>) -> Vec<Node<T>>
where
T: Hash,
{
// TODO: optimize this; we can re-use the same storage
let mut new_nodes = Vec::with_capacity(nodes.len() >> 1);
let mut chunks = nodes.into_iter().array_chunks();
for [left, right] in chunks.by_ref() {
new_nodes.push(Node::Branch(Box::new(Branch::new(left, Some(right)))));
}
if let Some(last_node) = chunks
.into_remainder()
.and_then(|mut last_node| last_node.next())
{
new_nodes.push(Node::Branch(Box::new(Branch::singleton(last_node))));
}
new_nodes
}
}
impl<T: Hash> Leaf<T> {
fn new(attrib: T) -> Self {
let hash = {
let mut hasher = DefaultHasher::new();
attrib.hash(&mut hasher);
hasher.finish()
};
Self { hash, attrib }
}
}
impl<T: Hash> Branch<T> {
fn new(left: Node<T>, right: Option<Node<T>>) -> Self {
let mut branch = Self {
hash: 0,
left,
right,
};
branch.hash = {
let mut hasher = DefaultHasher::new();
branch.hash(&mut hasher);
hasher.finish()
};
branch
}
fn singleton(node: Node<T>) -> Self {
Self::new(node, None)
}
}
impl<T> Hash for Branch<T> {
fn hash<H: Hasher>(&self, state: &mut H) {
state.write_u64(self.left.get_hash());
if let Some(right) = &self.right {
state.write_u64(right.get_hash());
} else {
state.write_u64(self.left.get_hash());
}
}
}
impl<T: Hash> Hash for Leaf<T> {
fn hash<H: Hasher>(&self, state: &mut H) {
self.attrib.hash(state);
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment