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