dense_map/
collection.rs

1// Copyright 2019 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5//! A collection of [`DenseMap`]s.
6//!
7//! Defines [`DenseMapCollection`], which is a generic map collection that can be
8//! keyed on [`DenseMapCollectionKey`], which is a two-level key structure.
9
10use std::num::NonZeroUsize;
11
12use crate::{DenseMap, EntryKey};
13
14/// A key that can index items in [`DenseMapCollection`].
15///
16/// A `DenseMapCollectionKey` is a key with two levels: `variant` and `id`. The
17/// number of `variant`s must be fixed and known at compile time, and is
18/// typically mapped to a number of `enum` variants (nested or not).
19pub trait DenseMapCollectionKey {
20    /// The number of variants this key supports.
21    const VARIANT_COUNT: NonZeroUsize;
22
23    /// Get the variant index for this key.
24    ///
25    /// # Panics
26    ///
27    /// Callers may assume that `get_variant` returns a value in the range `[0,
28    /// VARIANT_COUNT)`, and may panic if that assumption is violated.
29    fn get_variant(&self) -> usize;
30
31    /// Get the id index for this key.
32    fn get_id(&self) -> usize;
33}
34
35impl<O> EntryKey for O
36where
37    O: DenseMapCollectionKey,
38{
39    fn get_key_index(&self) -> usize {
40        <O as DenseMapCollectionKey>::get_id(self)
41    }
42}
43
44/// A vacant entry from a [`DenseMapCollection`].
45pub struct VacantEntry<'a, K, T> {
46    entry: crate::VacantEntry<'a, K, T>,
47    count: &'a mut usize,
48}
49
50impl<'a, K, T> VacantEntry<'a, K, T> {
51    /// Sets the value of the entry with the VacantEntry's key, and returns a
52    /// mutable reference to it.
53    pub fn insert(self, value: T) -> &'a mut T
54    where
55        K: EntryKey,
56    {
57        let Self { entry, count } = self;
58        *count += 1;
59        entry.insert(value)
60    }
61
62    /// Gets a reference to the key that would be used when inserting a value
63    /// through the `VacantEntry`.
64    pub fn key(&self) -> &K {
65        self.entry.key()
66    }
67
68    /// Take ownership of the key.
69    pub fn into_key(self) -> K {
70        self.entry.into_key()
71    }
72
73    /// Changes the key type of this `VacantEntry` to another key `X` that still
74    /// maps to the same index in a `DenseMap`.
75    ///
76    /// # Panics
77    ///
78    /// Panics if the resulting mapped key from `f` does not return the same
79    /// value for [`EntryKey::get_key_index`] as the old key did.
80    pub(crate) fn map_key<X, F>(self, f: F) -> VacantEntry<'a, X, T>
81    where
82        K: EntryKey,
83        X: EntryKey,
84        F: FnOnce(K) -> X,
85    {
86        let Self { entry, count } = self;
87        VacantEntry { entry: entry.map_key(f), count }
88    }
89}
90
91/// An occupied entry from a [`DenseMapCollection`].
92#[derive(Debug)]
93pub struct OccupiedEntry<'a, K, T> {
94    entry: crate::OccupiedEntry<'a, K, T>,
95    count: &'a mut usize,
96}
97
98impl<'a, K: EntryKey, T> OccupiedEntry<'a, K, T> {
99    /// Gets a reference to the key in the entry.
100    pub fn key(&self) -> &K {
101        self.entry.key()
102    }
103
104    /// Leaves the value in the map and produces the key.
105    pub fn into_key(self) -> K {
106        self.entry.into_key()
107    }
108
109    /// Gets a reference to the value in the entry.
110    pub fn get(&self) -> &T {
111        self.entry.get()
112    }
113
114    /// Gets a mutable reference to the value in the entry.
115    ///
116    /// If you need a reference to the `OccupiedEntry` which may outlive the
117    /// destruction of the entry value, see [`OccupiedEntry::into_mut`].
118    pub fn get_mut(&mut self) -> &mut T {
119        self.entry.get_mut()
120    }
121
122    /// Converts the `OccupiedEntry` into a mutable reference to the value in
123    /// the entry with a lifetime bound to the map itself.
124    ///
125    /// If you need multiple references to the `OccupiedEntry`, see
126    /// [`OccupiedEntry::get_mut`].
127    pub fn into_mut(self) -> &'a mut T {
128        self.entry.into_mut()
129    }
130
131    /// Sets the value of the entry, and returns the entry's old value.
132    pub fn insert(&mut self, value: T) -> T {
133        self.entry.insert(value)
134    }
135
136    /// Takes the value out of the entry, and returns it.
137    pub fn remove(self) -> T {
138        let Self { entry, count } = self;
139        *count -= 1;
140        entry.remove()
141    }
142
143    /// Changes the key type of this `OccupiedEntry` to another key `X` that
144    /// still maps to the same value.
145    ///
146    /// # Panics
147    ///
148    /// Panics if the resulting mapped key from `f` does not return the same
149    /// value for [`EntryKey::get_key_index`] as the old key did.
150    pub(crate) fn map_key<X, F>(self, f: F) -> OccupiedEntry<'a, X, T>
151    where
152        K: EntryKey,
153        X: EntryKey,
154        F: FnOnce(K) -> X,
155    {
156        let Self { entry, count } = self;
157        OccupiedEntry { entry: entry.map_key(f), count }
158    }
159}
160
161/// A view into an in-place entry in a map that can be vacant or occupied.
162pub enum Entry<'a, K, T> {
163    /// A vacant entry.
164    Vacant(VacantEntry<'a, K, T>),
165    /// An occupied entry.
166    Occupied(OccupiedEntry<'a, K, T>),
167}
168
169impl<'a, K: EntryKey, T> Entry<'a, K, T> {
170    /// Returns a reference to this entry's key.
171    pub fn key(&self) -> &K {
172        match self {
173            Entry::Occupied(e) => e.key(),
174            Entry::Vacant(e) => e.key(),
175        }
176    }
177
178    /// Ensures a value is in the entry by inserting `default` if empty, and
179    /// returns a mutable reference to the value in the entry.
180    pub fn or_insert(self, default: T) -> &'a mut T
181    where
182        K: EntryKey,
183    {
184        self.or_insert_with(|| default)
185    }
186
187    /// Ensures a value is in the entry by inserting the result of the function
188    /// `f` if empty, and returns a mutable reference to the value in the entry.
189    pub fn or_insert_with<F: FnOnce() -> T>(self, f: F) -> &'a mut T {
190        match self {
191            Entry::Occupied(e) => e.into_mut(),
192            Entry::Vacant(e) => e.insert(f()),
193        }
194    }
195
196    /// Ensures a value is in the entry by inserting the default value if empty,
197    /// and returns a mutable reference to the value in the entry.
198    pub fn or_default(self) -> &'a mut T
199    where
200        T: Default,
201        K: EntryKey,
202    {
203        self.or_insert_with(<T as Default>::default)
204    }
205
206    /// Provides in-place mutable access to an occupied entry before any
207    /// potential inserts into the map.
208    pub fn and_modify<F: FnOnce(&mut T)>(self, f: F) -> Self {
209        match self {
210            Entry::Occupied(mut e) => {
211                f(e.get_mut());
212                Entry::Occupied(e)
213            }
214            Entry::Vacant(e) => Entry::Vacant(e),
215        }
216    }
217
218    /// Removes the entry from the [`DenseMapCollection`].
219    ///
220    /// Returns [`Some`] if the entry was occupied, otherwise [`None`].
221    pub fn remove(self) -> Option<T> {
222        match self {
223            Entry::Vacant(_) => None,
224            Entry::Occupied(e) => Some(e.remove()),
225        }
226    }
227}
228
229/// An iterator wrapper used to implement ExactSizeIterator.
230///
231/// Wraps an iterator of type `I`, keeping track of the number of elements it
232/// is expected to produce.
233struct SizeAugmentedIterator<I> {
234    wrapped: I,
235    remaining: usize,
236}
237
238impl<I: Iterator> Iterator for SizeAugmentedIterator<I> {
239    type Item = I::Item;
240
241    fn next(&mut self) -> Option<Self::Item> {
242        let Self { wrapped, remaining } = self;
243        match wrapped.next() {
244            Some(v) => {
245                *remaining -= 1;
246                Some(v)
247            }
248            None => {
249                assert_eq!(remaining, &0);
250                None
251            }
252        }
253    }
254
255    fn size_hint(&self) -> (usize, Option<usize>) {
256        (self.remaining, Some(self.remaining))
257    }
258}
259
260impl<I: Iterator> ExactSizeIterator for SizeAugmentedIterator<I> {}
261
262/// A generic collection indexed by a [`DenseMapCollectionKey`].
263///
264/// `DenseMapCollection` provides the same performance guarantees as [`DenseMap`], but
265/// provides a two-level keying scheme that matches the pattern used in
266/// [`crate::DeviceDense`].
267pub struct DenseMapCollection<K: DenseMapCollectionKey, T> {
268    // TODO(brunodalbo): we define a vector container here because we can't just
269    // define a fixed array length based on an associated const in
270    // DenseMapCollectionKey. When rust issue #43408 gets resolved we can switch
271    // this to use the associated const and just have a fixed length array.
272    data: Vec<DenseMap<T>>,
273    count: usize,
274    _marker: core::marker::PhantomData<K>,
275}
276
277impl<K: DenseMapCollectionKey, T> DenseMapCollection<K, T> {
278    /// Creates a new empty `DenseMapCollection`.
279    pub fn new() -> Self {
280        let mut data = Vec::new();
281        data.resize_with(K::VARIANT_COUNT.get(), DenseMap::default);
282        Self { data, count: 0, _marker: core::marker::PhantomData }
283    }
284
285    fn get_map(&self, key: &K) -> &DenseMap<T> {
286        &self.data[key.get_variant()]
287    }
288
289    fn get_entry(&mut self, key: &K) -> Entry<'_, usize, T> {
290        let Self { data, count, _marker } = self;
291        match data[key.get_variant()].entry(key.get_id()) {
292            crate::Entry::Occupied(entry) => Entry::Occupied(OccupiedEntry { entry, count }),
293            crate::Entry::Vacant(entry) => Entry::Vacant(VacantEntry { entry, count }),
294        }
295    }
296
297    /// Returns `true` if the `DenseMapCollection` holds no items.
298    pub fn is_empty(&self) -> bool {
299        let Self { count, data: _, _marker } = self;
300        *count == 0
301    }
302
303    /// Returns a reference to the item indexed by `key`, or `None` if the `key`
304    /// doesn't exist.
305    pub fn get(&self, key: &K) -> Option<&T> {
306        self.get_map(key).get(key.get_id())
307    }
308
309    /// Returns a mutable reference to the item indexed by `key`, or `None` if
310    /// the `key` doesn't exist.
311    pub fn get_mut(&mut self, key: &K) -> Option<&mut T> {
312        match self.get_entry(key) {
313            Entry::Occupied(e) => Some(e.into_mut()),
314            Entry::Vacant(_) => None,
315        }
316    }
317
318    /// Removes item indexed by `key` from the container.
319    ///
320    /// Returns the removed item if it exists, or `None` otherwise.
321    pub fn remove(&mut self, key: &K) -> Option<T> {
322        match self.get_entry(key) {
323            Entry::Occupied(e) => Some(e.remove()),
324            Entry::Vacant(_) => None,
325        }
326    }
327
328    /// Inserts `item` at `key`.
329    ///
330    /// If the [`DenseMapCollection`] already contained an item indexed by `key`,
331    /// `insert` returns it, or `None` otherwise.
332    pub fn insert(&mut self, key: &K, item: T) -> Option<T> {
333        match self.get_entry(key) {
334            Entry::Occupied(mut e) => Some(e.insert(item)),
335            Entry::Vacant(e) => {
336                let _: &mut T = e.insert(item);
337                None
338            }
339        }
340    }
341
342    /// Creates an iterator over the containing items.
343    pub fn iter(&self) -> impl ExactSizeIterator<Item = &T> {
344        let Self { data, count, _marker } = self;
345        SizeAugmentedIterator {
346            wrapped: data.iter().flat_map(|m| m.key_ordered_iter()).map(|(_, v)| v),
347            remaining: *count,
348        }
349    }
350
351    /// Creates a mutable iterator over the containing items.
352    pub fn iter_mut(&mut self) -> impl ExactSizeIterator<Item = &mut T> {
353        let Self { data, count, _marker } = self;
354        SizeAugmentedIterator {
355            wrapped: data.iter_mut().flat_map(|m| m.key_ordered_iter_mut()).map(|(_, v)| v),
356            remaining: *count,
357        }
358    }
359
360    /// Creates an iterator over the maps in variant order.
361    pub fn iter_maps(&self) -> impl Iterator<Item = &DenseMap<T>> {
362        let Self { data, count: _, _marker } = self;
363        data.iter()
364    }
365
366    /// Gets the given key's corresponding entry in the map for in-place
367    /// manipulation.
368    pub fn entry(&mut self, key: K) -> Entry<'_, K, T> {
369        match self.get_entry(&key) {
370            Entry::Occupied(e) => Entry::Occupied(e.map_key(|_| key)),
371            Entry::Vacant(e) => Entry::Vacant(e.map_key(|_| key)),
372        }
373    }
374
375    /// Inserts a new entry, constructing a key with the provided function.
376    ///
377    /// # Panics
378    ///
379    /// The `make_key` function _must_ always construct keys of the same
380    /// variant, otherwise this method will panic.
381    pub fn push_entry(&mut self, make_key: fn(usize) -> K, value: T) -> OccupiedEntry<'_, K, T> {
382        let Self { count, data, _marker } = self;
383        let variant = make_key(0).get_variant();
384        let entry = data[variant].push_entry(value);
385        *count += 1;
386
387        let entry = entry.map_key(make_key);
388
389        let entry_variant = entry.key().get_variant();
390        assert_eq!(
391            entry_variant, variant,
392            "key variant is inconsistent; got both {variant} and {entry_variant}"
393        );
394        OccupiedEntry { entry, count }
395    }
396}
397
398impl<K: DenseMapCollectionKey, T> Default for DenseMapCollection<K, T> {
399    fn default() -> Self {
400        Self::new()
401    }
402}
403
404#[cfg(test)]
405mod tests {
406    use std::collections::HashSet;
407
408    use super::*;
409    use crate::testutil::assert_empty;
410
411    #[derive(Copy, Clone, Eq, PartialEq, Debug)]
412    enum FakeVariants {
413        A,
414        B,
415        C,
416    }
417
418    #[derive(Copy, Clone, Eq, PartialEq, Debug)]
419    struct FakeKey {
420        id: usize,
421        var: FakeVariants,
422    }
423
424    impl FakeKey {
425        const fn new(id: usize, var: FakeVariants) -> Self {
426            Self { id, var }
427        }
428    }
429
430    impl DenseMapCollectionKey for FakeKey {
431        const VARIANT_COUNT: NonZeroUsize = NonZeroUsize::new(3).unwrap();
432
433        fn get_variant(&self) -> usize {
434            match self.var {
435                FakeVariants::A => 0,
436                FakeVariants::B => 1,
437                FakeVariants::C => 2,
438            }
439        }
440
441        fn get_id(&self) -> usize {
442            self.id
443        }
444    }
445
446    type TestCollection = DenseMapCollection<FakeKey, i32>;
447
448    const KEY_A: FakeKey = FakeKey::new(0, FakeVariants::A);
449    const KEY_B: FakeKey = FakeKey::new(2, FakeVariants::B);
450    const KEY_C: FakeKey = FakeKey::new(4, FakeVariants::C);
451
452    #[test]
453    fn test_insert_and_get() {
454        let mut t = TestCollection::new();
455        let DenseMapCollection { data, count, _marker } = &t;
456        assert_empty(data[0].key_ordered_iter());
457        assert_empty(data[1].key_ordered_iter());
458        assert_empty(data[2].key_ordered_iter());
459        assert_eq!(count, &0);
460
461        assert_eq!(t.insert(&KEY_A, 1), None);
462        let DenseMapCollection { data, count, _marker } = &t;
463        assert!(!data[0].is_empty());
464        assert_eq!(count, &1);
465
466        assert_eq!(t.insert(&KEY_B, 2), None);
467        let DenseMapCollection { data, count, _marker } = &t;
468        assert!(!data[1].is_empty());
469        assert_eq!(count, &2);
470
471        assert_eq!(*t.get(&KEY_A).unwrap(), 1);
472        assert_eq!(t.get(&KEY_C), None);
473
474        *t.get_mut(&KEY_B).unwrap() = 3;
475        assert_eq!(*t.get(&KEY_B).unwrap(), 3);
476    }
477
478    #[test]
479    fn test_remove() {
480        let mut t = TestCollection::new();
481        assert_eq!(t.insert(&KEY_B, 15), None);
482        assert_eq!(t.remove(&KEY_B).unwrap(), 15);
483        let DenseMapCollection { data: _, count, _marker } = &t;
484        assert_eq!(count, &0);
485
486        assert_eq!(t.remove(&KEY_B), None);
487    }
488
489    #[test]
490    fn test_iter() {
491        let mut t = TestCollection::new();
492        assert_eq!(t.insert(&KEY_A, 15), None);
493        assert_eq!(t.insert(&KEY_B, -5), None);
494        assert_eq!(t.insert(&KEY_C, -10), None);
495        let mut c = 0;
496        let mut sum = 0;
497        for i in t.iter() {
498            c += 1;
499            sum += *i;
500        }
501        assert_eq!(c, 3);
502        assert_eq!(sum, 0);
503    }
504
505    #[test]
506    fn test_iter_len() {
507        let mut t = TestCollection::new();
508        assert_eq!(t.insert(&KEY_A, 1), None);
509        assert_eq!(t.insert(&KEY_B, 1), None);
510        assert_eq!(t.insert(&KEY_C, 1), None);
511        assert_eq!(t.iter().len(), 3);
512        assert_eq!(t.remove(&KEY_A), Some(1));
513        assert_eq!(t.iter().len(), 2);
514    }
515
516    #[test]
517    fn test_is_empty() {
518        let mut t = TestCollection::new();
519        assert!(t.is_empty());
520        assert_eq!(t.insert(&KEY_B, 15), None);
521        assert!(!t.is_empty());
522    }
523
524    #[test]
525    fn test_iter_mut() {
526        let mut t = TestCollection::new();
527        assert_eq!(t.insert(&KEY_A, 15), None);
528        assert_eq!(t.insert(&KEY_B, -5), None);
529        assert_eq!(t.insert(&KEY_C, -10), None);
530        for i in t.iter_mut() {
531            *i *= 2;
532        }
533        assert_eq!(*t.get(&KEY_A).unwrap(), 30);
534        assert_eq!(*t.get(&KEY_B).unwrap(), -10);
535        assert_eq!(*t.get(&KEY_C).unwrap(), -20);
536        assert_eq!(t.iter_mut().len(), 3);
537    }
538
539    #[test]
540    fn test_entry() {
541        let mut t = TestCollection::new();
542        assert_eq!(*t.entry(KEY_A).or_insert(2), 2);
543        assert_eq!(
544            *t.entry(KEY_A)
545                .and_modify(|v| {
546                    *v = 10;
547                })
548                .or_insert(5),
549            10
550        );
551        assert_eq!(
552            *t.entry(KEY_B)
553                .and_modify(|v| {
554                    *v = 10;
555                })
556                .or_insert(5),
557            5
558        );
559        assert_eq!(*t.entry(KEY_C).or_insert_with(|| 7), 7);
560
561        assert_eq!(*t.entry(KEY_C).key(), KEY_C);
562        assert_eq!(*t.get(&KEY_A).unwrap(), 10);
563        assert_eq!(*t.get(&KEY_B).unwrap(), 5);
564        assert_eq!(*t.get(&KEY_C).unwrap(), 7);
565    }
566
567    #[test]
568    fn push_entry_valid() {
569        let mut t = TestCollection::new();
570        assert_eq!(t.insert(&KEY_A, 0), None);
571        assert_eq!(t.insert(&KEY_B, 1), None);
572        assert_eq!(t.insert(&KEY_C, 2), None);
573
574        let make_key = |index| FakeKey { id: index, var: FakeVariants::A };
575
576        {
577            let entry = t.push_entry(make_key, 30);
578            assert_eq!(entry.key(), &FakeKey { id: 1, var: FakeVariants::A });
579            assert_eq!(entry.get(), &30);
580        }
581
582        {
583            let entry = t.push_entry(make_key, 20);
584            assert_eq!(entry.key(), &FakeKey { id: 2, var: FakeVariants::A });
585            assert_eq!(entry.get(), &20);
586        }
587
588        assert_eq!(t.iter().collect::<HashSet<_>>(), HashSet::from([&0, &1, &2, &30, &20]));
589    }
590
591    #[test]
592    #[should_panic(expected = "variant is inconsistent")]
593    fn push_entry_invalid_key_fn() {
594        let mut t = TestCollection::new();
595        assert_eq!(t.insert(&KEY_A, 0), None);
596
597        let bad_make_key = |index| FakeKey {
598            id: index,
599            var: if index % 2 == 0 { FakeVariants::A } else { FakeVariants::B },
600        };
601
602        let _ = t.push_entry(bad_make_key, 1);
603        let _ = t.push_entry(bad_make_key, 2);
604    }
605}