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
    fn deref(&self) -> &Self::Target {
19
        &self.0
20
    }
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
7173
    pub fn new() -> Self {
53
7173
        ProtectionSet {
54
7173
            roots: Vec::new(),
55
7173
            free: None,
56
7173
            number_of_insertions: 0,
57
7173
            size: 0,
58
7173
            generation_counter: GenerationCounter::new(),
59
7173
        }
60
7173
    }
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
948835403
    pub fn protect(&mut self, object: T) -> ProtectionIndex {
84
948835403
        self.number_of_insertions += 1;
85
948835403
        self.size += 1;
86

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

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

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

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

            
116
    /// Remove protection from the given object. Note that index must be the
117
    /// index returned by the [ProtectionSet::protect] call.
118
    ///
119
    /// # Safety
120
    ///
121
    /// The caller must ensure that `index` refers to a currently protected
122
    /// entry and that it is not unprotected more than once.
123
948458905
    pub unsafe fn unprotect(&mut self, index: ProtectionIndex) {
124
948458905
        let index = self.generation_counter.get_index(index.0);
125
948458905
        self.size -= 1;
126

            
127
        // SAFETY: `index` refers to an occupied entry. We must drop the stored
128
        // object before overwriting the union slot with freelist metadata.
129
948458905
        unsafe {
130
948458905
            ManuallyDrop::drop(&mut self.roots[index].object);
131
948458905
        }
132

            
133
948458905
        match self.free {
134
934921561
            Some(next) => {
135
934921561
                self.roots[index] = Entry { next };
136
934921561
            }
137
13537344
            None => {
138
13537344
                self.roots[index] = Entry { next: index };
139
13537344
            }
140
        };
141

            
142
948458905
        self.free = Some(index);
143

            
144
        // Postcondition: verify the object was correctly removed from protection
145
948458905
        debug_assert!(
146
948458905
            self.freelist_iter().any(|free_idx| free_idx == index),
147
            "Failed to unprotect object"
148
        );
149
948458905
    }
150

            
151
    /// Replaces the object at the given index with the new object.
152
    pub fn replace(&mut self, index: ProtectionIndex, object: T) {
153
        let index = self.generation_counter.get_index(index.0);
154

            
155
        debug_assert!(
156
            !self.freelist_iter().any(|free_idx| free_idx == index),
157
            "Index {index} does not point to a filled entry"
158
        );
159

            
160
        self.roots[index] = Entry {
161
            object: ManuallyDrop::new(object),
162
        };
163
    }
164

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

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

            
191
    /// Returns whether the protection set contains the given index.
192
    ///
193
    /// This operation is O(n) in the size of the freelist.
194
350010
    pub fn contains_root(&self, index: ProtectionIndex) -> bool {
195
        // A stale index (whose slot was reused) is not a live root; report it as
196
        // absent rather than panicking in `get_index`.
197
350010
        if !self.generation_counter.is_valid(index.0) {
198
            return false;
199
350010
        }
200

            
201
350010
        let idx = self.generation_counter.get_index(index.0);
202
525000008
        !self.freelist_iter().any(|free_idx| free_idx == idx)
203
350010
    }
204
}
205

            
206
impl<T> Drop for ProtectionSet<T> {
207
5402
    fn drop(&mut self) {
208
19495435
        for index in 0..self.roots.len() {
209
1101306270877
            if self.freelist_iter().any(|free_idx| free_idx == index) {
210
19118937
                continue;
211
376498
            }
212

            
213
            // SAFETY: indices not present in the freelist still contain a live
214
            // protected object that must be dropped.
215
376498
            unsafe {
216
376498
                ManuallyDrop::drop(&mut self.roots[index].object);
217
376498
            }
218
        }
219
5402
    }
220
}
221

            
222
impl<T> Index<ProtectionIndex> for ProtectionSet<T> {
223
    type Output = T;
224

            
225
250003
    fn index(&self, index: ProtectionIndex) -> &Self::Output {
226
        // Panics in debug builds when the generation is stale, i.e., the slot was freed
227
        // (and possibly reused) after the index was issued.
228
250003
        let idx = self.generation_counter.get_index(index.0);
229
250003
        debug_assert!(
230
312375000
            !self.freelist_iter().any(|free_idx| free_idx == idx),
231
            "Attempting to index free spot {}",
232
            index,
233
        );
234
        // SAFETY: The generational index ensures this slot was not freed after the index was issued.
235
250003
        unsafe { &self.roots[idx].object }
236
250003
    }
237
}
238

            
239
impl<T: fmt::Debug> fmt::Debug for ProtectionSet<T> {
240
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
241
        f.debug_map().entries(self.iter()).finish()
242
    }
