1
use std::fmt;
2
use std::hash::Hash;
3
use std::mem::ManuallyDrop;
4
use std::ops::Deref;
5
use std::ops::Index;
6

            
7
use merc_utilities::GenerationCounter;
8
use merc_utilities::GenerationalIndex;
9

            
10
/// A type-safe index for the ProtectionSet to prevent accidental use of wrong indices
11
#[repr(transparent)]
12
#[derive(Copy, Clone, Default, PartialEq, Eq, Hash)]
13
pub struct ProtectionIndex(GenerationalIndex<usize>);
14

            
15
impl Deref for ProtectionIndex {
16
    type Target = usize;
17

            
18
250003
    fn deref(&self) -> &Self::Target {
19
250003
        &self.0
20
250003
    }
21
}
22

            
23
impl fmt::Debug for ProtectionIndex {
24
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25
        write!(f, "ProtectionIndex({:?})", self.0)
26
    }
27
}
28

            
29
impl fmt::Display for ProtectionIndex {
30
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
31
        write!(f, "{}", self.0)
32
    }
33
}
34

            
35
/// A collection that assigns a unique index to every object added to it, and allows
36
/// removing objects while reusing their indices later. This is useful for managing
37
/// objects that must not be garbage collected, and as such it is called a protection set.
38
/// It is similar to a [`merc_collections::IndexedSet`], except that we cannot look up elements by value.
39
#[derive(Default)]
40
pub struct ProtectionSet<T> {
41
    roots: Vec<Entry<T>>, // The set of root active nodes.
42
    free: Option<usize>,
43
    number_of_insertions: u64,
44
    size: usize,
45

            
46
    /// The number of generations
47
    generation_counter: GenerationCounter,
48
}
49

            
50
impl<T> ProtectionSet<T> {
51
    /// Creates a new empty protection set.
52
6987
    pub fn new() -> Self {
53
6987
        ProtectionSet {
54
6987
            roots: Vec::new(),
55
6987
            free: None,
56
6987
            number_of_insertions: 0,
57
6987
            size: 0,
58
6987
            generation_counter: GenerationCounter::new(),
59
6987
        }
60
6987
    }
61

            
62
    /// Returns the number of insertions into the protection set.
63
100
    pub fn number_of_insertions(&self) -> u64 {
64
100
        self.number_of_insertions
65
100
    }
66

            
67
    /// Returns maximum number of active instances.
68
100
    pub fn maximum_size(&self) -> usize {
69
100
        self.roots.capacity()
70
100
    }
71

            
72
    /// Returns the number of roots in the protection set
73
100
    pub fn len(&self) -> usize {
74
100
        self.size
75
100
    }
76

            
77
    /// Returns whether the protection set is empty.
78
100
    pub fn is_empty(&self) -> bool {
79
100
        self.len() == 0
80
100
    }
81

            
82
    /// Adds the given object to the protection set and returns its index.
83
782556281
    pub fn protect(&mut self, object: T) -> ProtectionIndex {
84
782556281
        self.number_of_insertions += 1;
85
782556281
        self.size += 1;
86

            
87
782556281
        let index = match self.free {
88
768124434
            Some(first) => {
89
768124434
                let next = unsafe { self.roots[first].next };
90
768124434
                if first == next {
91
25104959
                    // The list is empty as its first element points to itself.
92
25104959
                    self.free = None;
93
743019475
                } else {
94
743019475
                    // Update free to be the next element in the list.
95
743019475
                    self.free = Some(next);
96
743019475
                }
97

            
98
768124434
                self.roots[first] = Entry {
99
768124434
                    object: ManuallyDrop::new(object),
100
768124434
                };
101
768124434
                first
102
            }
103
            None => {
104
                // If free list is empty insert new entry into roots.
105
14431847
                self.roots.push(Entry {
106
14431847
                    object: ManuallyDrop::new(object),
107
14431847
                });
108

            
109
14431847
                self.roots.len() - 1
110
            }
111
        };
112

            
113
782556281
        ProtectionIndex(self.generation_counter.create_index(index))
114
782556281
    }
115

            
116
    /// Remove protection from the given object. Note that index must be the
117
    /// index returned by the [ProtectionSet::protect] call.
118
782190727
    pub fn unprotect(&mut self, index: ProtectionIndex) {
119
782190727
        let index = self.generation_counter.get_index(index.0);
120
782190727
        self.size -= 1;
121

            
122
        // SAFETY: `index` refers to an occupied entry. We must drop the stored
123
        // object before overwriting the union slot with freelist metadata.
124
782190727
        unsafe {
125
782190727
            ManuallyDrop::drop(&mut self.roots[index].object);
126
782190727
        }
127

            
128
782190727
        match self.free {
129
757080017
            Some(next) => {
130
757080017
                self.roots[index] = Entry { next };
131
757080017
            }
132
25110710
            None => {
133
25110710
                self.roots[index] = Entry { next: index };
134
25110710
            }
135
        };
136

            
137
782190727
        self.free = Some(index);
138

            
139
        // Postcondition: verify the object was correctly removed from protection
140
782190727
        debug_assert!(
141
782190727
            self.freelist_iter().any(|free_idx| free_idx == index),
142
            "Failed to unprotect object"
143
        );
144
782190727
    }
145

            
146
    /// Replaces the object at the given index with the new object.
147
    pub fn replace(&mut self, index: ProtectionIndex, object: T) {
148
        let index = self.generation_counter.get_index(index.0);
149

            
150
        debug_assert!(
151
            !self.freelist_iter().any(|free_idx| free_idx == index),
152
            "Index {index} does not point to a filled entry"
153
        );
154

            
155
        self.roots[index] = Entry {
156
            object: ManuallyDrop::new(object),
157
        };
158
    }
159

            
160
    /// Returns an iterator over all root indices in the protection set.
161
    ///
162
    /// This is an O(n) operation in the size of the freelist, so it should be
163
    /// used with care. It is intended for debugging and testing purposes.
164
1553968
    pub fn iter(&self) -> ProtSetIter<'_, T> {
165
1553968
        let mut free_indices: Vec<usize> = self.freelist_iter().collect();
166
1553968
        free_indices.sort_unstable();
167
1553968
        ProtSetIter {
168
1553968
            current: 0,
169
1553968
            protection_set: self,
170
1553968
            generation_counter: &self.generation_counter,
171
1553968
            free_indices,
172
1553968
        }
173
1553968
    }
