1
use core::panic;
2
use std::fmt;
3
use std::hash::BuildHasher;
4
use std::hash::Hash;
5
use std::ops::Deref;
6
use std::ops::Index;
7
use std::ops::IndexMut;
8

            
9
use hashbrown::Equivalent;
10
use hashbrown::HashTable;
11
use rustc_hash::FxBuildHasher;
12

            
13
use merc_utilities::GenerationCounter;
14
use merc_utilities::GenerationalIndex;
15
use merc_utilities::cast;
16

            
17
/// A type-safe index for use with [IndexedSet]. Uses generational indices in debug builds to assert
18
/// correct usage of indices.
19
#[repr(transparent)]
20
#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Default)]
21
pub struct SetIndex(GenerationalIndex<usize>);
22

            
23
impl Deref for SetIndex {
24
    type Target = usize;
25

            
26
103400904
    fn deref(&self) -> &Self::Target {
27
103400904
        &self.0
28
103400904
    }
29
}
30

            
31
impl fmt::Debug for SetIndex {
32
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
33
        write!(f, "SetIndex({})", self.0)
34
    }
35
}
36

            
37
impl fmt::Display for SetIndex {
38
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
39
        write!(f, "{}", self.0)
40
    }
41
}
42

            
43
/// A set that assigns a unique index to every entry. The returned index can be used to access the inserted entry.
44
pub struct IndexedSet<T, S = FxBuildHasher> {
45
    /// The table of elements, which can be either filled or empty.
46
    table: Vec<IndexSetEntry<T>>,
47
    /// Indexes of the elements in the set. Stores only the index; the hash is recomputed via the
48
    /// stored `hasher` on lookup using [`HashTable`]'s explicit-hash API.
49
    index: HashTable<usize>,
50
    /// A list of free nodes, where the value is the first free node.
51
    free: Option<usize>,
52
    /// The number of generations
53
    generation_counter: GenerationCounter,
54
    /// The hasher used to compute hashes for elements
55
    hasher: S,
56
}
57

            
58
/// An entry in the indexed set, which can either be filled or empty.
59
enum IndexSetEntry<T> {
60
    Filled(T),
61
    Empty(usize),
62
}
63

            
64
impl<T, S: BuildHasher + Default> IndexedSet<T, S> {
65
    /// Creates a new empty IndexedSet with the default hasher.
66
3771
    pub fn new() -> IndexedSet<T, S> {
67
3771
        IndexedSet {
68
3771
            table: Vec::default(),
69
3771
            index: HashTable::new(),
70
3771
            free: None,
71
3771
            generation_counter: GenerationCounter::new(),
72
3771
            hasher: S::default(),
73
3771
        }
74
3771
    }
75
}
76

            
77
impl<T, S> IndexedSet<T, S> {
78
    /// Creates a new empty IndexedSet with the specified hasher.
79
    pub fn with_hasher(hash_builder: S) -> IndexedSet<T, S> {
80
        IndexedSet {
81
            table: Vec::default(),
82
            index: HashTable::new(),
83
            free: None,
84
            generation_counter: GenerationCounter::new(),
85
            hasher: hash_builder,
86
        }
87
    }
88

            
89
    /// Returns the number of elements in the set.
90
1357197
    pub fn len(&self) -> usize {
91
1357197
        self.index.len()
92
1357197
    }
93

            
94
    /// Returns true if the set is empty.
95
    pub fn is_empty(&self) -> bool {
96
        self.len() == 0
97
    }
98

            
99
    /// Returns a reference to the element at the given index, if it exists.
100
2200
    pub fn get(&self, index: SetIndex) -> Option<&T> {
101
2200
        if let Some(entry) = self.table.get(self.generation_counter.get_index(index.0)) {
102
2200
            match entry {
103
2200
                IndexSetEntry::Filled(element) => Some(element),
104
                IndexSetEntry::Empty(_) => None,
105
            }
106
        } else {
107
            None
108
        }
109
2200
    }
110

            
111
    /// Returns a reference to the element at the given raw table index, if a
112
    /// filled entry exists there.
113
    ///
114
    /// Unlike [Self::get] this takes a plain `usize` and performs no generation
115
    /// check; it only verifies that the slot is in bounds and filled.
116
3
    pub fn get_by_index(&self, index: usize) -> Option<&T> {
117
3
        if let Some(entry) = self.table.get(index) {
118
3
            match entry {
119
3
                IndexSetEntry::Filled(element) => Some(element),
120
                IndexSetEntry::Empty(_) => None,
121
            }
122
        } else {
123
            None
124
        }
125
3
    }
126

            
127
    /// Returns the capacity of the set.
128
    pub fn capacity(&self) -> usize {
129
        self.table.capacity()
130
    }
131

            
132
    /// Returns an iterator over the elements in the set.
133
23073
    pub fn iter(&self) -> Iter<'_, T, S> {
134
23073
        Iter {
135
23073
            reference: self,
136
23073
            index: 0,
137
23073
            generation_counter: &self.generation_counter,
138
23073
        }
139
23073
    }
