1
//! Provides a generational index implementation that offers generation checking
2
//! in debug builds while having zero runtime cost in release builds.
3

            
4
use std::fmt;
5
use std::hash::Hash;
6
use std::hash::Hasher;
7
use std::ops::Deref;
8

            
9
/// A generational index that stores both an index and a generation counter.
10
/// The generation is only tracked in debug builds to avoid overhead in release.
11
///
12
/// This allows detecting use-after-free scenarios in debug mode while
13
/// maintaining zero overhead in release mode.
14
#[repr(C)]
15
#[derive(Copy, Clone)]
16
pub struct GenerationalIndex<I: Copy + Into<usize> = usize> {
17
    /// The raw index value
18
    index: I,
19

            
20
    #[cfg(debug_assertions)]
21
    /// Generation counter, only available in debug builds
22
    generation: usize,
23
}
24

            
25
impl Default for GenerationalIndex<usize> {
26
1
    fn default() -> Self {
27
1
        GenerationalIndex {
28
1
            index: 0,
29
1
            #[cfg(debug_assertions)]
30
1
            generation: usize::MAX,
31
1
        }
32
1
    }
33
}
34

            
35
impl<I: Copy + Into<usize>> Deref for GenerationalIndex<I> {
36
    type Target = I;
37

            
38
    /// Deref implementation to access the underlying index value.
39
103400904
    fn deref(&self) -> &Self::Target {
40
103400904
        &self.index
41
103400904
    }
42
}
43

            
44
impl<I: Copy + Into<usize>> GenerationalIndex<I> {
45
    /// Creates a new generational index with the specified index.
46
    #[cfg(debug_assertions)]
47
1100395962
    fn new(index: I, generation: usize) -> Self {
48
1100395962
        Self { index, generation }
49
1100395962
    }
50

            
51
    /// Creates a new generational index with the specified index and generation.
52
    #[cfg(not(debug_assertions))]
53
    fn new(index: I) -> Self {
54
        Self { index }
55
    }
56
}
57

            
58
/// A counter that keeps track of generational indices.
59
/// This helps manage generations of indices to detect use-after-free and similar issues.
60
#[derive(Clone, Debug, Default)]
61
pub struct GenerationCounter {
62
    /// Current generation count, only stored in debug builds
63
    #[cfg(debug_assertions)]
64
    current_generation: Vec<usize>,
65
}
66

            
67
impl GenerationCounter {
68
    /// Creates a new generation counter.
69
106663
    pub fn new() -> Self {
70
        #[cfg(debug_assertions)]
71
        {
72
106663
            Self {
73
106663
                current_generation: Vec::new(),
74
106663
            }
75
        }
76

            
77
        #[cfg(not(debug_assertions))]
78
        Self {}
79
106663
    }
80
}
81

            
82
impl GenerationCounter {
83
    /// Creates a new generational index with the given index and the next generation.
84
960793278
    pub fn create_index<I>(&mut self, index: I) -> GenerationalIndex<I>
85
960793278
    where
86
960793278
        I: Copy + Into<usize>,
87
    {
88
        #[cfg(debug_assertions)]
89
        {
90
960793278
            let generation = if self.current_generation.len() <= index.into() {
91
31453340
                self.current_generation.resize(index.into() + 1, 0);
92
31453340
                0
93
            } else {
94
929339938
                let generation = &mut self.current_generation[index.into()];
95
929339938
                *generation = generation.wrapping_add(1);
96
929339938
                *generation
97
            };
98

            
99
960793278
            GenerationalIndex::new(index, generation)
100
        }
101

            
102
        #[cfg(not(debug_assertions))]
103
        {
104
            GenerationalIndex::new(index)
105
        }
106
960793278
    }
107

            
108
    /// Returns a generational index with the given index and the current generation.
109
139602684
    pub fn recall_index<I>(&self, index: I) -> GenerationalIndex<I>
110
139602684
    where
111
139602684
        I: Copy + Into<usize>,
112
    {
113
        #[cfg(debug_assertions)]
114
        {
115
139602684
            let idx = index.into();
116
139602684
            assert!(
117
139602684
                idx < self.current_generation.len(),
118
                "recall_index called with index {idx} that was never created"
119
            );
120
139602684
            GenerationalIndex::new(index, self.current_generation[idx])
121
        }
122
        #[cfg(not(debug_assertions))]
123
        {
124
            GenerationalIndex::new(index)
125
        }
126
139602684
    }
127

            
128
    /// Returns whether the given generational index is still valid, i.e. its
129
    /// generation matches the current generation of the underlying slot.
130
350010
    pub fn is_valid<I>(&self, index: GenerationalIndex<I>) -> bool
131
350010
    where
132
350010
        I: Copy + Into<usize>,
133
    {
134
        #[cfg(debug_assertions)]
135
        {
136
350010
            let idx = index.index.into();
137
350010
            assert!(
138
350010
                idx < self.current_generation.len(),
139
                "is_valid called with index {idx} that was never created"
140
            );
141
350010
            self.current_generation[idx] == index.generation
142
        }
143
        #[cfg(not(debug_assertions))]
144
        {
145
            let _ = index;
146
            true
147
        }
148
350010
    }
149

            
150
    /// Returns the underlying index, checks if the generation is correct.
151
949061118
    pub fn get_index<I>(&self, index: GenerationalIndex<I>) -> I
152
949061118
    where
153
949061118
        I: Copy + Into<usize> + fmt::Debug,
154
    {
155
        #[cfg(debug_assertions)]
156
        {
157
949061118
            let idx = index.index.into();
158
949061118
            assert!(
159
949061118
                idx < self.current_generation.len(),
160
                "get_index called with index {idx} that was never created"
161
            );
162
949061118
            if self.current_generation[idx] != index.generation {
163
                panic!("Attempting to access an invalid index: {index:?}");
164
949061118
            }
165
        }
166

            
167
949061118
        index.index
168
949061118
    }
169
}
170

            
171
// Standard trait implementations for GenerationalIndex
172

            
173
impl<I> PartialEq for GenerationalIndex<I>
174
where
175
    I: Copy + Into<usize> + Eq,
