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
1301453032
    fn deref(&self) -> &Self::Target {
29
1301453032
        &self.0
30
1301453032
    }
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
14049
    pub fn new() -> IndexedSet<T, S> {
68
14049
        IndexedSet {
69
14049
            table: Vec::default(),
70
14049
            index: HashSet::with_hasher(NoHasherBuilder),
71
14049
            free: None,
72
14049
            generation_counter: GenerationCounter::new(),
73
14049
            hasher: S::default(),
74
14049
        }
75
14049
    }
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
27665806
    pub fn len(&self) -> usize {
92
27665806
        self.table.len()
93
27665806
    }
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
128929278
    pub fn get(&self, index: SetIndex) -> Option<&T> {
102
128929278
        if let Some(entry) = self.table.get(self.generation_counter.get_index(index.0)) {
103
128929278
            match entry {
104
128929278
                IndexSetEntry::Filled(element) => Some(element),
105
                IndexSetEntry::Empty(_) => None,
106
            }
107
        } else {
108
            None
109
        }
110
128929278
    }
111

            
112
    /// Returns the capacity of the set.
113
    pub fn capacity(&self) -> usize {
114
        self.table.capacity()
115
    }
116

            
117
    /// Returns an iterator over the elements in the set.
118
5340
    pub fn iter(&self) -> Iter<'_, T, S> {
119
5340
        Iter {
120
5340
            reference: self,
121
5340
            index: 0,
122
5340
            generation_counter: &self.generation_counter,
123
5340
        }
124
5340
    }
125
}
126

            
127
impl<T: Clone, S> IndexedSet<T, S> {
128
    /// Returns a vector containing all elements of this indexed set.
129
4900
    pub fn to_vec(&self) -> Vec<T> {
130
33100
        self.iter().map(|(_, entry)| entry.clone()).collect()
131
4900
    }
132
}
133

            
134
impl<T: Hash + Eq, S: BuildHasher> IndexedSet<T, S> {
135
    /// Inserts the given element into the set
136
    ///
137
    /// Returns the corresponding index and a boolean indicating if the element was inserted.
138
    pub fn insert_equiv<'a, Q>(&mut self, value: &'a Q) -> (SetIndex, bool)
139
    where
140
        Q: Hash + Equivalent<T>,
141
        T: From<&'a Q>,
142
    {
143
        let equivalent = IndexValueEquivalent::new(value, &self.hasher, &self.table);
144

            
145
        if let Some(entry) = self.index.get(&equivalent) {
146
            // The element is already in the set, so return the index.
147
            return (SetIndex(self.generation_counter.recall_index(entry.index)), false);
148
        }
149

            
150
        let value: T = value.into();
151
        let hash = self.hasher.hash_one(&value);
152

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

            
155
        let index = match self.free {
156
            Some(first) => {
157
                let next = match self.table[first] {
158
                    IndexSetEntry::Empty(x) => x,
159
                    IndexSetEntry::Filled(_) => panic!("The free list contains a filled element"),
160
                };
161

            
162
                if first == next {
163
                    // The list is now empty as its first element points to itself.
164
                    self.free = None;
165
                } else {
166
                    // Update free to be the next element in the list.
167
                    self.free = Some(next);
168
                }
169

            
170
                self.table[first] = IndexSetEntry::Filled(value);
171
                first
172
            }
173
            None => {
174
                // No free positions so insert new.
175
                self.table.push(IndexSetEntry::Filled(value));
176
                self.table.len() - 1
177
            }
178
        };
179

            
180
        self.index.insert(IndexEntry::new(index, hash));
181
        (SetIndex(self.generation_counter.create_index(index)), true)
182
    }
183

            
184
    /// Inserts the given element into the set
185
    ///
186
    /// Returns the corresponding index and a boolean indicating if the element was inserted.
187
14029947
    pub fn insert(&mut self, value: T) -> (SetIndex, bool) {
188
14029947
        let equivalent = IndexValueEquivalent::new(&value, &self.hasher, &self.table);
189

            
190
14029947
        if let Some(entry) = self.index.get(&equivalent) {
191
            // The element is already in the set, so return the index.
192
5287683
            return (SetIndex(self.generation_counter.recall_index(entry.index)), false);
193
8742264
        }
194

            
195
8742264
        let hash = equivalent.hash();
196

            
197
8742264
        let index = match self.free {
198
            Some(first) => {
199
                let next = match self.table[first] {
200
                    IndexSetEntry::Empty(x) => x,
201
                    IndexSetEntry::Filled(_) => panic!("The free list contains a filled element"),
202
                };
203

            
204
                if first == next {
205
                    // The list is now empty as its first element points to itself.
206
                    self.free = None;
207
                } else {
208
                    // Update free to be the next element in the list.
209
                    self.free = Some(next);
210
                }
211

            
212
                self.table[first] = IndexSetEntry::Filled(value);
213
                first
214
            }
215
            None => {
216
                // No free positions so insert new.
217
8742264
                self.table.push(IndexSetEntry::Filled(value));
218
8742264
                self.table.len() - 1
219
            }
220
        };
221

            
222
8742264
        self.index.insert(IndexEntry::new(index, hash));
223
8742264
        (SetIndex(self.generation_counter.create_index(index)), true)
224
14029947
    }
