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

            
10
use hashbrown::Equivalent;
11
use hashbrown::HashSet;
12
use rustc_hash::FxBuildHasher;
13

            
14
use merc_utilities::GenerationCounter;
15
use merc_utilities::GenerationalIndex;
16
use merc_utilities::NoHasherBuilder;
17
use merc_utilities::cast;
18

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

            
25
impl Deref for SetIndex {
26
    type Target = usize;
27

            
28
2556201770
    fn deref(&self) -> &Self::Target {
29
2556201770
        &self.0
30
2556201770
    }
31
}
32

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

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

            
45
/// A set that assigns a unique index to every entry. The returned index can be used to access the inserted entry.
46
pub struct IndexedSet<T, S = FxBuildHasher> {
47
    /// The table of elements, which can be either filled or empty.
48
    table: Vec<IndexSetEntry<T>>,
49
    /// Indexes of the elements in the set, using NoHasher to directly use precomputed hashes.
50
    index: HashSet<IndexEntry, NoHasherBuilder>,
51
    /// A list of free nodes, where the value is the first free node.
52
    free: Option<usize>,
53
    /// The number of generations
54
    generation_counter: GenerationCounter,
55
    /// The hasher used to compute hashes for elements
56
    hasher: S,
57
}
58

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

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

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

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

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

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

            
112
    /// Returns a reference to the element at the given index, if it exists.
113
    ///
114
    /// Does not check if the index is valid, so it is the caller's responsibility to ensure this.
115
    pub fn get_by_index(&self, index: usize) -> Option<&T> {
116
        if let Some(entry) = self.table.get(index) {
117
            match entry {
118
                IndexSetEntry::Filled(element) => Some(element),
119
                IndexSetEntry::Empty(_) => None,
120
            }
121
        } else {
122
            None
123
        }
124
    }
125

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

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

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

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

            
159
        if let Some(entry) = self.index.get(&equivalent) {
160
            // The element is already in the set, so return the index.
161
            return (SetIndex(self.generation_counter.recall_index(entry.index)), false);
162
        }
163

            
164
        let value: T = value.into();
165
        let hash = self.hasher.hash_one(&value);
166

            
167
        debug_assert_eq!(hash, equivalent.hash(), "Hash values should be the same");
168

            
169
        let index = match self.free {
170
            Some(first) => {
171
                let next = match self.table[first] {
172
                    IndexSetEntry::Empty(x) => x,
173
                    IndexSetEntry::Filled(_) => panic!("The free list contains a filled element"),
174
                };
175

            
176
                if first == next {
177
                    // The list is now empty as its first element points to itself.
178
                    self.free = None;
179
                } else {
180
                    // Update free to be the next element in the list.
181
                    self.free = Some(next);
182
                }
183

            
184
                self.table[first] = IndexSetEntry::Filled(value);
185
                first
186
            }
187
            None => {
188
                // No free positions so insert new.
189
                self.table.push(IndexSetEntry::Filled(value));
190
                self.table.len() - 1
191
            }
192
        };
193

            
194
        self.index.insert(IndexEntry::new(index, hash));
195
        (SetIndex(self.generation_counter.create_index(index)), true)
196
    }
197

            
198
    /// Inserts the given element into the set
199
    ///
200
    /// Returns the corresponding index and a boolean indicating if the element was inserted.
201
13742280
    pub fn insert(&mut self, value: T) -> (SetIndex, bool) {
202
13742280
        let equivalent = IndexValueEquivalent::new(&value, &self.hasher, &self.table);
203

            
204
13742280
        if let Some(entry) = self.index.get(&equivalent) {
205
            // The element is already in the set, so return the index.
206
5626297
            return (SetIndex(self.generation_counter.recall_index(entry.index)), false);
207
8115983
        }
208

            
209
8115983
        let hash = equivalent.hash();
210

            
211
8115983
        let index = match self.free {
212
            Some(first) => {
213
                let next = match self.table[first] {
214
                    IndexSetEntry::Empty(x) => x,
215
                    IndexSetEntry::Filled(_) => panic!("The free list contains a filled element"),
216
                };
217

            
218
                if first == next {
219
                    // The list is now empty as its first element points to itself.
220
                    self.free = None;
221
                } else {
222
                    // Update free to be the next element in the list.
223
                    self.free = Some(next);
224
                }
225

            
226
                self.table[first] = IndexSetEntry::Filled(value);
227
                first
228
            }
229
            None => {
230
                // No free positions so insert new.
231
8115983
                self.table.push(IndexSetEntry::Filled(value));
232
8115983
                self.table.len() - 1
233
            }
234
        };
235

            
236
8115983
        self.index.insert(IndexEntry::new(index, hash));
237
8115983
        (SetIndex(self.generation_counter.create_index(index)), true)
238
13742280
    }
239

            
240
    /// Returns the index for the given element, or None if it does not exist.
241
2084295
    pub fn index<Q>(&self, key: &Q) -> Option<SetIndex>
242
2084295
    where
243
2084295
        Q: Hash + Equivalent<T>,