176
{
177
5278
    fn eq(&self, other: &Self) -> bool {
178
        #[cfg(debug_assertions)]
179
        {
180
            // A sentinel (default) index has the `usize::MAX` generation. It is
181
            // only equal to another sentinel index with the same raw index, and
182
            // never aliases a live generation. Handling it here keeps `eq`
183
            // reflexive (`default == default`), which is required for `Eq` and
184
            // for use as a hash-map key.
185
5278
            let self_sentinel = self.generation == usize::MAX;
186
5278
            let other_sentinel = other.generation == usize::MAX;
187
5278
            if self_sentinel || other_sentinel {
188
2
                return self_sentinel && other_sentinel && self.index == other.index;
189
5276
            }
190

            
191
5276
            debug_assert_eq!(
192
                self.generation, other.generation,
193
                "Comparing indices of different generations"
194
            );
195
        }
196

            
197
5275
        self.index == other.index
198
5277
    }
199
}
200

            
201
impl<I> Eq for GenerationalIndex<I> where I: Copy + Into<usize> + Eq {}
202

            
203
impl<I> PartialOrd for GenerationalIndex<I>
204
where
205
    I: Copy + Into<usize> + PartialOrd + Eq,
206
{
207
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
208
        // Skip the generation check for the sentinel (default) generation, to
209
        // stay consistent with `PartialEq`.
210
        #[cfg(debug_assertions)]
211
        if self.generation != usize::MAX && other.generation != usize::MAX {
212
            debug_assert_eq!(
213
                self.generation, other.generation,
214
                "Comparing indices of different generations"
215
            );
216
        }
217

            
218
        self.index.partial_cmp(&other.index)
219
    }
220
}
221

            
222
impl<I> Ord for GenerationalIndex<I>
223
where
224
    I: Copy + Into<usize> + Eq + Ord,
225
{
226
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
227
        // Skip the generation check for the sentinel (default) generation, to
228
        // stay consistent with `PartialEq`.
229
        #[cfg(debug_assertions)]
230
        if self.generation != usize::MAX && other.generation != usize::MAX {
231
            debug_assert_eq!(
232
                self.generation, other.generation,
233
                "Comparing indices of different generations"
234
            );
235
        }
236
        self.index.cmp(&other.index)
237
    }
238
}
239

            
240
impl<I> Hash for GenerationalIndex<I>
241
where
242
    I: Copy + Into<usize> + Hash,
243
{
244
2
    fn hash<H: Hasher>(&self, state: &mut H) {
245
2
        self.index.hash(state);
246
2
    }
247
}
248

            
249
impl<I> fmt::Debug for GenerationalIndex<I>
250
where
251
    I: Copy + Into<usize> + fmt::Debug,
252
{
253
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
254
        #[cfg(debug_assertions)]
255
        {
256
            write!(
257
                f,
258
                "GenerationalIndex(index: {:?}, generation: {})",
259
                self.index, self.generation
260
            )
261
        }
262
        #[cfg(not(debug_assertions))]
263
        {
264
            write!(f, "GenerationalIndex(index: {:?})", self.index)
265
        }
266
    }
267
}
268

            
269
impl fmt::Display for GenerationalIndex<usize> {
270
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
271
        write!(f, "{}", self.index)
272
    }
273
}
274

            
275
#[cfg(test)]
276
mod tests {
277
    #[cfg(debug_assertions)]
278
    use crate::GenerationCounter;
279
    use crate::GenerationalIndex;
280

            
281
    #[test]
282
1
    fn test_default_index_is_reflexive() {
283
        // A default (sentinel) index must equal itself, otherwise `Eq` is
284
        // violated and it cannot be used as a hash-map key.
285
1
        let default = GenerationalIndex::<usize>::default();
286
1
        assert_eq!(default, default);
287

            
288
1
        let mut set = std::collections::HashSet::new();
289
1
        set.insert(default);
290
1
        assert!(set.contains(&default));
291
1
    }
292

            
293
    #[test]
294
    #[should_panic]
295
    #[cfg(debug_assertions)]
296
1
    fn test_generational_index_equality() {
297
1
        let mut counter = GenerationCounter::new();
298
1
        let idx1 = counter.create_index(42usize);
299
1
        let idx2 = counter.create_index(42usize);
300
1
        let idx4 = counter.create_index(43usize);
301

            
302
1
        let idx3 = counter.recall_index(42usize);
303

            
304
1
        assert_ne!(idx1, idx4);
305
1
        assert_eq!(idx2, idx3);
306

            
307
        // This panics since idx1 and idx2 are from different generations
308
1
        assert_eq!(idx1, idx2);
309
    }
310
}