225

            
226
    /// Returns the index for the given element, or None if it does not exist.
227
1092183
    pub fn index<Q>(&self, key: &Q) -> Option<SetIndex>
228
1092183
    where
229
1092183
        Q: Hash + Equivalent<T>,
230
    {
231
1092183
        let equivalent = IndexValueEquivalent::new(key, &self.hasher, &self.table);
232

            
233
1092183
        self.index
234
1092183
            .get(&equivalent)
235
1092183
            .map(|entry| SetIndex(self.generation_counter.recall_index(entry.index)))
236
1092183
    }
237

            
238
    /// Erases all elements for which f(index, element) returns false. Allows
239
    /// modifying the given element (as long as the hash/equality does not change).
240
240
    pub fn retain_mut<F>(&mut self, mut f: F)
241
240
    where
242
240
        F: FnMut(SetIndex, &mut T) -> bool,
243
    {
244
432770
        for (index, element) in self.table.iter_mut().enumerate() {
245
432770
            if let IndexSetEntry::Filled(value) = element
246
223013
                && !f(SetIndex(self.generation_counter.recall_index(index)), value) {
247
                    // Find and remove the IndexEntry from the index set
248
578242958
                    let entry_to_remove = self.index.iter().find(|entry| entry.index == index).cloned();
249

            
250
213376
                    if let Some(entry) = entry_to_remove {
251
213376
                        self.index.remove(&entry);
252
213376
                    }
253

            
254
213376
                    match self.free {
255
213256
                        Some(next) => {
256
213256
                            *element = IndexSetEntry::Empty(next);
257
213256
                        }
258
120
                        None => {
259
120
                            *element = IndexSetEntry::Empty(index);
260
120
                        }
261
                    };
262
213376
                    self.free = Some(index);
263
219394
                };
264
        }
265
240
    }
266

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

            
271
1000
        if let Some(entry) = self.index.take(&equivalent) {
272
867
            let next = match self.free {
273
767
                Some(next) => next,
274
100
                None => entry.index,
275
            };
276

            
277
867
            self.table[entry.index] = IndexSetEntry::Empty(next);
278
867
            self.free = Some(entry.index);
279
867
            true
280
        } else {
281
            // The element was not found in the set.
282
133
            false
283
        }
284
1000
    }
285

            
286
    /// Returns true iff the set contains the given element.
287
1858638
    pub fn contains<Q>(&self, element: &Q) -> bool
288
1858638
    where
289
1858638
        Q: Hash + Equivalent<T>,
290
    {
291
        // Compute the hash using our hash_builder
292
1858638
        let equivalent = IndexValueEquivalent::new(element, &self.hasher, &self.table);
293
1858638
        self.index.contains(&equivalent)
294
1858638
    }
295
}
296

            
297
impl<T, S> fmt::Debug for IndexedSet<T, S>
298
where
299
    T: fmt::Debug,
300
{
301
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
302
        f.debug_list().entries(self.iter()).finish()
303
    }
304
}
305

            
306
impl<T, S: BuildHasher + Default> Default for IndexedSet<T, S> {
307
100
    fn default() -> IndexedSet<T, S> {
308
100
        IndexedSet::new()
309
100
    }
310
}
311

            
312
impl<T, S> Index<SetIndex> for IndexedSet<T, S> {
313
    type Output = T;
314

            
315
128913505
    fn index(&self, index: SetIndex) -> &Self::Output {
316
128913505
        cast!(&self.table[*index], IndexSetEntry::Filled)
317
128913505
    }
318
}
319

            
320
impl<T, S: BuildHasher> IndexMut<SetIndex> for IndexedSet<T, S> {
321
19154
    fn index_mut(&mut self, index: SetIndex) -> &mut Self::Output {
322
19154
        cast!(&mut self.table[*index], IndexSetEntry::Filled)
323
19154
    }
324
}
325

            
326
/// An entry in the index that stores both the index and a precomputed hash value
327
#[derive(Copy, Clone, PartialEq, Eq)]
328
struct IndexEntry {
329
    /// The index into the table
330
    index: usize,
331
    /// Precomputed hash value of the element at this index
332
    hash: u64,
333
}
334

            
335
impl IndexEntry {
336
    /// Creates a new IndexEntry with the given index and hash value
337
74688958
    fn new(index: usize, hash: u64) -> Self {
338
74688958
        Self { index, hash }
339
74688958
    }
340
}
341

            
342
impl Hash for IndexEntry {
343
23695882
    fn hash<H: Hasher>(&self, state: &mut H) {
344
        // Simply use the precomputed hash value
345
23695882
        state.write_u64(self.hash);
346
23695882
    }
347
}
348

            
349
/// An equivalent wrapper that allows looking up elements in a set using the original value.
350
/// This avoids duplicating the key in both the table and index.
351
struct IndexValueEquivalent<'a, T, Q> {
352
    value: &'a Q,