140
}
141

            
142
impl<T: Clone, S> IndexedSet<T, S> {
143
    /// Returns a vector containing all elements of this indexed set.
144
100
    pub fn to_vec(&self) -> Vec<T> {
145
300
        self.iter().map(|(_, entry)| entry.clone()).collect()
146
100
    }
147
}
148

            
149
impl<T: Hash + Eq, S: BuildHasher> IndexedSet<T, S> {
150
    /// Inserts the given element into the set
151
    ///
152
    /// Returns the corresponding index and a boolean indicating if the element was inserted.
153
    pub fn insert_equiv<'a, Q>(&mut self, value: &'a Q) -> (SetIndex, bool)
154
    where
155
        Q: Hash + Equivalent<T>,
156
        T: From<&'a Q>,
157
    {
158
        let hash = self.hasher.hash_one(value);
159

            
160
        if let Some(&existing) = self.index.find(hash, |&i| entry_matches(&self.table, i, value)) {
161
            // The element is already in the set, so return the index.
162
            return (SetIndex(self.generation_counter.recall_index(existing)), false);
163
        }
164

            
165
        let value: T = value.into();
166
        debug_assert_eq!(hash, self.hasher.hash_one(&value), "Hash values should be the same");
167

            
168
        let index = self.insert_into_table(value);
169
        let hasher_ref = &self.hasher;
170
        let table_ref = &self.table;
171
        self.index
172
            .insert_unique(hash, index, |&i| entry_hash(table_ref, hasher_ref, i));
173
        (SetIndex(self.generation_counter.create_index(index)), true)
174
    }
175

            
176
    /// Inserts the given element into the set
177
    ///
178
    /// Returns the corresponding index and a boolean indicating if the element was inserted.
179
764820
    pub fn insert(&mut self, value: T) -> (SetIndex, bool) {
180
764820
        let hash = self.hasher.hash_one(&value);
181

            
182
764820
        if let Some(&existing) = self.index.find(hash, |&i| entry_matches(&self.table, i, &value)) {
183
            // The element is already in the set, so return the index.
184
212600
            return (SetIndex(self.generation_counter.recall_index(existing)), false);
185
552220
        }
186

            
187
552220
        let index = self.insert_into_table(value);
188
552220
        let hasher_ref = &self.hasher;
189
552220
        let table_ref = &self.table;
190
552220
        self.index
191
853311
            .insert_unique(hash, index, |&i| entry_hash(table_ref, hasher_ref, i));
192
552220
        (SetIndex(self.generation_counter.create_index(index)), true)
193
764820
    }
194

            
195
    /// Returns the index for the given element, or None if it does not exist.
196
1887337
    pub fn index<Q>(&self, key: &Q) -> Option<SetIndex>
197
1887337
    where
198
1887337
        Q: Hash + Equivalent<T> + ?Sized,
199
    {
200
1887337
        let hash = self.hasher.hash_one(key);
201
1887337
        self.index
202
1905841
            .find(hash, |&i| entry_matches(&self.table, i, key))
203
1887337
            .map(|&i| SetIndex(self.generation_counter.recall_index(i)))
204
1887337
    }
205

            
206
    /// Erases all elements for which f(index, element) returns false. Allows
207
    /// modifying the given element (as long as the hash/equality does not change).
208
    pub fn retain_mut<F>(&mut self, mut f: F)
209
    where
210
        F: FnMut(SetIndex, &mut T) -> bool,
211
    {
212
        for (index, element) in self.table.iter_mut().enumerate() {
213
            if let IndexSetEntry::Filled(value) = element
214
                && !f(SetIndex(self.generation_counter.recall_index(index)), value)
215
            {
216
                let hash = self.hasher.hash_one(value);
217
                if let Ok(entry) = self.index.find_entry(hash, |&i| i == index) {
218
                    entry.remove();
219
                }
220

            
221
                match self.free {
222
                    Some(next) => {
223
                        *element = IndexSetEntry::Empty(next);
224
                    }
225
                    None => {
226
                        *element = IndexSetEntry::Empty(index);
227
                    }
228
                };
229
                self.free = Some(index);
230
            };
231
        }
232
    }
