1
use core::panic;
2
use std::fmt;
3
use std::hash::Hash;
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
/// Is is similar to a [ `crate::IndexedSet`], except that we cannot look up elements by value.
39
#[derive(Debug, 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
    /// The number of generations
46
    generation_counter: GenerationCounter,
47
}
48

            
49
/// TODO: Is it possible to get the size of entries down to a sizeof(NonZero<usize>)?
50
#[derive(Debug)]
51
enum Entry<T> {
52
    Filled(T),
53
    Free(usize),
54
}
55

            
56
impl<T> ProtectionSet<T> {
57
    /// Creates a new empty protection set.
58
3734
    pub fn new() -> Self {
59
3734
        ProtectionSet {
60
3734
            roots: Vec::new(),
61
3734
            free: None,
62
3734
            number_of_insertions: 0,
63
3734
            size: 0,
64
3734
            generation_counter: GenerationCounter::new(),
65
3734
        }
66
3734
    }
67

            
68
    /// Returns the number of insertions into the protection set.
69
100
    pub fn number_of_insertions(&self) -> u64 {
70
100
        self.number_of_insertions
71
100
    }
72

            
73
    /// Returns maximum number of active instances.
74
100
    pub fn maximum_size(&self) -> usize {
75
100
        self.roots.capacity()
76
100
    }
77

            
78
    /// Returns the number of roots in the protection set
79
100
    pub fn len(&self) -> usize {
80
100
        self.size
81
100
    }
82

            
83
    /// Returns whether the protection set is empty.
84
100
    pub fn is_empty(&self) -> bool {
85
100
        self.len() == 0
86
100
    }
87

            
88
    /// Returns an iterator over all root indices in the protection set.
89
340
    pub fn iter(&self) -> ProtSetIter<'_, T> {
90
340
        ProtSetIter {
91
340
            current: 0,
92
340
            protection_set: self,
93
340
            generation_counter: &self.generation_counter,
94
340
        }
95
340
    }
96

            
97
    /// Returns whether the protection set contains the given index.
98
350010
    pub fn contains_root(&self, index: ProtectionIndex) -> bool {
99
350010
        matches!(self.roots[self.generation_counter.get_index(index.0)], Entry::Filled(_))
100
350010
    }
101

            
102
    /// Adds the given object to the protection set and returns its index.
103
462415051
    pub fn protect(&mut self, object: T) -> ProtectionIndex {
104
462415051
        self.number_of_insertions += 1;
105
462415051
        self.size += 1;
106

            
107
462415051
        let index = match self.free {
108
453185207
            Some(first) => {
109
453185207
                match &self.roots[first] {
110
453185207
                    Entry::Free(next) => {
111
453185207
                        if first == *next {
112
19836463
                            // The list is empty as its first element points to itself.
113
19836463
                            self.free = None;
114
433348744
                        } else {
115
433348744
                            // Update free to be the next element in the list.
116
433348744
                            self.free = Some(*next);
117
433348744
                        }
118
                    }
119
                    Entry::Filled(_) => {
120
                        panic!("The free list should not point a filled entry");
121
                    }
122
                }
123

            
124
453185207
                self.roots[first] = Entry::Filled(object);
125
453185207
                first
126
            }
127
            None => {
128
                // If free list is empty insert new entry into roots.
129
9229844
                self.roots.push(Entry::Filled(object));
130
9229844
                let index = self.roots.len() - 1;
131

            
132
                // Postcondition: verify the object was correctly added
133
9229844
                debug_assert!(
134
9229844
                    matches!(self.roots[index], Entry::Filled(_)),
135
                    "Failed to add object to protection set"
136
                );
137

            
138
9229844
                index
139
            }
140
        };
141

            
142
462415051
        ProtectionIndex(self.generation_counter.create_index(index))
143
462415051
    }
144

            
145
    /// Remove protection from the given object. Note that index must be the
146
    /// index returned by the [ProtectionSet::protect] call.
147
462062749
    pub fn unprotect(&mut self, index: ProtectionIndex) {
148
462062749
        let index = self.generation_counter.get_index(index.0);
149

            
150
462062749
        debug_assert!(
151
462062749
            matches!(self.roots[index], Entry::Filled(_)),
152
            "Index {index} is does not point to a filled entry"
153
        );
154

            
155
462062749
        self.size -= 1;
156

            
157
462062749
        match self.free {
158
442222699
            Some(next) => {
159
442222699
                self.roots[index] = Entry::Free(next);
160
442222699
            }
161
19840050
            None => {
162
19840050
                self.roots[index] = Entry::Free(index);
163
19840050
            }
164
        };
165

            
166
462062749
        self.free = Some(index);
167

            
168
        // Postcondition: verify the object was correctly removed from protection
169
462062749
        debug_assert!(
170
462062749
            matches!(self.roots[index], Entry::Free(_)),
171
            "Failed to unprotect object"
172
        );
173
462062749
    }
