1
use alloc::vec::{self, Vec};
2
use core::fmt::Debug;
3
use core::ops::{Index, IndexMut};
4
use core::{mem, slice};
5

            
6
use crate::{Assert, Generation, LotId};
7

            
8
/// A collection of `T`, organized by generational [`LotId`]s.
9
///
10
/// This data type allows storing data of any type and receiving a [`LotId`]
11
/// that can later be used to look up the data.
12
///
13
/// This data type cannot hold more than 2^48 items, due how [`LotId`]s are
14
/// allocated.
15
///
16
/// ## Generation Checks
17
///
18
/// A [`LotId`] contains 16 bits representing a lot's generation. Each time a
19
/// lot is updated, the lot's generation is incremented (with wrapping).
20
///
21
/// The lot's generation is checked when retrieving data using a [`LotId`]. If a
22
/// generation mismatch is found, the data will not be returned.
23
///
24
/// While the chances of generation collision may be low, this is not a perfect
25
/// check. Care should still be taken to ensure stale [`LotId`]s aren't used
26
/// when other ways of validating the data don't exist.
27
#[derive(Clone)]
28
pub struct Lots<T> {
29
    slots: Vec<SlotData<T>>,
30
    free_indicies: Vec<usize>,
31
}
32

            
33
impl<T> Eq for Lots<T> where T: Eq {}
34

            
35
impl<T> PartialEq for Lots<T>
36
where
37
    T: PartialEq,
38
{
39
2
    fn eq(&self, other: &Self) -> bool {
40
2
        self.slots == other.slots
41
2
    }
42
}
43

            
44
impl<T> Default for Lots<T> {
45
    #[inline]
46
3
    fn default() -> Self {
47
3
        Self::new()
48
3
    }
49
}
50

            
51
impl<T> Debug for Lots<T>
52
where
53
    T: Debug,
54
{
55
1
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
56
1
        let mut map = f.debug_map();
57
2
        for (id, value) in self.entries() {
58
2
            map.entry(&id, &value);
59
2
        }
60
1
        map.finish()
61
1
    }
