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

extern crate flat_tree;

#[derive(Clone, Debug)]
pub struct Node {
    index: u64,
    parent: u64,
    hash: Vec<u8>,
    data: Option<Vec<u8>>,
    size: usize,
}

pub struct Generator {
    roots: Vec<Node>,
    blocks: u64,
    leaf: fn(&Node, &Vec<Node>) -> Vec<u8>,
    parent: fn(&Node, &Node) -> Vec<u8>,
}

impl Generator {
    pub fn new(leaf: fn(&Node, &Vec<Node>) -> Vec<u8>,
               parent: fn(&Node, &Node) -> Vec<u8>)
               -> Generator {
        Generator {
            roots: vec![],
            blocks: 0,
            leaf: leaf,
            parent: parent,
        }
    }

    pub fn next(&mut self, data: Vec<u8>) -> Vec<Node> {
        self.blocks += 1;
        let mut nodes = vec![];
        let index = 2 * self.blocks;
        let len = data.len();
        let mut leaf = Node {
            index: index,
            parent: flat_tree::parent(index),
            hash: vec![],
            data: Some(data),
            size: len,
        };
        let hash = (self.leaf)(&leaf, &self.roots);
        leaf.hash = hash;
        self.roots.push(leaf.clone());
        nodes.push(leaf.clone());

        while self.roots.len() > 1 {
            let ref left = self.roots[self.roots.len() - 2].clone();
            let ref right = self.roots[self.roots.len() - 1].clone();

            if left.parent != right.parent {
                break;
            }

            self.roots.pop();
            let leaf = Node {
                index: left.parent,
                parent: flat_tree::parent(left.parent),
                hash: (self.parent)(left, right),
                size: left.size + right.size,
                data: None,
            };
            let len = self.roots.len();
            self.roots[len - 1] = leaf.clone();
            nodes.push(leaf.clone());
        }

        nodes
    }
}

extern crate ring;

#[cfg(test)]
mod tests {
    use ring::digest;
    use super::{Generator, Node};

    #[test]
    fn it_works() {
        let mut gen = Generator::new(leaf, parent);
        let nodes = gen.next(b"Hello World".to_vec());
        println!("{:?}", nodes);
    }

    fn parent(a: &Node, b: &Node) -> Vec<u8> {
        let mut data = a.data.clone().unwrap();
        data.append(&mut b.data.clone().unwrap());

        digest::digest(&digest::SHA256, data.as_slice())
            .as_ref()
            .to_vec()
    }

    fn leaf(leaf: &Node, roots: &Vec<Node>) -> Vec<u8> {
        let data = leaf.data.clone().unwrap();
        digest::digest(&digest::SHA256, data.as_slice()).as_ref().to_vec()
    }
}