1
use std::hash::BuildHasher;
2
use std::hash::Hash;
3

            
4
use rustc_hash::FxBuildHasher;
5

            
6
use crate::ConcurrentAppendVec;
7
use crate::ShardedHashMap;
8

            
9
/// An append-only set mapping each distinct value of type `T` to a stable
10
/// `usize` index.
11
///
12
/// The set is thread-safe: values can be inserted and looked up concurrently
13
/// through `&self` from multiple threads. To keep concurrent insertion
14
/// contention-free the indices are unique and stable but **not dense**, so walk
15
/// [`ConcurrentIndexedSet::iter`] rather than `0..len`. There is no removal; an
16
/// index stays valid for the lifetime of the set (until
17
/// [`ConcurrentIndexedSet::clear`], which requires `&mut self`).
18
pub struct ConcurrentIndexedSet<T, S = FxBuildHasher> {
19
    /// Stores each distinct value at the index handed out for it. Indices are
20
    /// stable but may be sparse.
21
    values: ConcurrentAppendVec<T>,
22
    /// Hash index from a value's hash to its index in `values`. Stores only the
23
    /// index; the hash and equality closures dereference `values`, so the
24
    /// payload is not duplicated.
25
    table: ShardedHashMap<usize, S>,
26
}
27

            
28
impl<T, S: Default> ConcurrentIndexedSet<T, S> {
29
    /// Creates a new empty set with the default hasher.
30
246
    pub fn new() -> ConcurrentIndexedSet<T, S> {
31
246
        ConcurrentIndexedSet {
32
246
            values: ConcurrentAppendVec::new(),
33
246
            table: ShardedHashMap::with_hasher(S::default()),
34
246
        }
35
246
    }
36

            
37
    /// Creates a new empty set whose hash index has room for at least
38
    /// `capacity` values before reallocating.
39
    pub fn with_capacity(capacity: usize) -> ConcurrentIndexedSet<T, S> {
40
        ConcurrentIndexedSet {
41
            values: ConcurrentAppendVec::new(),
42
            table: ShardedHashMap::with_capacity_and_hasher(capacity, S::default()),
43
        }
44
    }
45
}
46

            
47
impl<T, S> ConcurrentIndexedSet<T, S>
48
where
49
    T: Hash + Eq + Clone,
50
    S: BuildHasher,