353
    hash: u64,
354
    table: &'a Vec<IndexSetEntry<T>>,
355
}
356

            
357
impl<T, Q> IndexValueEquivalent<'_, T, Q> {
358
8742264
    fn hash(&self) -> u64 {
359
        // This is a placeholder for the actual hash function
360
8742264
        self.hash
361
8742264
    }
362
}
363

            
364
impl<'a, T, Q: Hash> IndexValueEquivalent<'a, T, Q> {
365
    /// Creates a new IndexValueEquivalent with the given value and table.
366
16981768
    fn new<S: BuildHasher>(value: &'a Q, hasher: &S, table: &'a Vec<IndexSetEntry<T>>) -> Self {
367
        // Constructor allows for centralized creation logic
368
16981768
        Self {
369
16981768
            value,
370
16981768
            table,
371
16981768
            hash: hasher.hash_one(value),
372
16981768
        }
373
16981768
    }
374
}
375

            
376
impl<T, Q: Equivalent<T>> Equivalent<IndexEntry> for IndexValueEquivalent<'_, T, Q> {
377
7932935
    fn equivalent(&self, key: &IndexEntry) -> bool {
378
7932935
        if let Some(IndexSetEntry::Filled(element)) = self.table.get(key.index) {
379
7932935
            self.value.equivalent(element)
380
        } else {
381
            false
382
        }
383
7932935
    }
384
}
385

            
386
impl<T, Q> Hash for IndexValueEquivalent<'_, T, Q> {
387
16963016
    fn hash<H: Hasher>(&self, state: &mut H) {
388
16963016
        state.write_u64(self.hash);
389
16963016
    }
390
}
391

            
392
/// An iterator over the elements in the IndexedSet.
393
pub struct Iter<'a, T, S> {
394
    reference: &'a IndexedSet<T, S>,
395
    index: usize,
396
    generation_counter: &'a GenerationCounter,
397
}
398

            
399
impl<'a, T, S> Iterator for Iter<'a, T, S> {
400
    type Item = (SetIndex, &'a T);
401

            
402
53376
    fn next(&mut self) -> Option<Self::Item> {
403
477376
        while self.index < self.reference.table.len() {
404
472036
            let current_index = self.index;
405
472036
            self.index += 1;
406

            
407
472036
            if let IndexSetEntry::Filled(element) = &self.reference.table[current_index] {
408
48036
                return Some((SetIndex(self.generation_counter.recall_index(current_index)), element));
409
424000
            }
410
        }
411

            
412
5340
        None
413
53376
    }
414
}
415

            
416
impl<'a, T, S> IntoIterator for &'a IndexedSet<T, S> {
417
    type Item = (SetIndex, &'a T);
418
    type IntoIter = Iter<'a, T, S>;
419

            
420
440
    fn into_iter(self) -> Self::IntoIter {
421
440
        self.iter()
422
440
    }
423
}
424

            
425
#[cfg(test)]
426
mod tests {
427
    use std::collections::HashMap;
428

            
429
    use rand::RngExt;
430

            
431
    use merc_utilities::random_test;
432

            
433
    use crate::IndexedSet;
434
    use crate::SetIndex;
435

            
436
    #[test]
437
1
    fn test_random_indexed_set_construction() {
438
100
        random_test(100, |rng| {
439
100
            let mut input = vec![];
440
10000
            for _ in 0..100 {
441
10000
                input.push(rng.random_range(0..32) as usize);
442
10000
            }
443

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

            
446
            // Insert several elements and keep track of the resulting indices.
447
100
            let mut set: IndexedSet<usize> = IndexedSet::default();
448
10000
            for element in &input {
449
10000
                let index = set.insert(*element).0;
450
10000
                indices.insert(*element, index);
451
10000
            }
452

            
453
            // Check if the indices match the previously stored ones.
454
3083
            for (index, value) in &set {
455
3083
                assert_eq!(
456
3083
                    indices[value], index,
457
                    "The resulting index does not match the returned value"
458
                );
459
            }
460

            
461
            // Remove some elements from the set.
462
1000
            for value in &mut input.iter().take(10) {
463
1000
                set.remove(value);
464
1000
                indices.remove(value);
465
1000
            }
466

            
467
            // Check consistency of the indexed set after removals.
468
2216
            for (index, value) in &set {
469
2216
                assert_eq!(
470
2216
                    indices[value], index,
471
                    "The resulting index does not match the returned value"
472
                );
473
            }
474

            
475
2216
            for (value, index) in &indices {
476
2216
                assert!(
477
2216
                    set.get(*index) == Some(value),
478
                    "Index {} should still match element {:?}",
479
                    *index,
480
                    value
481
                );
482
            }
483

            
484
            // Check the contains function
485
10000
            for value in &input {
486
10000
                let contains = indices.contains_key(value);
487
10000
                assert_eq!(
488
10000
                    set.contains(value),
489
                    contains,
490
                    "The contains function returned an incorrect result for value {:?}",
491
                    value
492
                );
493
            }
494
100
        })
495
1
    }
496
}