rustix_bl/
left_threaded_avl_tree.rs

1// An attribute to hide warnings for unused code.
2#![allow(dead_code)]
3#![allow(unused_parens)]
4
5
6use std::collections::HashMap;
7use std::cmp;
8
9pub trait AVLTree {
10    fn empty() -> Self;
11    fn insert(&mut self, id: u32) -> bool;
12    fn increment_by_one(&mut self, id: u32) -> Option<u32>; // returns score if successful
13    fn remove(&mut self, id: u32) -> Option<u32>; // returns score if successful
14    fn extract_top(&self, n: usize) -> Vec<u32>;
15}
16
17#[derive(Debug, Serialize, Deserialize)]
18pub struct ScoredIdTreeMock {
19    ids: Vec<u32>,    //just to mock it
20    scores: Vec<u32>, //just to mock it
21}
22
23impl Default for ScoredIdTreeMock {
24    fn default() -> Self {
25        return ScoredIdTreeMock {
26            ids: Vec::new(),
27            scores: Vec::new(),
28        };
29    }
30}
31
32impl ScoredIdTreeMock {
33    fn index_of(&self, id: u32) -> Option<usize> {
34        for i in 0..(self.ids.len()) {
35            if self.ids[i] == id {
36                return Some(i);
37            }
38        }
39        return None;
40    }
41
42    fn score_sorted_copy(&self) -> Vec<u32> {
43        let mut hashmap: HashMap<u32, u32> = HashMap::new();
44        for i in 0..(self.ids.len()) {
45            hashmap.insert(self.ids[i], self.scores[i]);
46        }
47        let hm = &*(&mut hashmap);
48        let cmp = |a: &u32, b: &u32| hm.get(b).unwrap_or(&0u32).cmp(hm.get(a).unwrap_or(&0u32));
49        let mut x: Vec<u32> = self.ids.to_vec();
50        x.sort_by(cmp);
51        return x;
52    }
53}
54
55impl AVLTree for ScoredIdTreeMock {
56    fn empty() -> Self {
57        return ScoredIdTreeMock::default();
58    }
59
60    fn insert(&mut self, id: u32) -> bool {
61        if (!self.ids.contains(&id)) {
62            self.ids.push(id);
63            self.scores.push(0);
64            return true;
65        } else {
66            return false;
67        }
68    }
69
70    fn increment_by_one(&mut self, id: u32) -> Option<u32> {
71        let o = self.index_of(id);
72        return o.map(|i| {
73            self.scores[i] = self.scores[i] + 1;
74            self.scores[i]
75        });
76    }
77
78    fn remove(&mut self, id: u32) -> Option<u32> {
79        let o = self.index_of(id);
80        match o {
81            Some(i) => {
82                self.ids.remove(i);
83                return Some(self.scores.remove(i));
84            }
85            None => {
86                return None;
87            }
88        }
89    }
90
91    fn extract_top(&self, n: usize) -> Vec<u32> {
92        let n: usize = cmp::min(n, self.ids.len());
93        return self.score_sorted_copy()[0..(n)].to_vec();
94    }
95}
96
97#[cfg(test)]
98mod tests {
99
100    use left_threaded_avl_tree::ScoredIdTreeMock;
101    use left_threaded_avl_tree::AVLTree;
102
103    #[test]
104    fn always_works() {
105        assert!(true);
106    }
107    #[test]
108    fn basic_adding_works() {
109        let mut tree = ScoredIdTreeMock::empty();
110        assert!(tree.insert(1));
111        assert!(!tree.insert(1));
112        assert!(tree.insert(2));
113        assert!(!tree.insert(1));
114        assert!(tree.insert(3));
115        assert_eq!(tree.increment_by_one(2).unwrap(), 1);
116        assert_eq!(tree.increment_by_one(2).unwrap(), 2);
117        assert_eq!(tree.increment_by_one(2).unwrap(), 3);
118        assert_eq!(tree.increment_by_one(1).unwrap(), 1);
119        assert!(tree.increment_by_one(0).is_none());
120        let out = tree.extract_top(3);
121        assert_eq!(tree.ids.len(), 3);
122        assert_eq!(tree.scores.len(), 3);
123        assert_eq!(out.len(), 3);
124        assert_eq!(out[0], 2);
125        assert_eq!(out[1], 1);
126        assert_eq!(out[2], 3);
127    }
128    #[test]
129    fn basic_removal_works() {
130        let mut tree = ScoredIdTreeMock::default();
131        tree.insert(1);
132        tree.insert(1);
133        tree.insert(2);
134        tree.insert(1);
135        tree.insert(3);
136        tree.increment_by_one(2);
137        tree.increment_by_one(2);
138        tree.increment_by_one(1);
139        tree.increment_by_one(1);
140        tree.increment_by_one(1);
141        tree.increment_by_one(0);
142        tree.increment_by_one(1);
143        assert_eq!(tree.remove(2).unwrap(), 2);
144        assert_eq!(tree.remove(2).is_none(), true);
145        assert_eq!(tree.remove(1).unwrap(), 4);
146        let out = tree.extract_top(3);
147        assert_eq!(tree.ids.len(), 1);
148        assert_eq!(tree.scores.len(), 1);
149        assert_eq!(out.len(), 1);
150    }
151}