merkle.rs

use std::{convert::AsRef, env, fmt, fs::read_to_string};

use sha2::{Digest as ShaDigest, Sha256};

pub type Bytes = Vec<u8>;

#[derive(Clone, Debug)]
pub struct Plaintext(pub Bytes);

impl AsRef<[u8]> for Plaintext {
    fn as_ref(&self) -> &[u8] {
        self.0.as_ref()
    }
}

#[derive(Clone, Debug)]
pub struct Digest(pub Bytes);

impl AsRef<[u8]> for Digest {
    fn as_ref(&self) -> &[u8] {
        self.0.as_ref()
    }
}

impl fmt::Display for Digest {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", hex::encode(self.0.clone()))
    }
}

pub fn concat_digests(left: MerkleTree, right: MerkleTree) -> Plaintext {
    Plaintext(
        [left.root().0.clone(), right.root().0.clone()]
            .iter()
            .flatten()
            .cloned()
            .collect(),
    )
}

pub fn somehow_hash<T>(x: T) -> Digest
where
    T: AsRef<[u8]>,
{
    let mut hasher = Sha256::new();
    hasher.update(x);
    Digest(hasher.finalize().to_vec())
}

#[derive(Clone, Debug)]
pub enum MerkleTree {
    Internal((Box<MerkleTree>, Box<MerkleTree>)),
    Leaf(Box<Digest>),
}

impl fmt::Display for MerkleTree {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::Leaf(x) => write!(f, "{}", hex::encode(*x.clone())),
            Self::Internal((left_child, right_child)) => {
                write!(f, "({}, {})", left_child, right_child)
            }
        }
    }
}

impl MerkleTree {
    pub fn new<T>(elems: Vec<T>) -> Self
    where
        T: AsRef<[u8]>,
    {
        let mut levels: Vec<Vec<MerkleTree>> =
            vec![Self::compute_leaves(elems)];

        while levels[0].len() > 1 {
            levels.insert(0, Self::compute_next_level(&levels[0]));
        }

        levels[0][0].clone()
    }

    fn compute_leaves<T>(elems: Vec<T>) -> Vec<MerkleTree>
    where
        T: AsRef<[u8]>,
    {
        elems
            .iter()
            .map(somehow_hash)
            .map(|digest| MerkleTree::Leaf(Box::new(digest)))
            .collect()
    }

    fn compute_next_level(level: &[MerkleTree]) -> Vec<MerkleTree> {
        level
            .chunks(2)
            .map(|xs| {
                MerkleTree::Internal((
                    Box::new(xs[0].clone()),
                    Box::new(xs[1].clone()),
                ))
            })
            .collect()
    }

    pub fn root(&self) -> Digest {
        match self {
            Self::Internal((left_child, right_child)) => somehow_hash(
                concat_digests(*left_child.clone(), *right_child.clone()),
            ),
            Self::Leaf(x) => *x.clone(),
        }
    }

    pub fn verify(&self, _proof: Vec<MerkleTree>) -> bool {
        todo!()
    }

    pub fn flatten(root: MerkleTree) -> Vec<MerkleTree> {
        let mut xs: Vec<MerkleTree> = vec![root.clone()];

        if let Self::Internal((left_child, right_child)) = root {
            xs.extend(Self::flatten(*left_child));
            xs.extend(Self::flatten(*right_child));
        }

        xs
    }
}

fn main() {
    if env::args().len() != 2 {
        eprintln!("usage: merkle FILE");
    } else {
        MerkleTree::flatten(MerkleTree::new(
            read_to_string(env::args().nth(1).unwrap())
                .expect("I/O failed")
                .lines()
                .map(|line| line.bytes().collect::<Vec<u8>>())
                .collect(),
        ))
        .iter()
        .map(|x| x.root())
        .enumerate()
        .for_each(|(i, x)| println!("({}, {})", i, x));
    }
}