1 use std::collections::BTreeMap;
2 
3 pub struct TopN<K: Ord, V> {
4     limit: Option<u32>,
5     count: u32,
6     echelons: BTreeMap<K, Vec<V>>,
7 }
8 
9 impl<K: Ord, V> TopN<K, V> {
10 
new(limit: u32) -> TopN<K, V>11     pub fn new(limit: u32) -> TopN<K, V> {
12         debug_assert_ne!(limit, 0);
13         TopN {
14             limit: Some(limit),
15             count: 0,
16             echelons: BTreeMap::new()
17         }
18     }
19 
limitless() -> TopN<K, V>20     pub fn limitless() -> TopN<K, V> {
21         TopN {
22             limit: None,
23             count: 0,
24             echelons: BTreeMap::new()
25         }
26     }
27 
insert(&mut self, k: K, v: V) -> Option<V> where K: Clone28     pub fn insert(&mut self, k: K, v: V) -> Option<V> where K: Clone {
29         self.count += 1;
30         self.echelons.entry(k).or_insert(Vec::new()).push(v);
31 
32         if let Some(limit) = self.limit {
33             if limit < self.count {
34                 self.count -= 1;
35 
36                 let last_key = self.echelons.iter().next_back().unwrap().0.clone();
37 
38                 let mut last_echelon = self.echelons.remove(&last_key).unwrap();
39                 let popped = last_echelon.pop().unwrap();
40                 if !last_echelon.is_empty() {
41                     self.echelons.insert(last_key, last_echelon);
42                 }
43                 return Some(popped);
44             }
45         }
46         None
47     }
48 
49     // see: https://github.com/rust-lang/rfcs/blob/master/text/1522-conservative-impl-trait.md
50 //    pub fn values(&self) -> impl Iterator<Item=&V> {
51 //        self.echelons.values().flat_map(|v| v)
52 //    }
values(&self) -> Vec<V> where V: Clone53     pub fn values(&self) -> Vec<V> where V: Clone {
54         self.echelons.values().flat_map(|v| v.iter().cloned()).collect()
55     }
56 }
57 
58 #[cfg(test)]
59 mod tests {
60     use super::*;
61 
62     #[test]
test_insert_one()63     fn test_insert_one() {
64         let mut top_n = TopN::new(5);
65         top_n.insert("asdf", 1);
66     }
67 
68     #[test]
test_insert_to_limit()69     fn test_insert_to_limit() {
70         let mut top_n = TopN::new(2);
71         top_n.insert("asdf", 1);
72         top_n.insert("xyz", 2);
73     }
74 
75     #[test]
test_insert_past_limit_bigger_discarded()76     fn test_insert_past_limit_bigger_discarded() {
77         let mut top_n = TopN::new(2);
78         top_n.insert("a", 1);
79         top_n.insert("b", 2);
80         top_n.insert("z", -1);
81         assert_eq!(top_n.values(), vec![1, 2]);
82     }
83 
84     #[test]
test_insert_past_limit_equal_discarded()85     fn test_insert_past_limit_equal_discarded() {
86         let mut top_n = TopN::new(2);
87         top_n.insert("a", 1);
88         top_n.insert("b", 2);
89         top_n.insert("b", -1);
90         assert_eq!(top_n.values(), vec![1, 2]);
91     }
92 
93     #[test]
test_insert_past_limit_smaller_last_one_discarded()94     fn test_insert_past_limit_smaller_last_one_discarded() {
95         let mut top_n = TopN::new(2);
96         top_n.insert("b", "second");
97         top_n.insert("c", "last");
98         top_n.insert("a", "first");
99         assert_eq!(top_n.values(), vec!["first", "second"]);
100     }
101 
102     #[test]
test_insert_past_limit_comprehensive()103     fn test_insert_past_limit_comprehensive() {
104         let mut top_n = TopN::new(5);
105         top_n.insert("asdf", 1);
106         assert_eq!(top_n.values(), vec![1]);
107         top_n.insert("asdf", 3);
108         assert_eq!(top_n.values(), vec![1, 3]);
109         top_n.insert("asdf", 3);
110         assert_eq!(top_n.values(), vec![1, 3, 3]);
111         top_n.insert("xyz", 4);
112         assert_eq!(top_n.values(), vec![1, 3, 3, 4]);
113         top_n.insert("asdf", 2);
114         assert_eq!(top_n.values(), vec![1, 3, 3, 2, 4]);
115         top_n.insert("xyz", 5);
116         assert_eq!(top_n.values(), vec![1, 3, 3, 2, 4]);
117         top_n.insert("asdf", -1);
118         assert_eq!(top_n.values(), vec![1, 3, 3, 2, -1]);
119     }
120 
121     #[test]
test_limitless()122     fn test_limitless() {
123         let mut top_n = TopN::limitless();
124         top_n.insert("z", 3);
125         top_n.insert("y", 2);
126         top_n.insert("a", 1);
127         top_n.insert("a", 0);
128         assert_eq!(top_n.values(), vec![1, 0, 2, 3]);
129     }
130 }
131 
132