244
    {
245
2084295
        let equivalent = IndexValueEquivalent::new(key, &self.hasher, &self.table);
246

            
247
2084295
        self.index
248
2084295
            .get(&equivalent)
249
2084295
            .map(|entry| SetIndex(self.generation_counter.recall_index(entry.index)))
250
2084295
    }
251

            
252
    /// Erases all elements for which f(index, element) returns false. Allows
253
    /// modifying the given element (as long as the hash/equality does not change).
254
240
    pub fn retain_mut<F>(&mut self, mut f: F)
255
240
    where
256
240
        F: FnMut(SetIndex, &mut T) -> bool,
257
    {
258
434504
        for (index, element) in self.table.iter_mut().enumerate() {
259
434504
            if let IndexSetEntry::Filled(value) = element
260
223812
                && !f(SetIndex(self.generation_counter.recall_index(index)), value)
261
            {
262
                // Find and remove the IndexEntry from the index set
263
582737234
                let entry_to_remove = self.index.iter().find(|entry| entry.index == index).cloned();
264

            
265
214236
                if let Some(entry) = entry_to_remove {
266
214236
                    self.index.remove(&entry);
267
214236
                }
268

            
269
214236
                match self.free {
270
214116
                    Some(next) => {
271
214116
                        *element = IndexSetEntry::Empty(next);
272
214116
                    }
273
120
                    None => {
274
120
                        *element = IndexSetEntry::Empty(index);
275
120
                    }
276
                };
277
214236
                self.free = Some(index);
278
220268
            };
279
        }
280
240
    }
281

            
282
    /// Removes the given element from the set.
283
1000
    pub fn remove(&mut self, element: &T) -> bool {
284
1000
        let equivalent = IndexValueEquivalent::new(element, &self.hasher, &self.table);
285

            
286
1000
        if let Some(entry) = self.index.take(&equivalent) {
287
866
            let next = match self.free {
288
766
                Some(next) => next,
289
100
                None => entry.index,
290
            };
291

            
292
866
            self.table[entry.index] = IndexSetEntry::Empty(next);
293
866
            self.free = Some(entry.index);
294
866
            true
295
        } else {
296
            // The element was not found in the set.
297
134
            false
298
        }
299
1000
    }
300

            
301
    /// Removes all elements in this indexed set.
302
600
    pub fn clear(&mut self) {
303
600
        self.table.clear();
304
600
        self.index.clear();
305
600
        self.free = None;
306
600
        self.generation_counter = GenerationCounter::new();
307
600
    }
308

            
309
    /// Returns true iff the set contains the given element.
310
1931881
    pub fn contains<Q>(&self, element: &Q) -> bool
311
1931881
    where
312
1931881
        Q: Hash + Equivalent<T>,
313
    {
314
        // Compute the hash using our hash_builder
315
1931881
        let equivalent = IndexValueEquivalent::new(element, &self.hasher, &self.table);
316
1931881
        self.index.contains(&equivalent)
317
1931881
    }
318
}
319

            
320
impl<T, S> fmt::Debug for IndexedSet<T, S>
321
where
322
    T: fmt::Debug,
323
{
324
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
325
        f.debug_list().entries(self.iter()).finish()
326
    }
327
}
328

            
329
impl<T, S: BuildHasher + Default> Default for IndexedSet<T, S> {
330
100
    fn default() -> IndexedSet<T, S> {
331
100
        IndexedSet::new()
332
100
    }
333
}
334

            
335
impl<T, S> Index<SetIndex> for IndexedSet<T, S> {
336
    type Output = T;
337

            
338
213467406
    fn index(&self, index: SetIndex) -> &Self::Output {
339
213467406
        cast!(&self.table[*index], IndexSetEntry::Filled)
340
213467406
    }
341
}
342

            
343
impl<T, S: BuildHasher> IndexMut<SetIndex> for IndexedSet<T, S> {
344
19032
    fn index_mut(&mut self, index: SetIndex) -> &mut Self::Output {
345
19032
        cast!(&mut self.table[*index], IndexSetEntry::Filled)
346
19032
    }
347
}
348

            
349
/// An entry in the index that stores both the index and a precomputed hash value
350
#[derive(Copy, Clone, PartialEq, Eq)]
351
struct IndexEntry {
352
    /// The index into the table
353
    index: usize,
354
    /// Precomputed hash value of the element at this index
355
    hash: u64,
356
}
357

            
358
impl IndexEntry {
359
    /// Creates a new IndexEntry with the given index and hash value
360
112655299
    fn new(index: usize, hash: u64) -> Self {
361
112655299
        Self { index, hash }
362
112655299
    }
363
}
364

            
365
impl Hash for IndexEntry {
366
21871839
    fn hash<H: Hasher>(&self, state: &mut H) {
367
        // Simply use the precomputed hash value
368
21871839
        state.write_u64(self.hash);
369
21871839
    }
370
}
371

            
372
/// An equivalent wrapper that allows looking up elements in a set using the original value.
373
/// This avoids duplicating the key in both the table and index.
374
struct IndexValueEquivalent<'a, T, Q> {
375
    value: &'a Q,