62
}
63

            
64
impl<T> Lots<T> {
65
    /// Returns a new, empty collection.
66
    #[inline]
67
    #[must_use]
68
17
    pub const fn new() -> Self {
69
17
        Self {
70
17
            slots: Vec::new(),
71
17
            free_indicies: Vec::new(),
72
17
        }
73
17
    }
74

            
75
    /// Returns an empty collection that can hold `initial_capacity` values
76
    /// without reallocation.
77
    #[inline]
78
    #[must_use]
79
4
    pub fn with_capacity(initial_capacity: usize) -> Self {
80
4
        Self {
81
4
            slots: Vec::with_capacity(initial_capacity),
82
4
            free_indicies: Vec::new(),
83
4
        }
84
4
    }
85

            
86
    /// Adds `value` to the collection, returning the value's unique [`LotId`].
87
    #[inline]
88
    #[allow(clippy::must_use_candidate)]
89
67
    pub fn push(&mut self, value: T) -> LotId {
90
67
        if let Some((index, SlotData::Empty { generation })) = self
91
67
            .free_indicies
92
67
            .pop()
93
67
            .and_then(|index| self.slots.get(index).map(|lot| (index, lot)))
94
        {
95
13
            let generation = generation.next();
96
13
            self.slots[index] = SlotData::Populated {
97
13
                generation,
98
13
                contents: value,
99
13
            };
100
13
            LotId::new(generation, index).assert("invalid lot id")
101
        } else {
102
54
            let index = self.slots.len();
103
54
            let generation = Generation::first();
104
54
            self.slots.push(SlotData::Populated {
105
54
                generation,
106
54
                contents: value,
107
54
            });
108
54

            
109
54
            LotId::new(generation, index).assert("invalid lot id")
110
        }
111
67
    }
112

            
113
    /// Removes all values from the collection.
114
    #[inline]
115
2
    pub fn clear(&mut self) {
116
2
        self.drain();
117
2
    }
118

            
119
    /// Looks up a previously stored value by its [`LotId`]. If the value hasn't
120
    /// been removed, a reference will be returned.
121
    ///
122
    /// Note: It is possible, but unlikely, for a [`LotId`] that has been
123
    /// removed to be reused.
124
    #[inline]
125
    #[must_use]
126
122
    pub fn get(&self, id: LotId) -> Option<&T> {
127
122
        match self.slots.get(id.index()) {
128
            Some(SlotData::Populated {
129
119
                generation,
130
119
                contents,
131
119
            }) if *generation == id.generation() => Some(contents),
132
3
            _ => None,
133
        }
134
122
    }
135

            
136
    /// Looks up a previously stored value by its [`LotId`]. If the value hasn't
137
    /// been removed, a mutable reference will be returned.
138
    ///
139
    /// Note: It is possible, but unlikely, for a [`LotId`] that has been
140
    /// removed to be reused.
141
    #[inline]
142
    #[must_use]
143
12
    pub fn get_mut(&mut self, id: LotId) -> Option<&mut T> {
144
12
        match self.slots.get_mut(id.index()) {
145
            Some(SlotData::Populated {
146
10
                generation,
147
10
                contents,
148
10
            }) if *generation == id.generation() => Some(contents),
149
2
            _ => None,
150
        }
151
12
    }
152

            
153
    /// Removes a previously stored value by its [`LotId`]. If the value is
154
    /// present, it will be removed and returned.
155
    ///
156
    /// Note: It is possible, but unlikely, for a [`LotId`] that has been
157
    /// removed to be reused.
158
    #[inline]
159
28
    pub fn remove(&mut self, id: LotId) -> Option<T> {
160
28
        match self.slots.get_mut(id.index()) {
161
27
            Some(lot) if lot.generation() == id.generation() => {
162
27
                if let SlotData::Populated { .. } = lot {
163
25
                    let generation = lot.generation();
164
25
                    let SlotData::Populated { contents, .. } =
165
25
                        mem::replace(lot, SlotData::Empty { generation })
166
                    else {
167
                        unreachable!()
168
                    };
169
25
                    self.free_indicies.push(id.index());
170
25
                    return Some(contents);
171
2
                }
172
            }
173
1
            _ => {}
174
        }
175

            
176
3
        None
177
28
    }
178

            
179
    /// Returns an iterator that returns all the contained values in this
180
    /// collection as they're removed from the collection.
181
    ///
182
    /// Dropping the iterator will still result in the elements being removed.
183
    #[inline]
184
3
    pub fn drain(&mut self) -> Drain<'_, T, DrainAll> {
185
3
        Drain {
186
3
            map: self,
187
3
            index: 0,
188
3
            filter: DrainAll,
189
3
        }
190
3
    }
191

            
192
    /// Returns an iterator that invokes `filter` for each item in the
193
    /// collection. If `filter` returns true for that value, it will be removed
194
    /// and returned from the iterator. When false is returned, the value is
195
    /// kept in the collection.
196
    ///
197
    /// Dropping the iterator will still result in the filtered elements being
198
    /// removed.
199
    #[inline]
200
1
    pub fn drain_filter<Filter>(&mut self, filter: Filter) -> Drain<'_, T, Filter>
201
1
    where
202
1
        Filter: FnMut(&mut T) -> bool,
203
1
    {
204
1
        Drain {
205
1
            map: self,
206
1
            index: 0,
207
1
            filter,
208
1
        }
209
1
    }
210

            
211
    /// Returns the number of values contained in this collection.
212
    #[inline]
213
    #[must_use]
214
1
    pub fn len(&self) -> usize {
215
1
        self.slots.len() - self.free_indicies.len()
216
1
    }
217

            
218
    /// Returns true if this collection has no values.
219
    #[inline]
220
    #[must_use]
221
1
    pub fn is_empty(&self) -> bool {
222
1
        self.slots.len() == self.free_indicies.len()
223
1
    }
224

            
225
    /// Returns an iterator of references to all contained values.
226
    #[inline]
227
    #[must_use]
228
1
    pub fn iter(&self) -> Iter<'_, T> {
229
1
        self.into_iter()
230
1
    }
231

            
232
    /// Returns an iterator of exclusive references to all contained values.
233
    #[inline]
234
    #[must_use]