51
{
52
    /// Inserts `value` and returns its index together with a boolean that is
53
    /// true when the value was newly inserted and false when it was already
54
    /// present.
55
66016
    pub fn insert(&self, value: T) -> (usize, bool) {
56
66016
        let hash = self.table.hash(&value);
57
        // SAFETY (this and the hash closure below): the table only ever holds
58
        // indices returned by `push`, and every read here runs under the shard
59
        // lock that ordered that push's publish, so the slot is initialized and
60
        // visible — `get_unchecked` is sound.
61
66016
        let eq = |&index: &usize| unsafe { self.values.get_unchecked(index) } == &value;
62

            
63
        // Fast path: the value is already present, so no write lock is needed.
64
66016
        if let Some(index) = self.table.find(hash, eq) {
65
42677
            return (index, false);
66
23339
        }
67

            
68
        // The vacant branch runs while the shard for `hash` is write-locked, so
69
        // two threads inserting the same new value cannot both push and hand
70
        // out duplicate indices. `push` publishes the slot before returning, so
71
        // recording its index in the table afterwards guarantees a returned
72
        // index always resolves in `get_by_index`.
73
23339
        self.table.find_or_insert_with(
74
23339
            hash,
75
23339
            eq,
76
21557
            |&index| self.table.hash(unsafe { self.values.get_unchecked(index) }),
77
23339
            || self.values.push(value.clone()),
78
        )
79
66016
    }
80

            
81
    /// Returns the index of `value` if it is present, or `None` otherwise.
82
50356
    pub fn index(&self, value: &T) -> Option<usize> {
83
50356
        let hash = self.table.hash(value);
84
        // SAFETY: indices in the table were returned by `push` and published
85
        // under the shard lock this read takes, so `get_unchecked` is sound.
86
50356
        self.table
87
50365
            .find(hash, |&index| unsafe { self.values.get_unchecked(index) } == value)
88
50356
    }
89

            
90
    /// Returns true if `value` is present in the set.
91
25050
    pub fn contains(&self, value: &T) -> bool {
92
25050
        self.index(value).is_some()
93
25050
    }
94
}
95

            
96
impl<T, S> ConcurrentIndexedSet<T, S> {
97
    /// Returns a reference to the value at `index`, or `None` if the index was
98
    /// never handed out (out of range or a sparse gap).
99
64283
    pub fn get_by_index(&self, index: usize) -> Option<&T> {
100
64283
        self.values.get(index)
101
64283
    }
102

            
103
    /// Returns an iterator over the resident values. Order is unspecified, and
104
    /// indices are not exposed because they may be sparse.
105
1
    pub fn iter(&self) -> impl Iterator<Item = &T> {
106
1
        self.values.iter()
107
1
    }
108

            
109
    /// Returns the number of distinct values in the set.
110
20052
    pub fn len(&self) -> usize {
111
20052
        self.values.len()
112
20052
    }
113

            
114
    /// Returns true if the set is empty.
115
    pub fn is_empty(&self) -> bool {
116
        self.values.is_empty()
117
    }
118

            
119
    /// Removes all values, invalidating every previously returned index.
120
    pub fn clear(&mut self) {
121
        self.values.clear();
122
        self.table.clear();
123
    }
124
}
125

            
126
impl<T, S: Default> Default for ConcurrentIndexedSet<T, S> {
127
    fn default() -> ConcurrentIndexedSet<T, S> {
128
        ConcurrentIndexedSet::new()
129
    }
130
}
131

            
132
#[cfg(test)]
133
mod tests {
134
    use std::collections::HashMap;
135
    use std::collections::HashSet;
136
    use std::sync::Arc;
137
    use std::sync::atomic::AtomicUsize;
138
    use std::sync::atomic::Ordering;
139

            
140
    use rand::RngExt;
141

            
142
    use merc_utilities::random_test;
143
    use merc_utilities::random_test_threads;
144

            
145
    use super::ConcurrentIndexedSet;
146

            
147
    /// Inserts random values from a small space and checks the set against a
148
    /// `HashMap<value, index>` oracle after every step: the fresh-insert flag,
149
    /// the index assigned in insertion order (dense on a single thread), and
150
    /// that every read path (`index`, `contains`, `get_by_index`) round-trips
151
    /// against the model.
152
    #[test]
153
    #[cfg_attr(miri, ignore)]
154
1
    fn random_insert_dedup_and_roundtrip() {
155
50
        random_test(50, |rng| {
156
50
            let set: ConcurrentIndexedSet<u64> = ConcurrentIndexedSet::new();
157
50
            let mut model: HashMap<u64, usize> = HashMap::new();
158

            
159
50
            for _ in 0..500 {
160
25000
                let value = rng.random_range(0..32u64);
161
25000
                let next = model.len();
162
25000
                let expected_new = !model.contains_key(&value);
163

            
164
25000
                let (index, is_new) = set.insert(value);
165
25000
                assert_eq!(is_new, expected_new, "fresh-insert flag agrees with the model");
166
25000
                assert_eq!(
167
                    index,
168
25000
                    *model.entry(value).or_insert(next),
169
                    "equal values share the index assigned in insertion order"
170
                );
171

            
172
25000
                assert_eq!(set.get_by_index(index), Some(&value), "the index resolves to the value");
173
25000
                assert_eq!(set.index(&value), Some(index), "lookup is the inverse of get_by_index");
174
25000
                assert!(set.contains(&value));
175
            }
176

            
177
50
            assert_eq!(set.len(), model.len());
178

            
179
            // Values from outside the inserted space are absent.
180
50
            assert_eq!(set.index(&64), None);
181
50
            assert!(!set.contains(&64));
182
50
            assert_eq!(set.get_by_index(model.len()), None);
183
50
        });
184
1
    }
185

            
186
    /// Concurrently inserts random values from a fixed value space across
187
    /// several threads. Across all threads each value yields exactly one fresh
188
    /// insertion, so the count of fresh insertions must equal the set size, and
189
    /// every returned index resolves to the value that was inserted even under
190
    /// contention.
191
    #[test]
192
    #[cfg_attr(miri, ignore)]
193
1
    fn random_concurrent_insert_counts_each_once() {
194
1
        let set: Arc<ConcurrentIndexedSet<u64>> = Arc::new(ConcurrentIndexedSet::new());
195
1
        let inserted = Arc::new(AtomicUsize::new(0));
196
1
        let values = 256u64;
197

            
198
1
        random_test_threads(
199
            2000,
200
            8,
201
8
            || (Arc::clone(&set), Arc::clone(&inserted)),
202
16000
            move |rng, (set, inserted)| {
203
16000
                let value = rng.random_range(0..values);
204
16000
                let (index, is_new) = set.insert(value);
205
16000
                assert_eq!(
206
16000
                    set.get_by_index(index),
207
16000
                    Some(&value),
208
                    "the index resolves under contention"
209
                );
210
16000
                if is_new {
211
256
                    inserted.fetch_add(1, Ordering::Relaxed);
212
15744
                }
213
16000
            },
214
        );
215

            
216
1
        assert_eq!(
217
1
            inserted.load(Ordering::Relaxed),
218
1
            set.len(),
219
            "every fresh insertion is a distinct resident value"
220
        );
221

            
222
        // Indices may be sparse, so iterate the resident values rather than
223
        // `0..len`. Each must be in range and round-trip through `index` and
224
        // `get_by_index`.
225
1
        let mut residents = HashSet::new();
226
256
        for &value in set.iter() {
227
256
            assert!(value < values, "only inserted values are resident");
228
256
            assert!(residents.insert(value), "each value is resident once");
229
256
            let index = set.index(&value).expect("a resident value has an index");
230
256
            assert_eq!(set.get_by_index(index), Some(&value), "index is the inverse of lookup");
231
        }
232
1
        assert_eq!(residents.len(), set.len(), "iter visits every resident value once");
233
1
    }
234
}