netstack3_base/data_structures/
socketmap.rs

1// Copyright 2022 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//! Defines generic data structures used to implement common application socket
6//! functionality for multiple protocols.
7//!
8//! The core of this module is the [`SocketMap`] struct. It provides a map-like
9//! API for setting and getting values while maintaining extra information about
10//! the number of values of certain types present in the map.
11
12use core::fmt::Debug;
13use core::hash::Hash;
14use core::num::NonZeroUsize;
15
16use derivative::Derivative;
17use either::Either;
18use netstack3_hashmap::{HashMap, hash_map};
19
20/// A type whose values can "shadow" other values of the type.
21///
22/// An implementation of this trait defines a relationship between values of the
23/// type. For any value `s: S`, if `t` appears in
24/// `IterShadows::iter_shadows(s)`, then `s` shadows `t`.
25///
26/// This "shadows" relationship is similar to [`PartialOrd`] in that certain
27/// propreties must hold:
28///
29/// 1. transitivity: if `s.iter_shadows()` yields `t`, and `t.iter_shadows()`
30///    yields `u`, then `s.iter_shadows()` must also yield `u`.
31/// 2. anticyclic: `s` cannot shadow itself.
32///
33/// Produces an iterator that yields all the shadows of a given value. The order
34/// of iteration is unspecified.
35pub trait IterShadows {
36    /// The iterator returned by `iter_shadows`.
37    type IterShadows: Iterator<Item = Self>;
38    /// Produces the iterator for shadow values.
39    fn iter_shadows(&self) -> Self::IterShadows;
40}
41
42/// A type whose values can be used to produce "tag" values of a different type.
43///
44/// This can be used to provide a summary value, e.g. even or odd for an
45/// integer-like type.
46pub trait Tagged<A> {
47    /// The tag type.
48    type Tag: Copy + Eq + Debug;
49
50    /// Returns the tag value for `self` at the given address.
51    ///
52    /// This function must be deterministic, such that calling `Tagged::tag` on
53    /// the same values always returns the same tag value.
54    fn tag(&self, address: &A) -> Self::Tag;
55}
56
57/// A map that stores values and summarizes tag counts.
58///
59/// This provides a similar insertion/removal API to [`HashMap`] for individual
60/// key/value pairs. Unlike a regular `HashMap`, the key type `A` is required to
61/// implement [`IterShadows`], and `V` to implement [`Tagged`].
62///
63/// Since `A` implements `IterShadows`, a given value `a : A` has zero or more
64/// shadow values. Since the shadow relationship is transitive, we call any
65/// value `v` that is reachable by following shadows of `a` one of `a`'s
66/// "ancestors", and we say `a` is a "descendant" of `v`.
67///
68/// In addition to keys and values, this map stores the number of values
69/// present in the map for all descendants of each key. These counts are
70/// separated into buckets for different tags of type `V::Tag`.
71#[derive(Derivative, Debug)]
72#[derivative(Default(bound = ""))]
73pub struct SocketMap<A: Hash + Eq, V: Tagged<A>> {
74    map: HashMap<A, MapValue<V, V::Tag>>,
75    len: usize,
76}
77
78#[derive(Derivative, Debug)]
79#[derivative(Default(bound = ""))]
80struct MapValue<V, T> {
81    value: Option<V>,
82    descendant_counts: DescendantCounts<T>,
83}
84
85#[derive(Derivative, Debug)]
86#[derivative(Default(bound = ""))]
87struct DescendantCounts<T, const INLINE_SIZE: usize = 1> {
88    /// Holds unordered (tag, count) pairs.
89    ///
90    /// [`DescendantCounts`] maintains the invariant that tags are unique. The
91    /// ordering of tags is unspecified.
92    counts: smallvec::SmallVec<[(T, NonZeroUsize); INLINE_SIZE]>,
93}
94
95/// An entry for a key in a map that has a value.
96///
97/// This type maintains the invariant that, if an `OccupiedEntry(map, a)`
98/// exists, `SocketMap::get(map, a)` is `Some(v)`, i.e. the `HashMap` that
99/// [`SocketMap`] wraps contains a [`MapValue`] whose `value` field is
100/// `Some(v)`.
101pub struct OccupiedEntry<'a, A: Hash + Eq, V: Tagged<A>>(&'a mut SocketMap<A, V>, A);
102
103/// An entry for a key in a map that does not have a value.
104///
105/// This type maintains the invariant that, if a `VacantEntry(map, a)` exists,
106/// `SocketMap::get(map, a)` is `None`. This means that in the `HashMap` that
107/// `SocketMap` wraps, either there is no value for key `a` or there is a
108/// `MapValue` whose `value` field is `None`.
109#[cfg_attr(test, derive(Debug))]
110pub struct VacantEntry<'a, A: Hash + Eq, V: Tagged<A>>(&'a mut SocketMap<A, V>, A);
111
112/// An entry in a map that can be used to manipulate the value in-place.
113#[cfg_attr(test, derive(Debug))]
114pub enum Entry<'a, A: Hash + Eq, V: Tagged<A>> {
115    // NB: Both `OccupiedEntry` and `VacantEntry` store a reference to the map
116    // and a key directly since they need access to the entire map to update
117    // descendant counts. This means that any operation on them requires an
118    // additional map lookup with the same key. Experimentation suggests the
119    // compiler will optimize this duplicate lookup out, since it is the same
120    // one done by `SocketMap::entry` to produce the `Entry` in the first place.
121    /// An occupied entry.
122    Occupied(OccupiedEntry<'a, A, V>),
123    /// A vacant entry.
124    Vacant(VacantEntry<'a, A, V>),
125}
126
127impl<A, V> SocketMap<A, V>
128where
129    A: IterShadows + Hash + Eq,
130    V: Tagged<A>,
131{
132    /// Returns the number of entries in this `SocketMap`.
133    pub fn len(&self) -> usize {
134        self.len
135    }
136
137    /// Gets a reference to the value associated with the given key, if any.
138    pub fn get(&self, key: &A) -> Option<&V> {
139        let Self { map, len: _ } = self;
140        map.get(key).and_then(|MapValue { value, descendant_counts: _ }| value.as_ref())
141    }
142
143    /// Provides an [`Entry`] for the given key for in-place manipulation.
144    ///
145    /// This is similar to the API provided by [`HashMap::entry`]. Callers can
146    /// match on the result to perform different actions depending on whether
147    /// the map has a value for the key or not.
148    pub fn entry(&mut self, key: A) -> Entry<'_, A, V> {
149        let Self { map, len: _ } = self;
150        match map.get(&key) {
151            Some(MapValue { descendant_counts: _, value: Some(_) }) => {
152                Entry::Occupied(OccupiedEntry(self, key))
153            }
154            Some(MapValue { descendant_counts: _, value: None }) | None => {
155                Entry::Vacant(VacantEntry(self, key))
156            }
157        }
158    }
159
160    /// Removes the value for the given key if there is one.
161    ///
162    /// If there is a value for key `key`, removes it and returns it. Otherwise
163    /// returns None.
164    #[cfg(test)]
165    pub fn remove(&mut self, key: &A) -> Option<V>
166    where
167        A: Clone,
168    {
169        match self.entry(key.clone()) {
170            Entry::Vacant(_) => return None,
171            Entry::Occupied(o) => Some(o.remove()),
172        }
173    }
174
175    /// Returns counts of tags for values at keys that shadow `key`.
176    ///
177    /// This is equivalent to iterating over all keys in the map, filtering for
178    /// those keys for which `key` is one of their shadows, then calling
179    /// [`Tagged::tag`] on the value for each of those keys, and then computing
180    /// the number of occurrences for each tag.
181    pub fn descendant_counts(
182        &self,
183        key: &A,
184    ) -> impl ExactSizeIterator<Item = &'_ (V::Tag, NonZeroUsize)> {
185        let Self { map, len: _ } = self;
186        map.get(key)
187            .map(|MapValue { value: _, descendant_counts }| {
188                Either::Left(descendant_counts.into_iter())
189            })
190            .unwrap_or(Either::Right(core::iter::empty()))
191    }
192
193    /// Returns an iterator over the keys and values in the map.
194    pub fn iter(&self) -> impl Iterator<Item = (&'_ A, &'_ V)> {
195        let Self { map, len: _ } = self;
196        map.iter().filter_map(|(a, MapValue { value, descendant_counts: _ })| {
197            value.as_ref().map(|v| (a, v))
198        })
199    }
200
201    fn increment_descendant_counts(
202        map: &mut HashMap<A, MapValue<V, V::Tag>>,
203        shadows: A::IterShadows,
204        tag: V::Tag,
205    ) {
206        for shadow in shadows {
207            let MapValue { descendant_counts, value: _ } = map.entry(shadow).or_default();
208            descendant_counts.increment(tag);
209        }
210    }
211
212    fn update_descendant_counts(
213        map: &mut HashMap<A, MapValue<V, V::Tag>>,
214        shadows: A::IterShadows,
215        old_tag: V::Tag,
216        new_tag: V::Tag,
217    ) {
218        if old_tag != new_tag {
219            for shadow in shadows {
220                let counts = &mut map.get_mut(&shadow).unwrap().descendant_counts;
221                counts.increment(new_tag);
222                counts.decrement(old_tag);
223            }
224        }
225    }
226
227    fn decrement_descendant_counts(
228        map: &mut HashMap<A, MapValue<V, V::Tag>>,
229        shadows: A::IterShadows,
230        old_tag: V::Tag,
231    ) {
232        for shadow in shadows {
233            let mut entry = match map.entry(shadow) {
234                hash_map::Entry::Occupied(o) => o,
235                hash_map::Entry::Vacant(_) => unreachable!(),
236            };
237            let MapValue { descendant_counts, value } = entry.get_mut();
238            descendant_counts.decrement(old_tag);
239            if descendant_counts.is_empty() && value.is_none() {
240                let _: MapValue<_, _> = entry.remove();
241            }
242        }
243    }
244}
245
246impl<'a, K: Eq + Hash + IterShadows, V: Tagged<K>> OccupiedEntry<'a, K, V> {
247    /// Gets a reference to the key for the entry.
248    pub fn key(&self) -> &K {
249        let Self(SocketMap { map: _, len: _ }, key) = self;
250        key
251    }
252
253    /// Retrieves the value referenced by this entry.
254    pub fn get(&self) -> &V {
255        let Self(SocketMap { map, len: _ }, key) = self;
256        let MapValue { descendant_counts: _, value } = map.get(key).unwrap();
257        // unwrap() call is guaranteed safe by OccupiedEntry invariant.
258        value.as_ref().unwrap()
259    }
260
261    // NB: there is no get_mut because that would allow the caller to manipulate
262    // a value without updating the descendant tag counts.
263
264    /// Runs the provided callback on the value referenced by this entry.
265    ///
266    /// Returns the result of the callback.
267    pub fn map_mut<R>(&mut self, apply: impl FnOnce(&mut V) -> R) -> R {
268        let Self(SocketMap { map, len: _ }, key) = self;
269        // unwrap() calls are guaranteed safe by OccupiedEntry invariant.
270        let MapValue { descendant_counts: _, value } = map.get_mut(key).unwrap();
271        let value = value.as_mut().unwrap();
272
273        let old_tag = value.tag(key);
274        let r = apply(value);
275        let new_tag = value.tag(key);
276        SocketMap::update_descendant_counts(map, key.iter_shadows(), old_tag, new_tag);
277        r
278    }
279
280    /// Extracts the underlying [`SocketMap`] reference backing this entry.
281    pub fn into_map(self) -> &'a mut SocketMap<K, V> {
282        let Self(socketmap, _) = self;
283        socketmap
284    }
285
286    /// Removes the value from the map and returns it.
287    pub fn remove(self) -> V {
288        let (value, _map) = self.remove_from_map();
289        value
290    }
291
292    /// Gets a reference to the backing map.
293    pub fn get_map(&self) -> &SocketMap<K, V> {
294        let Self(socketmap, _) = self;
295        socketmap
296    }
297
298    /// Removes the value from the map and returns the value and map.
299    pub fn remove_from_map(self) -> (V, &'a mut SocketMap<K, V>) {
300        let Self(socketmap, key) = self;
301        let SocketMap { map, len } = socketmap;
302        let shadows = key.iter_shadows();
303        let mut entry = match map.entry(key) {
304            hash_map::Entry::Occupied(o) => o,
305            hash_map::Entry::Vacant(_) => unreachable!("OccupiedEntry not occupied"),
306        };
307        let tag = {
308            let MapValue { descendant_counts: _, value } = entry.get();
309            // unwrap() is guaranteed safe by OccupiedEntry invariant.
310            value.as_ref().unwrap().tag(entry.key())
311        };
312
313        let MapValue { descendant_counts, value } = entry.get_mut();
314        // unwrap() is guaranteed safe by OccupiedEntry invariant.
315        let value =
316            value.take().expect("OccupiedEntry invariant violated: expected Some, found None");
317        if descendant_counts.is_empty() {
318            let _: MapValue<V, V::Tag> = entry.remove();
319        }
320        SocketMap::decrement_descendant_counts(map, shadows, tag);
321        *len -= 1;
322        (value, socketmap)
323    }
324}
325
326impl<'a, K: Eq + Hash + IterShadows, V: Tagged<K>> VacantEntry<'a, K, V> {
327    /// Inserts a value for the key referenced by this entry.
328    ///
329    /// Returns a reference to the newly-inserted value.
330    pub fn insert(self, value: V) -> OccupiedEntry<'a, K, V>
331    where
332        K: Clone,
333    {
334        let Self(socket_map, key) = self;
335        let SocketMap { map, len } = socket_map;
336        let iter_shadows = key.iter_shadows();
337        let tag = value.tag(&key);
338        *len += 1;
339        SocketMap::increment_descendant_counts(map, iter_shadows, tag);
340        let MapValue { value: map_value, descendant_counts: _ } =
341            map.entry(key.clone()).or_default();
342        assert!(map_value.replace(value).is_none());
343        OccupiedEntry(socket_map, key)
344    }
345
346    /// Extracts the underlying [`SocketMap`] reference backing this entry.
347    pub fn into_map(self) -> &'a mut SocketMap<K, V> {
348        let Self(socketmap, _) = self;
349        socketmap
350    }
351
352    /// Gets a reference to the backing map.
353    pub fn get_map(&self) -> &SocketMap<K, V> {
354        let Self(socketmap, _) = self;
355        socketmap
356    }
357
358    /// Gets the descendant counts for this entry.
359    pub fn descendant_counts(&self) -> impl ExactSizeIterator<Item = &'_ (V::Tag, NonZeroUsize)> {
360        let Self(socket_map, key) = self;
361        socket_map.descendant_counts(&key)
362    }
363}
364
365impl<'a, A: Debug + Eq + Hash, V: Tagged<A>> Debug for OccupiedEntry<'a, A, V> {
366    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
367        let Self(_socket_map, key) = self;
368        f.debug_tuple("OccupiedEntry").field(&"_").field(key).finish()
369    }
370}
371
372impl<T: Eq, const INLINE_SIZE: usize> DescendantCounts<T, INLINE_SIZE> {
373    const ONE: NonZeroUsize = NonZeroUsize::new(1).unwrap();
374
375    /// Increments the count for the given tag.
376    fn increment(&mut self, tag: T) {
377        let Self { counts } = self;
378        match counts.iter_mut().find_map(|(t, count)| (t == &tag).then_some(count)) {
379            Some(count) => *count = NonZeroUsize::new(count.get() + 1).unwrap(),
380            None => counts.push((tag, Self::ONE)),
381        }
382    }
383
384    /// Decrements the count for the given tag.
385    ///
386    /// # Panics
387    ///
388    /// Panics if there is no count for the given tag.
389    fn decrement(&mut self, tag: T) {
390        let Self { counts } = self;
391        let (index, count) = counts
392            .iter_mut()
393            .enumerate()
394            .find_map(|(i, (t, count))| (t == &tag).then_some((i, count)))
395            .unwrap();
396        if let Some(new_count) = NonZeroUsize::new(count.get() - 1) {
397            *count = new_count
398        } else {
399            let _: (T, NonZeroUsize) = counts.swap_remove(index);
400        }
401    }
402
403    fn is_empty(&self) -> bool {
404        let Self { counts } = self;
405        counts.is_empty()
406    }
407}
408
409impl<'d, T, const INLINE_SIZE: usize> IntoIterator for &'d DescendantCounts<T, INLINE_SIZE> {
410    type Item = &'d (T, NonZeroUsize);
411    type IntoIter =
412        <&'d smallvec::SmallVec<[(T, NonZeroUsize); INLINE_SIZE]> as IntoIterator>::IntoIter;
413
414    fn into_iter(self) -> Self::IntoIter {
415        let DescendantCounts { counts } = self;
416        counts.into_iter()
417    }
418}
419
420#[cfg(test)]
421mod tests {
422    use alloc::vec;
423    use alloc::vec::Vec;
424
425    use assert_matches::assert_matches;
426    use proptest::prop_assert_eq;
427    use proptest::strategy::Strategy;
428
429    use super::*;
430
431    trait AsMap {
432        type K: Hash + Eq;
433        type V;
434        fn as_map(self) -> HashMap<Self::K, Self::V>;
435    }
436
437    impl<'d, K, V, I> AsMap for I
438    where
439        K: Hash + Eq + Clone + 'd,
440        V: 'd,
441        V: Clone + Into<usize>,
442        I: Iterator<Item = &'d (K, V)>,
443    {
444        type K = K;
445        type V = usize;
446        fn as_map(self) -> HashMap<Self::K, Self::V> {
447            self.map(|(k, v)| (k.clone(), v.clone().into())).collect()
448        }
449    }
450
451    #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
452    enum Address {
453        A(u8),
454        AB(u8, char),
455        ABC(u8, char, u8),
456    }
457    use Address::*;
458
459    impl IterShadows for Address {
460        type IterShadows = <Vec<Address> as IntoIterator>::IntoIter;
461        fn iter_shadows(&self) -> Self::IterShadows {
462            match self {
463                A(_) => vec![],
464                AB(a, _) => vec![A(*a)],
465                ABC(a, b, _) => vec![AB(*a, *b), A(*a)],
466            }
467            .into_iter()
468        }
469    }
470
471    #[derive(Eq, PartialEq, Clone, Copy, Debug)]
472    struct TV<T, V>(T, V);
473
474    impl<T: Copy + Eq + Debug, V> Tagged<Address> for TV<T, V> {
475        type Tag = T;
476
477        fn tag(&self, _: &Address) -> Self::Tag {
478            self.0
479        }
480    }
481
482    type TestSocketMap<T> = SocketMap<Address, TV<T, u8>>;
483
484    #[test]
485    fn insert_get_remove() {
486        let mut map = TestSocketMap::default();
487
488        assert_matches!(map.entry(ABC(1, 'c', 2)), Entry::Vacant(v) => v.insert(TV(0, 32)));
489        assert_eq!(map.get(&ABC(1, 'c', 2)), Some(&TV(0, 32)));
490
491        assert_eq!(map.remove(&ABC(1, 'c', 2)), Some(TV(0, 32)));
492        assert_eq!(map.get(&ABC(1, 'c', 2)), None);
493    }
494
495    #[test]
496    fn insert_remove_len() {
497        let mut map = TestSocketMap::default();
498        let TestSocketMap { len, map: _ } = map;
499        assert_eq!(len, 0);
500
501        assert_matches!(map.entry(ABC(1, 'c', 2)), Entry::Vacant(v) => v.insert(TV(0, 32)));
502        let TestSocketMap { len, map: _ } = map;
503        assert_eq!(len, 1);
504
505        assert_eq!(map.remove(&ABC(1, 'c', 2)), Some(TV(0, 32)));
506        let TestSocketMap { len, map: _ } = map;
507        assert_eq!(len, 0);
508    }
509
510    #[test]
511    fn entry_same_key() {
512        let mut map = TestSocketMap::default();
513
514        assert_matches!(map.entry(ABC(1, 'c', 2)), Entry::Vacant(v) => v.insert(TV(0, 32)));
515        let occupied = assert_matches!(map.entry(ABC(1, 'c', 2)), Entry::Occupied(o) => o);
516        assert_eq!(occupied.get(), &TV(0, 32));
517        let TestSocketMap { len, map: _ } = map;
518        assert_eq!(len, 1);
519    }
520
521    #[test]
522    fn multiple_insert_descendant_counts() {
523        let mut map = TestSocketMap::default();
524
525        assert_matches!(map.entry(ABC(1, 'c', 2)), Entry::Vacant(v) => v.insert(TV(1, 111)));
526        assert_matches!(map.entry(ABC(1, 'd', 2)), Entry::Vacant(v) => v.insert(TV(2, 111)));
527        assert_matches!(map.entry(AB(5, 'd')), Entry::Vacant(v) => v.insert(TV(1, 54)));
528        assert_matches!(map.entry(AB(1, 'd')),  Entry::Vacant(v) => v.insert(TV(3, 56)));
529        let TestSocketMap { len, map: _ } = map;
530        assert_eq!(len, 4);
531
532        assert_eq!(map.descendant_counts(&A(1)).as_map(), HashMap::from([(1, 1), (2, 1), (3, 1)]));
533        assert_eq!(map.descendant_counts(&AB(1, 'c')).as_map(), HashMap::from([(1, 1)]));
534        assert_eq!(map.descendant_counts(&AB(1, 'd')).as_map(), HashMap::from([(2, 1)]));
535
536        assert_eq!(map.descendant_counts(&A(5)).as_map(), HashMap::from([(1, 1)]));
537
538        assert_eq!(map.descendant_counts(&ABC(1, 'd', 2)).as_map(), HashMap::from([]));
539        assert_eq!(map.descendant_counts(&A(2)).as_map(), HashMap::from([]));
540    }
541
542    #[test]
543    fn entry_remove_no_shadows() {
544        let mut map = TestSocketMap::default();
545
546        assert_matches!(map.entry(ABC(16, 'c', 8)), Entry::Vacant(v) => v.insert(TV(3, 111)));
547
548        let entry = assert_matches!(map.entry(ABC(16, 'c', 8)), Entry::Occupied(o) => o);
549        assert_eq!(entry.remove(), TV(3, 111));
550        let TestSocketMap { map, len } = map;
551        assert_eq!(len, 0);
552        assert_eq!(map.len(), 0);
553    }
554
555    #[test]
556    fn entry_remove_with_shadows() {
557        let mut map = TestSocketMap::default();
558
559        assert_matches!(map.entry(ABC(16, 'c', 8)), Entry::Vacant(v) => v.insert(TV(2, 112)));
560        assert_matches!(map.entry(AB(16, 'c')), Entry::Vacant(v) => v.insert(TV(1, 111)));
561        assert_matches!(map.entry(A(16)), Entry::Vacant(v) => v.insert(TV(0, 110)));
562
563        let entry = assert_matches!(map.entry(AB(16, 'c')), Entry::Occupied(o) => o);
564        assert_eq!(entry.remove(), TV(1, 111));
565        let TestSocketMap { map, len } = map;
566        assert_eq!(len, 2);
567        assert_eq!(map.len(), 3);
568    }
569
570    #[test]
571    fn remove_ancestor_value() {
572        let mut map = TestSocketMap::default();
573        assert_matches!(map.entry(ABC(2, 'e', 1)), Entry::Vacant(v) => v.insert(TV(20, 100)));
574        assert_matches!(map.entry(AB(2, 'e')), Entry::Vacant(v) => v.insert(TV(20, 100)));
575        assert_eq!(map.remove(&AB(2, 'e')), Some(TV(20, 100)));
576
577        assert_eq!(map.descendant_counts(&A(2)).as_map(), HashMap::from([(20, 1)]));
578    }
579
580    fn key_strategy() -> impl Strategy<Value = Address> {
581        let a_strategy = 1..5u8;
582        let b_strategy = proptest::char::range('a', 'e');
583        let c_strategy = 1..5u8;
584        (a_strategy, proptest::option::of((b_strategy, proptest::option::of(c_strategy)))).prop_map(
585            |(a, b)| match b {
586                None => A(a),
587                Some((b, None)) => AB(a, b),
588                Some((b, Some(c))) => ABC(a, b, c),
589            },
590        )
591    }
592
593    fn value_strategy() -> impl Strategy<Value = TV<u8, u8>> {
594        (20..25u8, 100..105u8).prop_map(|(t, v)| TV(t, v))
595    }
596
597    #[derive(Debug, Copy, Clone, Eq, PartialEq)]
598    enum Operation {
599        Entry(Address, TV<u8, u8>),
600        Remove(Address),
601    }
602
603    impl Operation {
604        fn apply(
605            self,
606            socket_map: &mut TestSocketMap<u8>,
607            reference: &mut HashMap<Address, TV<u8, u8>>,
608        ) {
609            match self {
610                Operation::Entry(a, v) => match (socket_map.entry(a), reference.entry(a)) {
611                    (Entry::Occupied(mut s), hash_map::Entry::Occupied(mut h)) => {
612                        assert_eq!(s.map_mut(|value| core::mem::replace(value, v)), h.insert(v))
613                    }
614                    (Entry::Vacant(s), hash_map::Entry::Vacant(h)) => {
615                        let _: OccupiedEntry<'_, _, _> = s.insert(v);
616                        let _: &mut TV<_, _> = h.insert(v);
617                    }
618                    (Entry::Occupied(_), hash_map::Entry::Vacant(_)) => {
619                        panic!("socketmap has a value for {:?} but reference does not", a)
620                    }
621                    (Entry::Vacant(_), hash_map::Entry::Occupied(_)) => {
622                        panic!("socketmap has no value for {:?} but reference does", a)
623                    }
624                },
625                Operation::Remove(a) => assert_eq!(socket_map.remove(&a), reference.remove(&a)),
626            }
627        }
628    }
629
630    fn operation_strategy() -> impl Strategy<Value = Operation> {
631        proptest::prop_oneof!(
632            (key_strategy(), value_strategy()).prop_map(|(a, v)| Operation::Entry(a, v)),
633            key_strategy().prop_map(|a| Operation::Remove(a)),
634        )
635    }
636
637    fn validate_map(
638        map: TestSocketMap<u8>,
639        reference: HashMap<Address, TV<u8, u8>>,
640    ) -> Result<(), proptest::test_runner::TestCaseError> {
641        let map_values: HashMap<_, _> = map.iter().map(|(a, v)| (*a, *v)).collect();
642        assert_eq!(map_values, reference);
643        let TestSocketMap { len, map: _ } = map;
644        assert_eq!(len, reference.len());
645
646        let TestSocketMap { map: inner_map, len: _ } = &map;
647        for (key, entry) in inner_map {
648            let descendant_values = map
649                .iter()
650                .filter(|(k, _)| k.iter_shadows().any(|s| s == *key))
651                .map(|(_, value)| value);
652
653            // Fold values into a map from tag to count.
654            let expected_tag_counts = descendant_values.fold(HashMap::new(), |mut m, v| {
655                *m.entry(v.tag(key)).or_default() += 1;
656                m
657            });
658
659            let MapValue { descendant_counts, value: _ } = entry;
660            prop_assert_eq!(
661                expected_tag_counts,
662                descendant_counts.into_iter().as_map(),
663                "key = {:?}",
664                key
665            );
666        }
667        Ok(())
668    }
669
670    proptest::proptest! {
671        #![proptest_config(proptest::test_runner::Config {
672            // Add all failed seeds here.
673            failure_persistence: proptest_support::failed_seeds_no_std!(),
674            ..proptest::test_runner::Config::default()
675        })]
676
677        #[test]
678        fn test_arbitrary_operations(operations in proptest::collection::vec(operation_strategy(), 10)) {
679            let mut map = TestSocketMap::default();
680            let mut reference = HashMap::new();
681            for op in operations {
682                op.apply(&mut map, &mut reference);
683            }
684
685            // After all operations have completed, check invariants for
686            // SocketMap.
687            validate_map(map, reference)?;
688        }
689
690    }
691}
netstack3_base/data_structures/socketmap.rs

netstack3_base/data_structures/
socketmap.rs