235
1
    pub fn iter_mut(&mut self) -> IterMut<'_, T> {
236
1
        self.into_iter()
237
1
    }
238

            
239
    /// Returns an `Iterator<Item = (LotId, &T)>` for all contained values.
240
    #[inline]
241
    #[must_use]
242
2
    pub fn entries(&self) -> EntryIter<'_, T> {
243
2
        EntryIter(self.slots.iter().enumerate())
244
2
    }
245
}
246

            
247
impl<T> Index<LotId> for Lots<T> {
248
    type Output = T;
249

            
250
    #[inline]
251
57
    fn index(&self, id: LotId) -> &Self::Output {
252
57
        self.get(id).expect("id is not valid")
253
57
    }
254
}
255

            
256
impl<T> IndexMut<LotId> for Lots<T> {
257
    #[inline]
258
2
    fn index_mut(&mut self, id: LotId) -> &mut Self::Output {
259
2
        self.get_mut(id).expect("id is not valid")
260
2
    }
261
}
262

            
263
impl<T> FromIterator<T> for Lots<T> {
264
    #[inline]
265
2
    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
266
2
        let iter = iter.into_iter();
267
2
        let mut map = Self::with_capacity(iter.size_hint().0);
268
7
        for item in iter {
269
5
            map.push(item);
270
5
        }
271
2
        map
272
2
    }