233

            
234
    /// Removes the given element from the set.
235
1000
    pub fn remove(&mut self, element: &T) -> bool {
236
1000
        let hash = self.hasher.hash_one(element);
237

            
238
1000
        if let Ok(entry) = self.index.find_entry(hash, |&i| entry_matches(&self.table, i, element)) {
239
873
            let (removed_index, _) = entry.remove();
240
873
            let next = match self.free {
241
773
                Some(next) => next,
242
100
                None => removed_index,
243
            };
244

            
245
873
            self.table[removed_index] = IndexSetEntry::Empty(next);
246
873
            self.free = Some(removed_index);
247
873
            true
248
        } else {
249
            // The element was not found in the set.
250
127
            false
251
        }
252
1000
    }
253

            
254
    /// Removes all elements in this indexed set.
255
602
    pub fn clear(&mut self) {
256
602
        self.table.clear();
257
602
        self.index.clear();
258
602
        self.free = None;
259
602
        self.generation_counter = GenerationCounter::new();
260
602
    }
261

            
262
    /// Returns true iff the set contains the given element.
263
1542894
    pub fn contains<Q>(&self, element: &Q) -> bool
264
1542894
    where
265
1542894
        Q: Hash + Equivalent<T>,
266
    {
267
1542894
        let hash = self.hasher.hash_one(element);
268
1542894
        self.index
269
1542894
            .find(hash, |&i| entry_matches(&self.table, i, element))
270
1542894
            .is_some()
271
1542894
    }
272

            
273
    /// Inserts `value` into the `table`, reusing a free slot if one is available, and returns
274
    /// its index. Does not modify the secondary hash index.
275
552220
    fn insert_into_table(&mut self, value: T) -> usize {
276
552220
        match self.free {
277
            Some(first) => {
278
                let next = match self.table[first] {
279
                    IndexSetEntry::Empty(x) => x,
280
                    IndexSetEntry::Filled(_) => panic!("The free list contains a filled element"),
281
                };
282

            
283
                if first == next {
284
                    // The list is now empty as its first element points to itself.
285
                    self.free = None;
286
                } else {
287
                    // Update free to be the next element in the list.
288
                    self.free = Some(next);
289
                }
290

            
291
                self.table[first] = IndexSetEntry::Filled(value);
292
                first
293
            }
294
            None => {
295
                // No free positions so insert new.
296
552220
                self.table.push(IndexSetEntry::Filled(value));
297
552220
                self.table.len() - 1
298
            }
299
        }
300
552220
    }
301
}
302

            
303
impl<T, S> fmt::Debug for IndexedSet<T, S>
304
where
305
    T: fmt::Debug,
306
{
307
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
308
        f.debug_list().entries(self.iter()).finish()
309
    }
310
}
311

            
312
impl<T, S: BuildHasher + Default> Default for IndexedSet<T, S> {
313
100
    fn default() -> IndexedSet<T, S> {
314
100
        IndexedSet::new()
315
100
    }
316
}
317

            
318
impl<T, S> Index<SetIndex> for IndexedSet<T, S> {
319
    type Output = T;
320

            
321
    fn index(&self, index: SetIndex) -> &Self::Output {
322
        // Go through the generation counter so a stale index is detected in
323
        // debug builds, consistent with [Self::get].
324
        let raw = self.generation_counter.get_index(index.0);
325
        cast!(&self.table[raw], IndexSetEntry::Filled)
326
    }
327
}
328

            
329
impl<T, S: BuildHasher> IndexMut<SetIndex> for IndexedSet<T, S> {
330
    fn index_mut(&mut self, index: SetIndex) -> &mut Self::Output {
331
        // Go through the generation counter so a stale index is detected in
332
        // debug builds, consistent with [Self::get].
333
        let raw = self.generation_counter.get_index(index.0);
334
        cast!(&mut self.table[raw], IndexSetEntry::Filled)
335
    }
336
}
337

            
338
/// Returns whether the element stored at `table[index]` is equivalent to `value`.
339
///
340
/// Returns `false` for free slots, which should never appear in the secondary hash index.
341
2633288
fn entry_matches<T, Q: Equivalent<T> + ?Sized>(table: &[IndexSetEntry<T>], index: usize, value: &Q) -> bool {
342
2633288
    match &table[index] {
343
2633288
        IndexSetEntry::Filled(element) => value.equivalent(element),
344
        IndexSetEntry::Empty(_) => false,
345
    }
346
2633288
}
347

            
348
/// Computes the hash of the filled element stored at `table[index]` using `hasher`.
349
///
350
/// Used as the rehasher callback for [`HashTable`] operations.
351
852505
fn entry_hash<T: Hash, S: BuildHasher>(table: &[IndexSetEntry<T>], hasher: &S, index: usize) -> u64 {
352
852505
    match &table[index] {
353
852505
        IndexSetEntry::Filled(element) => hasher.hash_one(element),
354
        IndexSetEntry::Empty(_) => panic!("entry_hash called on an empty slot"),
355
    }
356
852505
}
357

            
358
/// An iterator over the elements in the IndexedSet.
359
pub struct Iter<'a, T, S> {
360
    reference: &'a IndexedSet<T, S>,