376
    hash: u64,
377
    table: &'a Vec<IndexSetEntry<T>>,
378
}
379

            
380
impl<T, Q> IndexValueEquivalent<'_, T, Q> {
381
8115983
    fn hash(&self) -> u64 {
382
        // This is a placeholder for the actual hash function
383
8115983
        self.hash
384
8115983
    }
385
}
386

            
387
impl<'a, T, Q: Hash> IndexValueEquivalent<'a, T, Q> {
388
    /// Creates a new IndexValueEquivalent with the given value and table.
389
17759456
    fn new<S: BuildHasher>(value: &'a Q, hasher: &S, table: &'a Vec<IndexSetEntry<T>>) -> Self {
390
        // Constructor allows for centralized creation logic
391
17759456
        Self {
392
17759456
            value,
393
17759456
            table,
394
17759456
            hash: hasher.hash_one(value),
395
17759456
        }
396
17759456
    }
397
}
398

            
399
impl<T, Q: Equivalent<T>> Equivalent<IndexEntry> for IndexValueEquivalent<'_, T, Q> {
400
9205480
    fn equivalent(&self, key: &IndexEntry) -> bool {
401
9205480
        if let Some(IndexSetEntry::Filled(element)) = self.table.get(key.index) {
402
9205480
            self.value.equivalent(element)
403
        } else {
404
            false
405
        }
406
9205480
    }
407
}
408

            
409
impl<T, Q> Hash for IndexValueEquivalent<'_, T, Q> {
410
17749108
    fn hash<H: Hasher>(&self, state: &mut H) {
411
17749108
        state.write_u64(self.hash);
412
17749108
    }
413
}
414

            
415
/// An iterator over the elements in the IndexedSet.
416
pub struct Iter<'a, T, S> {
417
    reference: &'a IndexedSet<T, S>,
418
    index: usize,
419
    generation_counter: &'a GenerationCounter,
420
}
421

            
422
impl<'a, T, S> Iterator for Iter<'a, T, S> {
423
    type Item = (SetIndex, &'a T);
424

            
425
329797
    fn next(&mut self) -> Option<Self::Item> {
426
755591
        while self.index < self.reference.table.len() {
427
732280
            let current_index = self.index;
428
732280
            self.index += 1;
429

            
430
732280
            if let IndexSetEntry::Filled(element) = &self.reference.table[current_index] {
431
306486
                return Some((SetIndex(self.generation_counter.recall_index(current_index)), element));
432
425794
            }
433
        }
434

            
435
23311
        None
436
329797
    }
437
}
438

            
439
impl<'a, T, S> IntoIterator for &'a IndexedSet<T, S> {
440
    type Item = (SetIndex, &'a T);
441
    type IntoIter = Iter<'a, T, S>;
442

            
443
440
    fn into_iter(self) -> Self::IntoIter {
444
440
        self.iter()
445
440
    }
446
}
447

            
448
#[cfg(test)]
449
mod tests {
450
    use std::collections::HashMap;
451

            
452
    use rand::RngExt;
453

            
454
    use merc_utilities::random_test;
455

            
456
    use crate::IndexedSet;
457
    use crate::SetIndex;
458

            
459
    #[test]
460
1
    fn test_random_indexed_set_construction() {
461
100
        random_test(100, |rng| {
462
100
            let mut input = vec![];
463
10000
            for _ in 0..100 {
464
10000
                input.push(rng.random_range(0..32) as usize);
465
10000
            }
466

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

            
469
            // Insert several elements and keep track of the resulting indices.
470
100
            let mut set: IndexedSet<usize> = IndexedSet::default();
471
10000
            for element in &input {
472
10000
                let index = set.insert(*element).0;
473
10000
                indices.insert(*element, index);
474
10000
            }
475

            
476
            // Check if the indices match the previously stored ones.
477
3065
            for (index, value) in &set {
478
3065
                assert_eq!(
479
3065
                    indices[value], index,
480
                    "The resulting index does not match the returned value"
481
                );
482
            }
483

            
484
            // Remove some elements from the set.
485
1000
            for value in &mut input.iter().take(10) {
486
1000
                set.remove(value);
487
1000
                indices.remove(value);
488
1000
            }
489

            
490
            // Check consistency of the indexed set after removals.
491
2199
            for (index, value) in &set {
492
2199
                assert_eq!(
493
2199
                    indices[value], index,
494
                    "The resulting index does not match the returned value"
495
                );
496
            }
497

            
498
2199
            for (value, index) in &indices {
499
2199
                assert!(
500
2199
                    set.get(*index) == Some(value),
501
                    "Index {} should still match element {:?}",
502
                    *index,
503
                    value
504
                );
505
            }
506

            
507
            // Check the contains function
508
10000
            for value in &input {
509
10000
                let contains = indices.contains_key(value);
510
10000
                assert_eq!(
511
10000
                    set.contains(value),
512
                    contains,
513
                    "The contains function returned an incorrect result for value {:?}",
514
                    value
515
                );
516
            }
517
100
        })
518
1
    }
519
}