273
}
274

            
275
#[derive(Debug, Clone, Eq, PartialEq)]
276
enum SlotData<T> {
277
    Empty { generation: Generation },
278
    Populated { generation: Generation, contents: T },
279
}
280

            
281
impl<T> SlotData<T> {
282
    #[inline]
283
52
    pub const fn generation(&self) -> Generation {
284
52
        match self {
285
52
            SlotData::Empty { generation } | SlotData::Populated { generation, .. } => *generation,
286
52
        }
287
52
    }
288
}
289

            
290
/// An iterator over all values contained in a [`Lots<T>`].
291
#[derive(Clone)]
292
pub struct Iter<'a, T>(slice::Iter<'a, SlotData<T>>);
293

            
294
impl<'a, T> Iterator for Iter<'a, T> {
295
    type Item = &'a T;
296

            
297
    #[inline]
298
4
    fn next(&mut self) -> Option<Self::Item> {
299
        loop {
300
5
            match self.0.next()? {
301
1
                SlotData::Empty { .. } => {}
302
3
                SlotData::Populated { contents, .. } => return Some(contents),
303
            }
304
        }
305
4
    }
306
}
307

            
308
/// An iterator providing exclusive access to all values contained in a
309
/// [`Lots<T>`].
310
pub struct IterMut<'a, T>(slice::IterMut<'a, SlotData<T>>);
311

            
312
impl<'a, T> Iterator for IterMut<'a, T> {
313
    type Item = &'a mut T;
314

            
315
    #[inline]
316
4
    fn next(&mut self) -> Option<Self::Item> {
317
        loop {
318
5
            match self.0.next()? {
319
1
                SlotData::Empty { .. } => {}
320
3
                SlotData::Populated { contents, .. } => return Some(contents),
321
            }
322
        }
323
4
    }
324
}
325

            
326
/// An iterator over a [`Lots<T>`] that returns each contained value and its
327
/// associated [`LotId`].
328
#[derive(Clone)]
329
pub struct EntryIter<'a, T>(core::iter::Enumerate<slice::Iter<'a, SlotData<T>>>);
330

            
331
impl<'a, T> Iterator for EntryIter<'a, T> {
332
    type Item = (LotId, &'a T);
333

            
334
    #[inline]
335
7
    fn next(&mut self) -> Option<Self::Item> {
336
        loop {
337
8
            match self.0.next()? {
338
1
                (_, SlotData::Empty { .. }) => {}
339
                (
340
5
                    index,
341
5
                    SlotData::Populated {
342
5
                        generation,
343
5
                        contents,
344
5
                    },
345
5
                ) => {
346
5
                    return Some((
347
5
                        LotId::new(*generation, index).expect("stored lots have valid ids"),
348
5
                        contents,
349
5
                    ))
350
                }
351
            }
352
        }
353
7
    }
354
}
355

            
356
/// An iterator over values being remoed from a [`Lots<T>`].
357
pub struct Drain<'a, T, Filter>
358
where
359
    Filter: DrainFilter<T>,
360
{
361
    map: &'a mut Lots<T>,
362
    index: usize,
363
    filter: Filter,
364
}
365

            
366
impl<T, Filter> Iterator for Drain<'_, T, Filter>
367
where
368
    Filter: DrainFilter<T>,
369
{
370
    type Item = T;
371

            
372
    #[inline]
373
10
    fn next(&mut self) -> Option<Self::Item> {
374
        loop {
375
17
            let lot = self.map.slots.get_mut(self.index)?;
376
            if let SlotData::Populated {
377
6
                generation,
378
6
                contents,
379
11
            } = lot
380
            {
381
6
                if self.filter.filter(contents) {
382
4
                    let generation = *generation;
383
4
                    let SlotData::Populated { contents, .. } =
384
4
                        mem::replace(lot, SlotData::Empty { generation })
385
                    else {
386
                        unreachable!("just matched")
387
                    };
388
4
                    self.map.free_indicies.push(self.index);
389
4
                    return Some(contents);
390
2
                }
391
5
            }
392

            
393
7
            self.index += 1;
394
        }
395
10
    }
396
}
397

            
398
impl<'a, T, Filter> Drop for Drain<'a, T, Filter>
399
where
400
    Filter: DrainFilter<T>,
401
{
402
    #[inline]
403
4
    fn drop(&mut self) {
404
        // Exhaust the iterator on drop to ensure we fully execute the drain.
405
5
        for _ in self {}
406
4
    }
407
}
408

            
409
/// A filter for a [`Drain`] iterator.
410
pub trait DrainFilter<T> {
411
    /// Return true if the value should be removed from the collection and
412
    /// returned from the [`Drain`] iterator.
413
    fn filter(&mut self, value: &mut T) -> bool;
414
}
415

            
416
impl<T, F> DrainFilter<T> for F
417
where
418
    F: FnMut(&mut T) -> bool,
419
{
420
    #[inline]
421
8
    fn filter(&mut self, value: &mut T) -> bool {
422
8
        self(value)
423
8
    }
424
}
425

            
426
/// A [`DrainFilter`] that drains all elements from a collection.
427
pub struct DrainAll;
428

            
429
impl<T> DrainFilter<T> for DrainAll {
430
    #[inline]
431
3
    fn filter(&mut self, _value: &mut T) -> bool {
432
3
        true
433
3
    }
434
}
435

            
436
impl<'a, T> IntoIterator for &'a Lots<T> {
437
    type IntoIter = Iter<'a, T>;
438
    type Item = &'a T;
439

            
440
    #[inline]
441
1
    fn into_iter(self) -> Self::IntoIter {
442
1
        Iter(self.slots.iter())
443
1
    }
444
}
445

            
446
impl<'a, T> IntoIterator for &'a mut Lots<T> {
447
    type IntoIter = IterMut<'a, T>;
448
    type Item = &'a mut T;
449

            
450
    #[inline]
451
1
    fn into_iter(self) -> Self::IntoIter {
452
1
        IterMut(self.slots.iter_mut())
453
1
    }
454
}
455

            
456
/// An iterator that removes all values from the collection and frees the
457
/// underlying collection.
458
pub struct IntoIter<T>(vec::IntoIter<SlotData<T>>);
459

            
460
impl<T> IntoIterator for Lots<T> {
461
    type IntoIter = IntoIter<T>;
462
    type Item = T;
463

            
464
    #[inline]
465
2
    fn into_iter(self) -> Self::IntoIter {
466
2
        IntoIter(self.slots.into_iter())
467
2
    }
468
}
469

            
470
impl<T> Iterator for IntoIter<T> {
471
    type Item = T;
472

            
473
    #[inline]
474
7
    fn next(&mut self) -> Option<Self::Item> {
475
        loop {
476
10
            match self.0.next()? {
477
5
                SlotData::Populated { contents, .. } => return Some(contents),
478
3
                SlotData::Empty { .. } => {}
479
            }
480
        }
481
7
    }
482
}
483

            
484
#[test]
485
1
fn slot_data_size() {
486
1
    assert_eq!(mem::size_of::<SlotData<u16>>(), 4);
487
1
}
488

            
489
#[test]
490
1
fn basics() {
491
1
    let mut map = Lots::new();
492
1
    let first = map.push(1);
493
1
    assert_eq!(map[first], 1);
494
1
    assert_eq!(map.len(), 1);
495
1
    map[first] = 2;
496
1
    assert_eq!(map[first], 2);
497
1
    let drain = map.drain().collect::<Vec<_>>();
498
1
    assert_eq!(drain, &[2]);
499
1
    assert!(map.is_empty());
500

            
501
1
    assert!(map.get(first).is_none());
502
1
    assert!(map.get_mut(first).is_none());
503
1
    assert!(map.remove(first).is_none());
504
1
    let second = map.push(1);
505
1
    assert_eq!(map.remove(second), Some(1));
506
1
}
507

            
508
#[test]
509
1
fn slot_reuse() {
510
1
    let mut map = Lots::default();
511
1
    let first = map.push(1);
512
1
    assert_eq!(first.generation().get(), 1);
513
1
    assert_eq!(first.index(), 0);
514
1
    assert_eq!(map.get(first), Some(&1));
515
1
    assert_eq!(map.remove(first), Some(1));
516
1
    assert_eq!(map.get(first), None);
517

            
518
1
    let second = map.push(2);
519
1
    assert_eq!(second.index(), 0);
520
1
    assert_eq!(second.generation().get(), 2);
521
1
    assert_eq!(map.get(second), Some(&2));
522
1
    map.clear();
523
1

            
524
1
    let third = map.push(3);
525
1
    assert_eq!(third.index(), 0);
526
1
    assert_eq!(third.generation().get(), 3);
527
1
    assert_eq!(map.get(third), Some(&3));
528
1
}
529

            
530
#[test]
531
#[allow(clippy::explicit_iter_loop)] // this was chosen for code coverage
532
1
fn iteration() {
533
1
    let mut map = Lots::default();
534
1
    map.push(1);
535
1
    let two = map.push(2);
536
1
    map.push(3);
537
1
    map.push(4);
538
1

            
539
1
    // Create a gap for the iterator.
540
1
    map.remove(two);
541
1

            
542
1
    let values = map.iter().copied().collect::<Vec<_>>();
543
1
    assert_eq!(values, &[1, 3, 4]);
544
3
    for value in map.iter_mut() {
545
3
        *value += 1;
546
3
    }
547
1
    let values = map.into_iter().collect::<Vec<_>>();
548
1
    assert_eq!(values, &[2, 4, 5]);
549
1
}
550

            
551
#[test]
552
1
fn drain_filter() {
553
1
    let mut map = Lots::default();
554
1
    map.push(1_i32);
555
1
    map.push(2);
556
1
    map.push(3);
557
1
    map.push(4);
558
4
    let odd = map.drain_filter(|v| *v % 2 == 1).collect::<Vec<_>>();
559
1
    assert_eq!(odd, &[1, 3]);
560
1
    assert_eq!(map.into_iter().collect::<Vec<_>>(), &[2, 4]);
561
1
}
562

            
563
#[test]
564
1
fn dbg() {
565
1
    let mut map = Lots::from_iter([1, 2, 3]);
566
1
    map.remove(map.entries().last().unwrap().0);
567
1
    assert_eq!(alloc::format!("{map:?}"), "{LotId(0g1): 1, LotId(1g1): 2}");
568
1
}
569

            
570
#[test]
571
1
fn clone_and_eq() {
572
1
    let map = Lots::from_iter([1, 2]);
573
1
    let mut map2 = map.clone();
574
1
    assert_eq!(map, map2);
575
1
    map2.push(3);
576
1
    assert_ne!(map, map2);
577
1
}
578

            
579
#[test]
580
1
fn out_of_bounds_id() {
581
1
    let mut map = Lots::new();
582
1
    let bad_key = map.push(1);
583
1
    assert!(Lots::<i32>::new().remove(bad_key).is_none());
584
1
}