tofuri_tree/
lib.rs

1use std::cmp::Ordering;
2use std::collections::HashMap;
3use tofuri_core::*;
4#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
5pub struct Branch {
6    pub hash: Hash,
7    pub height: usize,
8    pub timestamp: u32,
9}
10impl Branch {
11    fn new(hash: Hash, height: usize, timestamp: u32) -> Branch {
12        Branch { hash, height, timestamp }
13    }
14}
15#[derive(Default, Debug, Clone)]
16pub struct Tree {
17    branches: Vec<Branch>,
18    hashes: HashMap<Hash, Hash>,
19}
20impl Tree {
21    pub fn main(&self) -> Option<&Branch> {
22        self.branches.first()
23    }
24    pub fn size(&self) -> usize {
25        self.hashes.len()
26    }
27    pub fn hashes(&self, trust_fork_after_blocks: usize) -> (Vec<Hash>, Vec<Hash>) {
28        let mut trusted = vec![];
29        if let Some(main) = self.main() {
30            let mut hash = main.hash;
31            loop {
32                trusted.push(hash);
33                match self.get(&hash) {
34                    Some(previous_hash) => hash = *previous_hash,
35                    None => break,
36                };
37            }
38        }
39        if let Some(hash) = trusted.last() {
40            if hash != &[0; 32] {
41                panic!("broken chain")
42            }
43            trusted.pop();
44        }
45        trusted.reverse();
46        let len = trusted.len();
47        let start = if len < trust_fork_after_blocks { 0 } else { len - trust_fork_after_blocks };
48        let dynamic = trusted.drain(start..len).collect();
49        (trusted, dynamic)
50    }
51    pub fn hashes_dynamic(&self, trust_fork_after_blocks: usize) -> Vec<Hash> {
52        let mut vec = vec![];
53        if let Some(main) = self.main() {
54            let mut hash = main.hash;
55            for _ in 0..trust_fork_after_blocks {
56                vec.push(hash);
57                match self.get(&hash) {
58                    Some(previous_hash) => hash = *previous_hash,
59                    None => break,
60                };
61            }
62        }
63        if let Some(hash) = vec.last() {
64            if hash == &[0; 32] {
65                vec.pop();
66            }
67        }
68        vec.reverse();
69        vec
70    }
71    pub fn get(&self, hash: &Hash) -> Option<&Hash> {
72        self.hashes.get(hash)
73    }
74    pub fn insert(&mut self, hash: Hash, previous_hash: Hash, timestamp: u32) -> Option<bool> {
75        if self.hashes.insert(hash, previous_hash).is_some() {
76            return None;
77        }
78        if let Some(index) = self.branches.iter().position(|a| a.hash == previous_hash) {
79            // extend branch
80            self.branches[index] = Branch::new(hash, self.branches[index].height + 1, timestamp);
81            Some(false)
82        } else {
83            // new branch
84            self.branches.push(Branch::new(hash, self.height(&previous_hash), timestamp));
85            Some(true)
86        }
87    }
88    pub fn sort_branches(&mut self) {
89        self.branches.sort_by(|a, b| match b.height.cmp(&a.height) {
90            Ordering::Equal => a.timestamp.cmp(&b.timestamp),
91            x => x,
92        });
93    }
94    pub fn clear(&mut self) {
95        self.branches.clear();
96        self.hashes.clear();
97    }
98    fn height(&self, previous_hash: &Hash) -> usize {
99        let mut hash = previous_hash;
100        let mut height = 0;
101        loop {
102            match self.hashes.get(hash) {
103                Some(previous_hash) => {
104                    hash = previous_hash;
105                    height += 1;
106                }
107                None => {
108                    if hash != &[0; 32] {
109                        panic!("broken chain")
110                    }
111                    break;
112                }
113            };
114        }
115        height
116    }
117}
118#[cfg(test)]
119mod tests {
120    use super::*;
121    #[test]
122    fn test() {
123        let mut tree = Tree::default();
124        tree.insert([0x11; 32], [0x00; 32], 1);
125        tree.insert([0x22; 32], [0x11; 32], 1);
126        tree.insert([0x33; 32], [0x22; 32], 1);
127        assert_eq!(tree.size(), 3);
128        tree.insert([0x44; 32], [0x33; 32], 1);
129        tree.insert([0x55; 32], [0x22; 32], 1);
130        tree.insert([0x66; 32], [0x00; 32], 1);
131        tree.insert([0x77; 32], [0x55; 32], 0);
132        assert_eq!(tree.main(), Some(&Branch::new([0x44; 32], 3, 1)));
133        tree.sort_branches();
134        assert_eq!(tree.main(), Some(&Branch::new([0x77; 32], 3, 0)));
135        assert_eq!(tree.size(), 7);
136    }
137}
tofuri_tree/lib.rs

tofuri_tree/
lib.rs