174

            
175
    /// Replaces the object at the given index with the new object.
176
    pub fn replace(&mut self, index: ProtectionIndex, object: T) {
177
        let index = self.generation_counter.get_index(index.0);
178

            
179
        debug_assert!(
180
            matches!(self.roots[index], Entry::Filled(_)),
181
            "Index {index} is does not point to a filled entry"
182
        );
183

            
184
        self.roots[index] = Entry::Filled(object);
185
    }
186
}
187

            
188
impl<T> Index<ProtectionIndex> for ProtectionSet<T> {
189
    type Output = T;
190

            
191
250003
    fn index(&self, index: ProtectionIndex) -> &Self::Output {
192
250003
        match &self.roots[*index] {
193
250003
            Entry::Filled(value) => value,
194
            Entry::Free(_) => {
195
                panic!("Attempting to index free spot {}", index);
196
            }
197
        }
198
250003
    }
199
}
200

            
201
pub struct ProtSetIter<'a, T> {
202
    current: usize,
203
    protection_set: &'a ProtectionSet<T>,
204
    generation_counter: &'a GenerationCounter,
205
}
206

            
207
impl<'a, T> Iterator for ProtSetIter<'a, T> {
208
    type Item = (ProtectionIndex, &'a T);
209

            
210
351180
    fn next(&mut self) -> Option<Self::Item> {
211
        // Find the next valid entry, return it when found or None when end of roots is reached.
212
501620
        while self.current < self.protection_set.roots.len() {
213
501280
            let idx = self.current;
214
501280
            self.current += 1;
215

            
216
501280
            if let Entry::Filled(object) = &self.protection_set.roots[idx] {
217
350840
                return Some((ProtectionIndex(self.generation_counter.recall_index(idx)), object));
218
150440
            }
219
        }
220

            
221
340
        None
222
351180
    }
223
}
224

            
225
impl<'a, T> IntoIterator for &'a ProtectionSet<T> {
226
    type Item = (ProtectionIndex, &'a T);
227
    type IntoIter = ProtSetIter<'a, T>;
228

            
229
    fn into_iter(self) -> Self::IntoIter {
230
        self.iter()
231
    }
232
}
233

            
234
#[cfg(test)]
235
mod tests {
236
    use super::*;
237

            
238
    use rand::Rng;
239

            
240
    use merc_utilities::random_test;
241
    use merc_utilities::test_logger;
242

            
243
    #[test]
244
    #[cfg_attr(miri, ignore)]
245
1
    fn test_random_protection_set() {
246
100
        random_test(100, |rng| {
247
100
            let mut protection_set = ProtectionSet::<usize>::new();
248

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

            
252
500000
            for _ in 0..5000 {
253
500000
                indices.push(protection_set.protect(rng.random_range(0..1000)));
254
500000
            }
255

            
256
            // Unprotect a number of roots.
257
250000
            for index in 0..2500 {
258
250000
                assert!(protection_set[indices[index]] <= 1000);
259
250000
                protection_set.unprotect(indices[index]);
260
250000
                indices.remove(index);
261
            }
262

            
263
            // Protect more to test the freelist
264
100000
            for _ in 0..1000 {
265
100000
                indices.push(protection_set.protect(rng.random_range(0..1000)));
266
100000
            }
267

            
268
350000
            for index in &indices {
269
350000
                assert!(
270
350000
                    protection_set.contains_root(*index),
271
                    "All indices that are not unprotected should occur in the protection set"
272
                );
273
            }
274

            
275
100
            assert_eq!(
276
100
                protection_set.iter().count(),
277
100
                6000 - 2500,
278
                "This is the number of roots remaining"
279
            );
280
100
            assert_eq!(protection_set.number_of_insertions(), 6000);
281
100
            assert!(protection_set.maximum_size() >= 5000);
282
100
            assert!(!protection_set.is_empty());
283
100
        });
284
1
    }
285

            
286
    #[test]
287
1
    fn test_protection_set_basic() {
288
1
        test_logger();
289

            
290
1
        let mut set = ProtectionSet::<String>::new();
291

            
292
        // Protect some values
293
1
        let idx1 = set.protect(String::from("value1"));
294
1
        let idx2 = set.protect(String::from("value2"));
295

            
296
        // Verify contains_root works
297
1
        assert!(set.contains_root(idx1));
298
1
        assert!(set.contains_root(idx2));
299

            
300
        // Test indexing
301
1
        assert_eq!(set[idx1], "value1");
302
1
        assert_eq!(set[idx2], "value2");
303

            
304
        // Test unprotect
305
1
        set.unprotect(idx1);
306
1
        assert!(!set.contains_root(idx1));
307
1
        assert!(set.contains_root(idx2));
308

            
309
        // Re-use freed slot
310
1
        let idx3 = set.protect(String::from("value3"));
311
1
        assert_eq!(set[idx3], "value3");
312
1
    }
313
}