merkle.rs

use std::sync::Arc;

use ring::digest;

static HASH_ALGORITH: &'static digest::Algorithm = &digest::SHA256;

pub enum TreeNode {
    Empty { hash: digest::Digest },
    Leaf { hash: digest::Digest, data: Vec<u8> },
    SubTree {
        hash: digest::Digest,
        left: Arc<Box<TreeNode>>,
        right: Arc<Box<TreeNode>>,
    },
}

fn hash_empty() -> digest::Digest {
    return digest::Context::new(HASH_ALGORITH).finish();
}

fn hash_leaf(data: &Vec<u8>) -> digest::Digest {
    let mut ctx = digest::Context::new(HASH_ALGORITH);
    ctx.update(b"\x00");
    ctx.update(&data);
    return ctx.finish();
}

fn hash_children(left: &Arc<Box<TreeNode>>, right: &Arc<Box<TreeNode>>) -> digest::Digest {
    let mut ctx = digest::Context::new(HASH_ALGORITH);
    ctx.update(b"\x01");
    ctx.update(left.hash().as_ref());
    ctx.update(right.hash().as_ref());
    return ctx.finish();
}

impl TreeNode {
    pub fn new_empty() -> TreeNode {
        return TreeNode::Empty { hash: hash_empty() };
    }

    pub fn new_leaf(data: Vec<u8>) -> TreeNode {
        return TreeNode::Leaf {
            hash: hash_leaf(&data),
            data: data,
        };
    }

    pub fn new_tree(left: Arc<Box<TreeNode>>, right: Arc<Box<TreeNode>>) -> TreeNode {
        return TreeNode::SubTree {
            hash: hash_children(&left, &right),
            left: left,
            right: right,
        };
    }

    pub fn hash(&self) -> digest::Digest {
        match self {
            &TreeNode::Empty { hash } => hash,
            &TreeNode::Leaf { hash, .. } => hash,
            &TreeNode::SubTree { hash, .. } => hash,
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::{atomic, mpsc, Arc};
    use std::thread;

    use rand::{thread_rng, Rng};
    use test;

    use super::TreeNode;

    fn make_leaves() -> Vec<Arc<Box<TreeNode>>> {
        (0..102400)
            .map(|_| {
                let mut data = [0u8; 2048];
                thread_rng().fill_bytes(&mut data);
                Arc::new(Box::new(TreeNode::new_leaf(data.to_vec())))
            })
            .collect()
    }

    #[bench]
    fn bench_serial_build(b: &mut test::Bencher) {
        let leaves = make_leaves();
        b.iter(|| {
            let mut root = Arc::new(Box::new(TreeNode::new_empty()));
            for leaf in leaves.iter() {
                root = Arc::new(Box::new(TreeNode::new_tree(root, leaf.clone())));
            }
            test::black_box(root);
        })
    }

    #[bench]
    fn bench_parallel_build(b: &mut test::Bencher) {
        let leaves = make_leaves();
        b.iter(|| {
            const PARALLELISM: usize = 4;
            let (tree_tx, tree_rx) = mpsc::channel();
            let mut producers = vec![];
            let live_threads = Arc::new(atomic::AtomicUsize::new(PARALLELISM));
            for _ in 0..PARALLELISM {
                let (tx, rx) = mpsc::channel();
                producers.push(tx);
                let tree_tx = tree_tx.clone();
                let live_threads = live_threads.clone();
                thread::spawn(move || {
                    let mut finished = false;
                    while !finished {
                        let mut root = Arc::new(Box::new(TreeNode::new_empty()));
                        let mut found = false;
                        loop {
                            match rx.try_recv() {
                                Ok(leaf) => {
                                    root = Arc::new(Box::new(TreeNode::new_tree(root, leaf)));
                                    found = true;
                                }
                                Err(mpsc::TryRecvError::Empty) => {
                                    break;
                                }
                                Err(mpsc::TryRecvError::Disconnected) => {
                                    finished = true;
                                    break;
                                }
                            }
                        }
                        if found {
                            tree_tx.send(root).unwrap();
                        }
                    }
                    live_threads.fetch_sub(1, atomic::Ordering::Relaxed);
                });
            }

            for leaf in leaves.iter() {
                let idx = (leaf.hash().as_ref()[0] as usize) % producers.len();
                producers[idx].send(leaf.clone()).unwrap();
            }
            drop(producers);
            let mut root = Arc::new(Box::new(TreeNode::new_empty()));
            while live_threads.load(atomic::Ordering::Relaxed) > 0 {
                match tree_rx.try_recv() {
                    Ok(leaf) => {
                        root = Arc::new(Box::new(TreeNode::new_tree(root, leaf)));
                    }
                    Err(_) => {}
                }
            }
        })
    }
}