blob: 14c14fc6e403814005ec8b343496e3b8251fe442 [file] [log] [blame]
Jakub Kotura425e552020-12-21 17:28:15 +01001
2use std::collections::HashMap;
3use std::collections::hash_map::{Entry};
4use std::hash::Hash;
5use std::fmt;
6
7/// An iterator adapter to filter out duplicate elements.
8///
9/// See [`.unique_by()`](../trait.Itertools.html#method.unique) for more information.
10#[derive(Clone)]
11#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
12pub struct UniqueBy<I: Iterator, V, F> {
13 iter: I,
14 // Use a hashmap for the entry API
15 used: HashMap<V, ()>,
16 f: F,
17}
18
19impl<I, V, F> fmt::Debug for UniqueBy<I, V, F>
20 where I: Iterator + fmt::Debug,
21 V: fmt::Debug + Hash + Eq,
22{
23 debug_fmt_fields!(UniqueBy, iter, used);
24}
25
26/// Create a new `UniqueBy` iterator.
27pub fn unique_by<I, V, F>(iter: I, f: F) -> UniqueBy<I, V, F>
28 where V: Eq + Hash,
29 F: FnMut(&I::Item) -> V,
30 I: Iterator,
31{
32 UniqueBy {
33 iter,
34 used: HashMap::new(),
35 f,
36 }
37}
38
39// count the number of new unique keys in iterable (`used` is the set already seen)
40fn count_new_keys<I, K>(mut used: HashMap<K, ()>, iterable: I) -> usize
41 where I: IntoIterator<Item=K>,
42 K: Hash + Eq,
43{
44 let iter = iterable.into_iter();
45 let current_used = used.len();
46 used.extend(iter.map(|key| (key, ())));
47 used.len() - current_used
48}
49
50impl<I, V, F> Iterator for UniqueBy<I, V, F>
51 where I: Iterator,
52 V: Eq + Hash,
53 F: FnMut(&I::Item) -> V
54{
55 type Item = I::Item;
56
Joel Galenson6f798712021-04-01 17:03:06 -070057 fn next(&mut self) -> Option<Self::Item> {
Jakub Kotura425e552020-12-21 17:28:15 +010058 while let Some(v) = self.iter.next() {
59 let key = (self.f)(&v);
60 if self.used.insert(key, ()).is_none() {
61 return Some(v);
62 }
63 }
64 None
65 }
66
67 #[inline]
68 fn size_hint(&self) -> (usize, Option<usize>) {
69 let (low, hi) = self.iter.size_hint();
70 ((low > 0 && self.used.is_empty()) as usize, hi)
71 }
72
73 fn count(self) -> usize {
74 let mut key_f = self.f;
75 count_new_keys(self.used, self.iter.map(move |elt| key_f(&elt)))
76 }
77}
78
Joel Galenson6f798712021-04-01 17:03:06 -070079impl<I, V, F> DoubleEndedIterator for UniqueBy<I, V, F>
80 where I: DoubleEndedIterator,
81 V: Eq + Hash,
82 F: FnMut(&I::Item) -> V
83{
84 fn next_back(&mut self) -> Option<Self::Item> {
85 while let Some(v) = self.iter.next_back() {
86 let key = (self.f)(&v);
87 if self.used.insert(key, ()).is_none() {
88 return Some(v);
89 }
90 }
91 None
92 }
93}
94
Jakub Kotura425e552020-12-21 17:28:15 +010095impl<I> Iterator for Unique<I>
96 where I: Iterator,
97 I::Item: Eq + Hash + Clone
98{
99 type Item = I::Item;
100
Joel Galenson6f798712021-04-01 17:03:06 -0700101 fn next(&mut self) -> Option<Self::Item> {
Jakub Kotura425e552020-12-21 17:28:15 +0100102 while let Some(v) = self.iter.iter.next() {
103 if let Entry::Vacant(entry) = self.iter.used.entry(v) {
104 let elt = entry.key().clone();
105 entry.insert(());
106 return Some(elt);
107 }
108 }
109 None
110 }
111
112 #[inline]
113 fn size_hint(&self) -> (usize, Option<usize>) {
114 let (low, hi) = self.iter.iter.size_hint();
115 ((low > 0 && self.iter.used.is_empty()) as usize, hi)
116 }
117
118 fn count(self) -> usize {
119 count_new_keys(self.iter.used, self.iter.iter)
120 }
121}
122
Joel Galenson6f798712021-04-01 17:03:06 -0700123impl<I> DoubleEndedIterator for Unique<I>
124 where I: DoubleEndedIterator,
125 I::Item: Eq + Hash + Clone
126{
127 fn next_back(&mut self) -> Option<Self::Item> {
128 while let Some(v) = self.iter.iter.next_back() {
129 if let Entry::Vacant(entry) = self.iter.used.entry(v) {
130 let elt = entry.key().clone();
131 entry.insert(());
132 return Some(elt);
133 }
134 }
135 None
136 }
137}
138
Jakub Kotura425e552020-12-21 17:28:15 +0100139/// An iterator adapter to filter out duplicate elements.
140///
141/// See [`.unique()`](../trait.Itertools.html#method.unique) for more information.
142#[derive(Clone)]
143#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
144pub struct Unique<I: Iterator> {
145 iter: UniqueBy<I, I::Item, ()>,
146}
147
148impl<I> fmt::Debug for Unique<I>
149 where I: Iterator + fmt::Debug,
150 I::Item: Hash + Eq + fmt::Debug,
151{
152 debug_fmt_fields!(Unique, iter);
153}
154
155pub fn unique<I>(iter: I) -> Unique<I>
156 where I: Iterator,
157 I::Item: Eq + Hash,
158{
159 Unique {
160 iter: UniqueBy {
161 iter,
162 used: HashMap::new(),
163 f: (),
164 }
165 }
166}