174

            
175
    /// Returns an iterator over all free indices in the freelist.
176
    ///
177
    /// This is intended for use in `debug_assert!`s to verify that an entry is
178
    /// filled (not in the freelist) or free (in the freelist).
179
798776551
    fn freelist_iter(&self) -> FreeListIter<'_, T> {
180
798776551
        FreeListIter {
181
798776551
            current: self.free,
182
798776551
            protection_set: self,
183
798776551
        }
184
798776551
    }
185

            
186
    /// Returns whether the protection set contains the given index.
187
    ///
188
    /// This operation is O(n) in the size of the freelist.
189
350010
    pub fn contains_root(&self, index: ProtectionIndex) -> bool {
190
350010
        let idx = self.generation_counter.get_index(index.0);
191
525000008
        !self.freelist_iter().any(|free_idx| free_idx == idx)
192
350010
    }
193
}
194

            
195
impl<T> Drop for ProtectionSet<T> {
196
6025
    fn drop(&mut self) {
197
14431843
        for index in 0..self.roots.len() {
198
766838044177
            if self.freelist_iter().any(|free_idx| free_idx == index) {
199
14066289
                continue;
200
365554
            }
201

            
202
            // SAFETY: indices not present in the freelist still contain a live
203
            // protected object that must be dropped.
204
365554
            unsafe {
205
365554
                ManuallyDrop::drop(&mut self.roots[index].object);
206
365554
            }
207
        }
208
6025
    }
209
}
210

            
211
impl<T> Index<ProtectionIndex> for ProtectionSet<T> {
212
    type Output = T;
213

            
214
250003
    fn index(&self, index: ProtectionIndex) -> &Self::Output {
215
250003
        let idx = *index;
216
250003
        debug_assert!(
217
312375000
            !self.freelist_iter().any(|free_idx| free_idx == idx),
218
            "Attempting to index free spot {}",
219
            index,
220
        );
221
        // SAFETY: The generational index ensures this slot was not freed after the index was issued.
222
250003
        unsafe { &self.roots[idx].object }
223
250003
    }
224
}
225

            
226
impl<T: fmt::Debug> fmt::Debug for ProtectionSet<T> {
227
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
228
        f.debug_map().entries(self.iter()).finish()
229
    }
230
}
231

            
232
/// An entry in the protection set, which can either be filled with an object or
233
/// free and point to the next free entry.
234
union Entry<T> {
235
    object: ManuallyDrop<T>,
236

            
237
    // The next free entry in the free list. If this points to itself, the free list is empty.
238
    next: usize,
239
}
240

            
241
/// Check that the Entry is the same size as a usize.
242
#[cfg(not(debug_assertions))]
243
const _: () = assert!(std::mem::size_of::<Entry<usize>>() == std::mem::size_of::<usize>());
244

            
245
/// An iterator over the filled entries in a protection set.
246
pub struct ProtSetIter<'a, T> {
247
    current: usize,
248
    protection_set: &'a ProtectionSet<T>,
249
    generation_counter: &'a GenerationCounter,
250
    /// Sorted free indices collected at iterator construction to skip free slots.
251
    free_indices: Vec<usize>,
