1
use core::panic;
2
use std::ops::Index;
3

            
4
/// The protection set keeps track of nodes that should not be garbage
5
/// collected since they are being referenced by instances.
6
#[derive(Debug, Default)]
7
pub struct ProtectionSet<T> {
8
    roots: Vec<Entry<T>>, // The set of root active nodes.
9
    free: Option<usize>,
10
    number_of_insertions: u64,
11
    size: usize,
12
}
13

            
14
#[derive(Debug)]
15
enum Entry<T> {
16
    Filled(T),
17
    Free(usize),
18
}
19

            
20
impl<T> ProtectionSet<T> {
21
1247
    pub fn new() -> Self {
22
1247
        ProtectionSet {
23
1247
            roots: vec![],
24
1247
            free: None,
25
1247
            number_of_insertions: 0,
26
1247
            size: 0,
27
1247
        }
28
1247
    }
29

            
30
    /// Returns the number of insertions into the protection set.
31
1593
    pub fn number_of_insertions(&self) -> u64 {
32
1593
        self.number_of_insertions
33
1593
    }
34

            
35
    /// Returns maximum number of active instances.
36
1593
    pub fn maximum_size(&self) -> usize {
37
1593
        self.roots.capacity()
38
1593
    }
39

            
40
    /// Returns the number of roots in the protection set
41
2801
    pub fn len(&self) -> usize {
42
2801
        self.size
43
2801
    }
44

            
45
    /// Returns whether the protection set is empty.
46
1
    pub fn is_empty(&self) -> bool {
47
1
        self.len() == 0
48
1
    }
49

            
50
    /// Returns an iterator over all root indices in the protection set.
51
3565
    pub fn iter(&self) -> ProtSetIter<T> {
52
3565
        ProtSetIter {
53
3565
            current: 0,
54
3565
            protection_set: self,
55
3565
        }
56
3565
    }
57

            
58
    /// Protect the given root node to prevent garbage collection.
59
305364586
    pub fn protect(&mut self, object: T) -> usize {
60
305364586
        self.number_of_insertions += 1;
61
305364586
        self.size += 1;
62
305364586

            
63
305364586
        match self.free {
64
296493424
            Some(first) => {
65
296493424
                match &self.roots[first] {
66
296493424
                    Entry::Free(next) => {
67
296493424
                        if first == *next {
68
2479486
                            // The list is empty as its first element points to itself.
69
2479486
                            self.free = None;
70
294013938
                        } else {
71
294013938
                            // Update free to be the next element in the list.
72
294013938
                            self.free = Some(*next);
73
294013938
                        }
74
                    }
75
                    Entry::Filled(_) => {
76
                        panic!("The free list should not point a filled entry");
77
                    }
78
                }
79

            
80
296493424
                self.roots[first] = Entry::Filled(object);
81
296493424
                first
82
            }
83
            None => {
84
                // If free list is empty insert new entry into roots.
85
8871162
                self.roots.push(Entry::Filled(object));
86
8871162
                self.roots.len() - 1
87
            }
88
        }
89
305364586
    }
90

            
91
    /// Remove protection from the given LDD node. Note that index must be the
92
    /// index returned by the [ProtectionSet::protect] call.
93
305361086
    pub fn unprotect(&mut self, index: usize) {
94
305361086
        debug_assert!(
95
305361086
            matches!(self.roots[index], Entry::Filled(_)),
96
            "Index {index} is not valid"
97
        );
98
305361086
        self.size -= 1;
99
305361086

            
100
305361086
        match self.free {
101
302880855
            Some(next) => {
102
302880855
                self.roots[index] = Entry::Free(next);
103
302880855
            }
104
2480231
            None => {
105
2480231
                self.roots[index] = Entry::Free(index);
106
2480231
            }
107
        };
108

            
109
305361086
        self.free = Some(index);
110
305361086
    }
111
}
112

            
113
impl<T> Index<usize> for ProtectionSet<T> {
114
    type Output = T;
115

            
116
2500
    fn index(&self, index: usize) -> &Self::Output {
117
2500
        match &self.roots[index] {
118
2500
            Entry::Filled(value) => value,
119
            Entry::Free(_) => {
120
                panic!("Attempting to index free spot {}", index);
121
            }
122
        }
123
2500
    }
124
}
125

            
126
pub struct ProtSetIter<'a, T> {
127
    current: usize,
128
    protection_set: &'a ProtectionSet<T>,
129
}
130

            
131
impl<'a, T> Iterator for ProtSetIter<'a, T> {
132
    type Item = (&'a T, usize);
133

            
134
49131
    fn next(&mut self) -> Option<Self::Item> {
135
        // Find the next valid entry, return it when found or None when end of roots is reached.
136
61006
        while self.current < self.protection_set.roots.len() {
137
57441
            if let Entry::Filled(object) = &self.protection_set.roots[self.current] {
138
45566
                self.current += 1;
139
45566
                return Some((object, self.current - 1));
140
11875
            } else {
141
11875
                self.current += 1;
142
11875
            }
143
        }
144

            
145
3565
        None
146
49131
    }
147
}
148

            
149
impl<'a, T> IntoIterator for &'a ProtectionSet<T> {
150
    type Item = (&'a T, usize);
151
    type IntoIter = ProtSetIter<'a, T>;
152

            
153
398
    fn into_iter(self) -> Self::IntoIter {
154
398
        self.iter()
155
398
    }
156
}
157

            
158
#[cfg(test)]
159
mod tests {
160
    use rand::rngs::StdRng;
161
    use rand::Rng;
162
    use rand::SeedableRng;
163

            
164
    use super::*;
165

            
166
    #[test]
167
1
    fn test_protection_set() {
168
1
        let mut protection_set = ProtectionSet::<usize>::new();
169
1

            
170
1
        // Protect a number of indices and record their roots.
171
1
        let mut indices: Vec<usize> = Vec::new();
172
1

            
173
1
        let seed: u64 = rand::rng().random();
174
1
        println!("seed: {}", seed);
175
1
        let mut rng = StdRng::seed_from_u64(seed);
176

            
177
5001
        for _ in 0..5000 {
178
5000
            indices.push(protection_set.protect(rng.random_range(0..1000)));
179
5000
        }
180

            
181
        // Unprotect a number of roots.
182
2501
        for index in 0..2500 {
183
2500
            assert!(protection_set[indices[index]] <= 1000);
184
2500
            protection_set.unprotect(indices[index]);
185
2500
            indices.remove(index);
186
        }
187

            
188
        // Protect more to test the freelist
189
1001
        for _ in 0..1000 {
190
1000
            indices.push(protection_set.protect(rng.random_range(0..1000)));
191
1000
        }
192

            
193
3501
        for index in &indices {
194
3500
            assert!(
195
3500
                matches!(protection_set.roots[*index], Entry::Filled(_)),
196
                "All indices that are not unprotected should occur in the protection set"
197
            );
198
        }
199

            
200
1
        assert_eq!(
201
1
            protection_set.iter().count(),
202
1
            6000 - 2500,
203
            "This is the number of roots remaining"
204
        );
205
1
        assert_eq!(protection_set.number_of_insertions(), 6000);
206
1
        assert!(protection_set.maximum_size() >= 5000);
207
1
        assert!(!protection_set.is_empty());
208

            
209
1
        println!("{:?}", protection_set);
210
1

            
211
1
        // TODO: Fix this test.
212
1
        // for root in protection_set.iter() {
213
1
        //     assert!(indices.contains(root.0), "Root must be valid");
214
1
        // }
215
1
    }
216
}