starlark_map/
unordered_set.rs

1/*
2 * Copyright 2019 The Starlark in Rust Authors.
3 * Copyright (c) Facebook, Inc. and its affiliates.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 *
9 *     https://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 */
17
18//! `HashSet` that does not expose insertion order.
19
20use std::hash::Hash;
21
22use allocative::Allocative;
23
24use crate::unordered_map;
25use crate::unordered_map::UnorderedMap;
26use crate::Equivalent;
27use crate::Hashed;
28use crate::StarlarkHashValue;
29
30/// `HashSet` that does not expose insertion order.
31#[derive(Clone, Allocative, Debug)]
32pub struct UnorderedSet<T> {
33    map: UnorderedMap<T, ()>,
34}
35
36impl<T> Default for UnorderedSet<T> {
37    #[inline]
38    fn default() -> UnorderedSet<T> {
39        UnorderedSet::new()
40    }
41}
42
43impl<T> UnorderedSet<T> {
44    /// Create a new empty set.
45    #[inline]
46    pub const fn new() -> UnorderedSet<T> {
47        UnorderedSet {
48            map: UnorderedMap::new(),
49        }
50    }
51
52    /// Create a new empty set with the specified capacity.
53    #[inline]
54    pub fn with_capacity(n: usize) -> UnorderedSet<T> {
55        UnorderedSet {
56            map: UnorderedMap::with_capacity(n),
57        }
58    }
59
60    /// Insert a value into the set.
61    #[inline]
62    pub fn insert(&mut self, k: T) -> bool
63    where
64        T: Hash + Eq,
65    {
66        self.map.insert(k, ()).is_none()
67    }
68
69    /// Clear the set, removing all values.
70    #[inline]
71    pub fn clear(&mut self) {
72        self.map.clear();
73    }
74
75    /// Is the set empty?
76    #[inline]
77    pub fn is_empty(&self) -> bool {
78        self.map.is_empty()
79    }
80
81    /// Get the number of elements in the set.
82    #[inline]
83    pub fn len(&self) -> usize {
84        self.map.len()
85    }
86
87    /// Does the set contain the specified value?
88    #[inline]
89    pub fn contains<Q>(&self, value: &Q) -> bool
90    where
91        Q: Hash + Equivalent<T> + ?Sized,
92    {
93        self.map.contains_key(value)
94    }
95
96    /// Does the set contain the specified value?
97    #[inline]
98    pub fn contains_hashed<Q>(&self, value: Hashed<&Q>) -> bool
99    where
100        Q: Equivalent<T> + ?Sized,
101    {
102        self.map.contains_key_hashed(value)
103    }
104
105    /// Lower-level access to the underlying map.
106    #[inline]
107    pub fn raw_entry_mut(&mut self) -> RawEntryBuilderMut<T> {
108        RawEntryBuilderMut {
109            entry: self.map.raw_entry_mut(),
110        }
111    }
112
113    /// This function is private.
114    fn iter(&self) -> impl Iterator<Item = &T> {
115        self.map.entries_unordered().map(|(k, _)| k)
116    }
117
118    /// Get the entries in the set, sorted.
119    pub fn entries_sorted(&self) -> Vec<&T>
120    where
121        T: Ord,
122    {
123        let mut entries = Vec::from_iter(self.iter());
124        entries.sort_unstable();
125        entries
126    }
127}
128
129impl<T: Eq + Hash> PartialEq for UnorderedSet<T> {
130    #[inline]
131    fn eq(&self, other: &Self) -> bool {
132        self.map == other.map
133    }
134}
135
136impl<T: Eq + Hash> Eq for UnorderedSet<T> {}
137
138impl<T: Eq + Hash> FromIterator<T> for UnorderedSet<T> {
139    #[inline]
140    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> UnorderedSet<T> {
141        UnorderedSet {
142            map: UnorderedMap::from_iter(iter.into_iter().map(|v| (v, ()))),
143        }
144    }
145}
146
147/// Builder for [`RawEntryMut`].
148pub struct RawEntryBuilderMut<'a, T> {
149    entry: unordered_map::RawEntryBuilderMut<'a, T, ()>,
150}
151
152impl<'a, T> RawEntryBuilderMut<'a, T> {
153    /// Find the entry for a key.
154    #[inline]
155    pub fn from_entry<Q>(self, entry: &Q) -> RawEntryMut<'a, T>
156    where
157        Q: Hash + Equivalent<T> + ?Sized,
158    {
159        let entry = Hashed::new(entry);
160        self.from_entry_hashed(entry)
161    }
162
163    /// Find the entry for a key.
164    #[inline]
165    pub fn from_entry_hashed<Q>(self, entry: Hashed<&Q>) -> RawEntryMut<'a, T>
166    where
167        Q: ?Sized + Equivalent<T>,
168    {
169        self.from_hash(entry.hash(), |k| entry.key().equivalent(k))
170    }
171
172    /// Find the entry by hash and equality function.
173    #[inline]
174    pub fn from_hash<F>(self, hash: StarlarkHashValue, is_match: F) -> RawEntryMut<'a, T>
175    where
176        F: for<'b> FnMut(&'b T) -> bool,
177    {
178        match self.entry.from_hash(hash, is_match) {
179            unordered_map::RawEntryMut::Occupied(e) => {
180                RawEntryMut::Occupied(RawOccupiedEntryMut { entry: e })
181            }
182            unordered_map::RawEntryMut::Vacant(e) => {
183                RawEntryMut::Vacant(RawVacantEntryMut { entry: e })
184            }
185        }
186    }
187}
188
189/// Reference to an occupied entry in a [`UnorderedSet`].
190pub struct RawOccupiedEntryMut<'a, T> {
191    entry: unordered_map::RawOccupiedEntryMut<'a, T, ()>,
192}
193
194/// Reference to a vacant entry in a [`UnorderedSet`].
195pub struct RawVacantEntryMut<'a, T> {
196    entry: unordered_map::RawVacantEntryMut<'a, T, ()>,
197}
198
199/// Reference to an entry in a [`UnorderedSet`].
200pub enum RawEntryMut<'a, T> {
201    /// Occupied entry.
202    Occupied(RawOccupiedEntryMut<'a, T>),
203    /// Vacant entry.
204    Vacant(RawVacantEntryMut<'a, T>),
205}
206
207impl<'a, T> RawOccupiedEntryMut<'a, T> {
208    /// Remove the entry.
209    #[inline]
210    pub fn remove(self) -> T {
211        self.entry.remove_entry().0
212    }
213
214    /// Replace the entry.
215    #[inline]
216    pub fn insert(&mut self, value: T) -> T {
217        self.entry.insert_key(value)
218    }
219}
220
221impl<'a, T> RawVacantEntryMut<'a, T> {
222    /// Insert an entry to the set. This function computes the hash of the key.
223    #[inline]
224    pub fn insert(self, value: T)
225    where
226        T: Hash,
227    {
228        let value = Hashed::new(value);
229        self.insert_hashed(value);
230    }
231
232    /// Insert an entry to the set.
233    #[inline]
234    pub fn insert_hashed(self, value: Hashed<T>)
235    where
236        T: Hash,
237    {
238        self.entry.insert_hashed(value, ());
239    }
240}
241
242#[cfg(test)]
243mod tests {
244    use crate::unordered_set::UnorderedSet;
245
246    #[test]
247    fn test_insert() {
248        let mut set = UnorderedSet::new();
249        assert!(set.insert(10));
250        assert!(!set.insert(10));
251        assert!(set.insert(20));
252        assert!(!set.insert(20));
253        assert_eq!(set.len(), 2);
254    }
255
256    #[test]
257    fn test_entries_sorted() {
258        let mut set = UnorderedSet::new();
259        set.insert(20);
260        set.insert(10);
261        set.insert(30);
262        assert_eq!(set.entries_sorted(), vec![&10, &20, &30]);
263    }
264}