252
}
253

            
254
impl<'a, T> Iterator for ProtSetIter<'a, T> {
255
    type Item = (ProtectionIndex, &'a T);
256

            
257
96291254
    fn next(&mut self) -> Option<Self::Item> {
258
        // Find the next valid entry, return it when found or None when end of roots is reached.
259
609967946
        while self.current < self.protection_set.roots.len() {
260
608413978
            let idx = self.current;
261
608413978
            self.current += 1;
262

            
263
608413978
            if self.free_indices.binary_search(&idx).is_err() {
264
                // SAFETY: idx is not in the free list, so it holds a valid filled object.
265
94737286
                let object = unsafe { &*self.protection_set.roots[idx].object };
266
94737286
                return Some((ProtectionIndex(self.generation_counter.recall_index(idx)), object));
267
513676692
            }
268
        }
269

            
270
1553968
        None
271
96291254
    }
272
}
273

            
274
/// An iterator over the free indices in the protection set freelist.
275
struct FreeListIter<'a, T> {
276
    current: Option<usize>,
277
    protection_set: &'a ProtectionSet<T>,
278
}
279

            
280
impl<'a, T> Iterator for FreeListIter<'a, T> {
281
    type Item = usize;
282

            
283
768973806132
    fn next(&mut self) -> Option<Self::Item> {
284
768973806132
        let curr = self.current?;
285
        // SAFETY: curr is a valid index in the free list, so it stores a `next` pointer.
286
768971286601
        let next = unsafe { self.protection_set.roots[curr].next };
287
768971286601
        if next == curr {
288
26892087
            // Sentinel: last free entry points to itself.
289
26892087
            self.current = None;
290
768944394514
        } else {
291
768944394514
            self.current = Some(next);
292
768944394514
        }
293
768971286601
        Some(curr)
294
768973806132
    }
295
}
296

            
297
impl<'a, T> IntoIterator for &'a ProtectionSet<T> {
298
    type Item = (ProtectionIndex, &'a T);
299
    type IntoIter = ProtSetIter<'a, T>;
300

            
301
    fn into_iter(self) -> Self::IntoIter {
302
        self.iter()
303
    }
304
}
305

            
306
#[cfg(test)]
307
mod tests {
308
    use rand::RngExt;
309

            
310
    use merc_utilities::random_test;
311
    use merc_utilities::test_logger;
312

            
313
    use crate::ProtectionIndex;
314
    use crate::ProtectionSet;
315

            
316
    #[test]
317
    #[cfg_attr(miri, ignore)]
318
1
    fn test_random_protection_set() {
319
100
        random_test(100, |rng| {
320
100
            let mut protection_set = ProtectionSet::<usize>::new();
321

            
322
            // Protect a number of indices and record their roots.
323
100
            let mut indices: Vec<ProtectionIndex> = Vec::new();
324

            
325
500000
            for _ in 0..5000 {
326
500000
                indices.push(protection_set.protect(rng.random_range(0..1000)));
327
500000
            }
328

            
329
            // Unprotect a number of roots.
330
250000
            for index in 0..2500 {
331
250000
                assert!(protection_set[indices[index]] <= 1000);
332
250000
                protection_set.unprotect(indices[index]);
333
250000
                indices.remove(index);
334
            }
335

            
336
            // Protect more to test the freelist
337
100000
            for _ in 0..1000 {
338
100000
                indices.push(protection_set.protect(rng.random_range(0..1000)));
339
100000
            }
340

            
341
350000
            for index in &indices {
342
350000
                assert!(
343
350000
                    protection_set.contains_root(*index),
344
                    "All indices that are not unprotected should occur in the protection set"
345
                );
346
            }
347

            
348
100
            assert_eq!(
349
100
                protection_set.iter().count(),
350
100
                6000 - 2500,
351
                "This is the number of roots remaining"
352
            );
353
100
            assert_eq!(protection_set.number_of_insertions(), 6000);
354
100
            assert!(protection_set.maximum_size() >= 5000);
355
100
            assert!(!protection_set.is_empty());
356
100
        });
357
1
    }
358

            
359
    #[test]
360
1
    fn test_protection_set_basic() {
361
1
        test_logger();
362

            
363
1
        let mut set = ProtectionSet::<String>::new();
364

            
365
        // Protect some values
366
1
        let idx1 = set.protect(String::from("value1"));
367
1
        let idx2 = set.protect(String::from("value2"));
368

            
369
        // Verify contains_root works
370
1
        assert!(set.contains_root(idx1));
371
1
        assert!(set.contains_root(idx2));
372

            
373
        // Test indexing
374
1
        assert_eq!(set[idx1], "value1");
375
1
        assert_eq!(set[idx2], "value2");
376

            
377
        // Test unprotect
378
1
        set.unprotect(idx1);
379
1
        assert!(!set.contains_root(idx1));
380
1
        assert!(set.contains_root(idx2));
381

            
382
        // Re-use freed slot
383
1
        let idx3 = set.protect(String::from("value3"));
384
1
        assert_eq!(set[idx3], "value3");
385
1
    }
386
}