243
}
244

            
245
/// An entry in the protection set, which can either be filled with an object or
246
/// free and point to the next free entry.
247
union Entry<T> {
248
    object: ManuallyDrop<T>,
249

            
250
    // The next free entry in the free list. If this points to itself, the free list is empty.
251
    next: usize,
252
}
253

            
254
/// Check that the Entry is the same size as a usize; the union of
255
/// `ManuallyDrop<usize>` and `usize` is pointer-sized in every build.
256
const _: () = assert!(std::mem::size_of::<Entry<usize>>() == std::mem::size_of::<usize>());
257

            
258
/// An iterator over the filled entries in a protection set.
259
pub struct ProtSetIter<'a, T> {
260
    current: usize,
261
    protection_set: &'a ProtectionSet<T>,
262
    generation_counter: &'a GenerationCounter,
263
    /// Sorted free indices collected at iterator construction to skip free slots.
264
    free_indices: Vec<usize>,
265
}
266

            
267
impl<'a, T> Iterator for ProtSetIter<'a, T> {
268
    type Item = (ProtectionIndex, &'a T);
269

            
270
87197819
    fn next(&mut self) -> Option<Self::Item> {
271
        // Find the next valid entry, return it when found or None when end of roots is reached.
272
572024952
        while self.current < self.protection_set.roots.len() {
273
569791930
            let idx = self.current;
274
569791930
            self.current += 1;
275

            
276
569791930
            if self.free_indices.binary_search(&idx).is_err() {
277
                // SAFETY: idx is not in the free list, so it holds a valid filled object.
278
84964797
                let object = unsafe { &*self.protection_set.roots[idx].object };
279
84964797
                return Some((ProtectionIndex(self.generation_counter.recall_index(idx)), object));
280
484827133
            }
281
        }
282

            
283
2233022
        None
284
87197819
    }
285
}
286

            
287
/// An iterator over the free indices in the protection set freelist.
288
struct FreeListIter<'a, T> {
289
    current: Option<usize>,
290
    protection_set: &'a ProtectionSet<T>,
291
}
292

            
293
impl<'a, T> Iterator for FreeListIter<'a, T> {
294
    type Item = usize;
295

            
296
1103580141449
    fn next(&mut self) -> Option<Self::Item> {
297
1103580141449
        let curr = self.current?;
298
        // SAFETY: curr is a valid index in the free list, so it stores a `next` pointer.
299
1103576931920
        let next = unsafe { self.protection_set.roots[curr].next };
300
1103576931920
        if next == curr {
301
15682569
            // Sentinel: last free entry points to itself.
302
15682569
            self.current = None;
303
1103561249351
        } else {
304
1103561249351
            self.current = Some(next);
305
1103561249351
        }
306
1103576931920
        Some(curr)
307
1103580141449
    }
308
}
309

            
310
impl<'a, T> IntoIterator for &'a ProtectionSet<T> {
311
    type Item = (ProtectionIndex, &'a T);
312
    type IntoIter = ProtSetIter<'a, T>;
313

            
314
    fn into_iter(self) -> Self::IntoIter {
315
        self.iter()
316
    }
317
}
318

            
319
#[cfg(test)]
320
mod tests {
321
    use rand::RngExt;
322

            
323
    use merc_utilities::random_test;
324
    use merc_utilities::test_logger;
325

            
326
    use super::ProtectionIndex;
327
    use super::ProtectionSet;
328

            
329
    #[test]
330
    #[cfg_attr(miri, ignore)]
331
1
    fn test_random_protection_set() {
332
100
        random_test(100, |rng| {
333
100
            let mut protection_set = ProtectionSet::<usize>::new();
334

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

            
338
500000
            for _ in 0..5000 {
339
500000
                indices.push(protection_set.protect(rng.random_range(0..1000)));
340
500000
            }
341

            
342
            // Unprotect a number of roots.
343
250000
            for index in 0..2500 {
344
250000
                assert!(protection_set[indices[index]] <= 1000);
345
                // SAFETY: each recorded index is protected exactly once and is
346
                // unprotected only here, so there is no double-unprotect.
347
250000
                unsafe {
348
250000
                    protection_set.unprotect(indices[index]);
349
250000
                }
350
250000
                indices.remove(index);
351
            }
352

            
353
            // Protect more to test the freelist
354
100000
            for _ in 0..1000 {
355
100000
                indices.push(protection_set.protect(rng.random_range(0..1000)));
356
100000
            }
357

            
358
350000
            for index in &indices {
359
350000
                assert!(
360
350000
                    protection_set.contains_root(*index),
361
                    "All indices that are not unprotected should occur in the protection set"
362
                );
363
            }
364

            
365
100
            assert_eq!(
366
100
                protection_set.iter().count(),
367
100
                6000 - 2500,
368
                "This is the number of roots remaining"
369
            );
370
100
            assert_eq!(protection_set.number_of_insertions(), 6000);
371
100
            assert!(protection_set.maximum_size() >= 5000);
372
100
            assert!(!protection_set.is_empty());
373
100
        });
374
1
    }
375

            
376
    #[test]
377
1
    fn test_protection_set_basic() {
378
1
        test_logger();
379

            
380
1
        let mut set = ProtectionSet::<String>::new();
381

            
382
        // Protect some values
383
1
        let idx1 = set.protect(String::from("value1"));
384
1
        let idx2 = set.protect(String::from("value2"));
385

            
386
        // Verify contains_root works
387
1
        assert!(set.contains_root(idx1));
388
1
        assert!(set.contains_root(idx2));
389

            
390
        // Test indexing
391
1
        assert_eq!(set[idx1], "value1");
392
1
        assert_eq!(set[idx2], "value2");
393

            
394
        // Test unprotect
395
        // SAFETY: `idx1` is currently protected and is unprotected only once.
396
1
        unsafe {
397
1
            set.unprotect(idx1);
398
1
        }
399
1
        assert!(!set.contains_root(idx1));
400
1
        assert!(set.contains_root(idx2));
401

            
402
        // Re-use freed slot
403
1
        let idx3 = set.protect(String::from("value3"));
404
1
        assert_eq!(set[idx3], "value3");
405
1
    }
406
}
407

            
408
#[cfg(kani)]
409
mod verification {
410
    use super::*;
411

            
412
    /// Covers the protect/index/replace/unprotect round trip and the empty-set
413
    /// pre/post conditions in a single harness.
414
    #[kani::proof]
415
    #[kani::unwind(3)]
416
    fn protection_set_roundtrip() {
417
        let mut ps: ProtectionSet<u32> = ProtectionSet::new();
418
        assert!(ps.is_empty());
419

            
420
        let v1: u32 = kani::any();
421
        let v2: u32 = kani::any();
422

            
423
        let idx = ps.protect(v1);
424
        assert_eq!(ps.len(), 1);
425
        assert!(ps.contains_root(idx));
426
        assert_eq!(ps[idx], v1);
427

            
428
        ps.replace(idx, v2);
429
        assert_eq!(ps[idx], v2);
430

            
431
        // SAFETY: `idx` is currently protected and is unprotected only once.
432
        unsafe {
433
            ps.unprotect(idx);
434
        }
435
        assert!(ps.is_empty());
436
        assert!(!ps.contains_root(idx));
437
        // The insertion counter is monotonic and is not affected by unprotect.
438
        assert_eq!(ps.number_of_insertions(), 1);
439
    }
440

            
441
    /// The generation counter must invalidate stale indices even when their
442
    /// underlying slot is reused — this is the load-bearing safety invariant.
443
    #[kani::proof]
444
    #[kani::unwind(3)]
445
    fn protection_set_stale_index_after_slot_reuse() {
446
        let mut ps: ProtectionSet<u32> = ProtectionSet::new();
447
        let v1: u32 = kani::any();
448
        let v2: u32 = kani::any();
449

            
450
        let i1 = ps.protect(v1);
451
        // SAFETY: `i1` is currently protected and is unprotected only once.
452
        unsafe {
453
            ps.unprotect(i1);
454
        }
455
        let i2 = ps.protect(v2);
456

            
457
        assert!(ps.contains_root(i2));
458
        assert!(!ps.contains_root(i1));
459
        assert_eq!(ps[i2], v2);
460
    }
461
}