361
    index: usize,
362
    generation_counter: &'a GenerationCounter,
363
}
364

            
365
impl<'a, T, S> Iterator for Iter<'a, T, S> {
366
    type Item = (SetIndex, &'a T);
367

            
368
320017
    fn next(&mut self) -> Option<Self::Item> {
369
320890
        while self.index < self.reference.table.len() {
370
297817
            let current_index = self.index;
371
297817
            self.index += 1;
372

            
373
297817
            if let IndexSetEntry::Filled(element) = &self.reference.table[current_index] {
374
296944
                return Some((SetIndex(self.generation_counter.recall_index(current_index)), element));
375
873
            }
376
        }
377

            
378
23073
        None
379
320017
    }
380
}
381

            
382
impl<'a, T, S> IntoIterator for &'a IndexedSet<T, S> {
383
    type Item = (SetIndex, &'a T);
384
    type IntoIter = Iter<'a, T, S>;
385

            
386
200
    fn into_iter(self) -> Self::IntoIter {
387
200
        self.iter()
388
200
    }
389
}
390

            
391
#[cfg(test)]
392
mod tests {
393
    use rand::RngExt;
394
    use std::collections::HashMap;
395

            
396
    use merc_utilities::random_test;
397

            
398
    use crate::IndexedSet;
399
    use crate::SetIndex;
400

            
401
    #[test]
402
1
    fn test_random_indexed_set_construction() {
403
100
        random_test(100, |rng| {
404
100
            let mut input = vec![];
405
10000
            for _ in 0..100 {
406
10000
                input.push(rng.random_range(0..32) as usize);
407
10000
            }
408

            
409
100
            let mut indices: HashMap<usize, SetIndex> = HashMap::default();
410

            
411
            // Insert several elements and keep track of the resulting indices.
412
100
            let mut set: IndexedSet<usize> = IndexedSet::default();
413
10000
            for element in &input {
414
10000
                let index = set.insert(*element).0;
415
10000
                indices.insert(*element, index);
416
10000
            }
417

            
418
            // Check if the indices match the previously stored ones.
419
3073
            for (index, value) in &set {
420
3073
                assert_eq!(
421
3073
                    indices[value], index,
422
                    "The resulting index does not match the returned value"
423
                );
424
            }
425

            
426
            // Remove some elements from the set.
427
1000
            for value in &mut input.iter().take(10) {
428
1000
                set.remove(value);
429
1000
                indices.remove(value);
430
1000
            }
431

            
432
            // Check consistency of the indexed set after removals.
433
2200
            for (index, value) in &set {
434
2200
                assert_eq!(
435
2200
                    indices[value], index,
436
                    "The resulting index does not match the returned value"
437
                );
438
            }
439

            
440
2200
            for (value, index) in &indices {
441
2200
                assert!(
442
2200
                    set.get(*index) == Some(value),
443
                    "Index {} should still match element {:?}",
444
                    *index,
445
                    value
446
                );
447
            }
448

            
449
            // Check the contains function
450
10000
            for value in &input {
451
10000
                let contains = indices.contains_key(value);
452
10000
                assert_eq!(
453
10000
                    set.contains(value),
454
                    contains,
455
                    "The contains function returned an incorrect result for value {:?}",
456
                    value
457
                );
458
            }
459
100
        })
460
1
    }
461
}