chore: Use EnumSet instead of a Boolean EnumMap (mozilla#2443)

* chore: Use `EnumSet` instead of a Boolean `EnumMap` For `PacketNumberSpaceSet`. May be faster? * Remove tests as suggested by @martinthomson * Fix * clippy --------- Signed-off-by: Lars Eggert <[email protected]>
larseggert · Feb 15, 2025 · 98fbc0a · 98fbc0a
1 parent ef6e6f4
commit 98fbc0a
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 111 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -31,6 +31,7 @@ rust-version = "1.76.0"
 [workspace.dependencies]
 # Checked against https://searchfox.org/mozilla-central/source/Cargo.lock 2024-11-11
 enum-map = { version = "2.7", default-features = false }
+enumset = { version = "1.1", default-features = false }
 log = { version = "0.4", default-features = false }
 qlog = { version = "0.13", default-features = false }
 quinn-udp = { version = "0.5.6", default-features = false, features = ["direct-log", "fast-apple-datapath"] }

diff --git a/neqo-http3/Cargo.toml b/neqo-http3/Cargo.toml
@@ -17,7 +17,7 @@ workspace = true
 
 [dependencies]
 # Checked against https://searchfox.org/mozilla-central/source/Cargo.lock 2024-11-11
-enumset = { version = "1.1", default-features = false }
+enumset = { workspace = true }
 log = { workspace = true }
 neqo-common = { path = "./../neqo-common" }
 neqo-crypto = { path = "./../neqo-crypto" }

diff --git a/neqo-transport/Cargo.toml b/neqo-transport/Cargo.toml
@@ -18,6 +18,7 @@ workspace = true
 [dependencies]
 # Checked against https://searchfox.org/mozilla-central/source/Cargo.lock 2024-11-11
 enum-map = { workspace = true }
+enumset = { workspace = true }
 indexmap = { version = "2.2", default-features = false } # See https://github.com/mozilla/neqo/issues/1858
 log = { workspace = true }
 neqo-common = { path = "../neqo-common" }

diff --git a/neqo-transport/src/recovery/mod.rs b/neqo-transport/src/recovery/mod.rs
@@ -21,6 +21,7 @@ use std::{
 };
 
 use enum_map::{enum_map, EnumMap};
+use enumset::enum_set;
 use neqo_common::{qdebug, qinfo, qlog::NeqoQlog, qtrace, qwarn};
 pub use sent::SentPacket;
 use sent::SentPackets;
@@ -94,7 +95,7 @@ impl SendProfile {
     #[must_use]
     pub fn new_pto(pn_space: PacketNumberSpace, mtu: usize, probe: PacketNumberSpaceSet) -> Self {
         debug_assert!(mtu > ACK_ONLY_SIZE_LIMIT);
-        debug_assert!(probe[pn_space]);
+        debug_assert!(probe.contains(pn_space));
         Self {
             limit: mtu,
             pto: Some(pn_space),
@@ -108,7 +109,7 @@ impl SendProfile {
     /// that has the PTO timer armed.
     #[must_use]
     pub fn should_probe(&self, space: PacketNumberSpace) -> bool {
-        self.probe[space]
+        self.probe.contains(space)
     }
 
     /// Determine whether an ACK-only packet should be sent for the given packet
@@ -444,7 +445,7 @@ impl PtoState {
     }
 
     pub fn new(space: PacketNumberSpace, probe: PacketNumberSpaceSet) -> Self {
-        debug_assert!(probe[space]);
+        debug_assert!(probe.contains(space));
         Self {
             space,
             count: 1,
@@ -454,7 +455,7 @@ impl PtoState {
     }
 
     pub fn pto(&mut self, space: PacketNumberSpace, probe: PacketNumberSpaceSet) {
-        debug_assert!(probe[space]);
+        debug_assert!(probe.contains(space));
         self.space = space;
         self.count += 1;
         self.packets = Self::pto_packet_count(space);
@@ -688,7 +689,7 @@ impl LossRecovery {
         // So maybe fire a PTO.
         if let Some(pto) = self.pto_time(rtt, PacketNumberSpace::ApplicationData) {
             if pto < now {
-                let probes = PacketNumberSpaceSet::from(&[PacketNumberSpace::ApplicationData]);
+                let probes = enum_set!(PacketNumberSpace::ApplicationData);
                 self.fire_pto(PacketNumberSpace::ApplicationData, probes, now);
             }
         }
@@ -827,7 +828,7 @@ impl LossRecovery {
         let mut allow_probes = PacketNumberSpaceSet::default();
         for pn_space in PacketNumberSpace::iter() {
             if let Some(t) = self.pto_time(rtt, pn_space) {
-                allow_probes[pn_space] = true;
+                allow_probes.insert(pn_space);
                 if t <= now {
                     qdebug!("[{self}] PTO timer fired for {pn_space}");
                     if let Some(space) = self.spaces.get_mut(pn_space) {

diff --git a/neqo-transport/src/tracking.rs b/neqo-transport/src/tracking.rs
@@ -9,14 +9,14 @@
 use std::{
     cmp::min,
     collections::VecDeque,
-    ops::{Index, IndexMut},
     time::{Duration, Instant},
 };
 
 use enum_map::{enum_map, Enum, EnumMap};
+use enumset::{EnumSet, EnumSetType};
 use neqo_common::{qdebug, qinfo, qtrace, qwarn, IpTosEcn};
 use neqo_crypto::Epoch;
-use strum::{EnumIter, IntoEnumIterator as _};
+use strum::EnumIter;
 
 use crate::{
     ecn,
@@ -27,7 +27,7 @@ use crate::{
     Error, Res,
 };
 
-#[derive(Clone, Copy, Debug, PartialEq, PartialOrd, Ord, Eq, Enum, EnumIter)]
+#[derive(Debug, PartialOrd, Ord, EnumSetType, Enum, EnumIter)]
 pub enum PacketNumberSpace {
     Initial,
     Handshake,
@@ -66,63 +66,7 @@ impl From<PacketType> for PacketNumberSpace {
     }
 }
 
-#[derive(Clone, Copy, Default)]
-pub struct PacketNumberSpaceSet {
-    spaces: EnumMap<PacketNumberSpace, bool>,
-}
-
-impl PacketNumberSpaceSet {
-    pub fn all() -> Self {
-        Self {
-            spaces: enum_map! {
-                PacketNumberSpace::Initial => true,
-                PacketNumberSpace::Handshake => true,
-                PacketNumberSpace::ApplicationData => true,
-            },
-        }
-    }
-}
-
-impl Index<PacketNumberSpace> for PacketNumberSpaceSet {
-    type Output = bool;
-
-    fn index(&self, index: PacketNumberSpace) -> &Self::Output {
-        &self.spaces[index]
-    }
-}
-
-impl IndexMut<PacketNumberSpace> for PacketNumberSpaceSet {
-    fn index_mut(&mut self, index: PacketNumberSpace) -> &mut Self::Output {
-        &mut self.spaces[index]
-    }
-}
-
-impl<T: AsRef<[PacketNumberSpace]>> From<T> for PacketNumberSpaceSet {
-    fn from(spaces: T) -> Self {
-        let mut v = Self::default();
-        for sp in spaces.as_ref() {
-            v[*sp] = true;
-        }
-        v
-    }
-}
-
-impl std::fmt::Debug for PacketNumberSpaceSet {
-    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
-        let mut first = true;
-        f.write_str("(")?;
-        for sp in PacketNumberSpace::iter() {
-            if self[sp] {
-                if !first {
-                    f.write_str("+")?;
-                    first = false;
-                }
-                std::fmt::Display::fmt(&sp, f)?;
-            }
-        }
-        f.write_str(")")
-    }
-}
+pub type PacketNumberSpaceSet = EnumSet<PacketNumberSpace>;
 
 impl std::fmt::Display for PacketNumberSpace {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
@@ -679,8 +623,8 @@ mod tests {
     use test_fixture::now;
 
     use super::{
-        AckTracker, Duration, Instant, PacketNumberSpace, PacketNumberSpaceSet, RecoveryToken,
-        RecvdPackets, MAX_TRACKED_RANGES,
+        AckTracker, Duration, Instant, PacketNumberSpace, RecoveryToken, RecvdPackets,
+        MAX_TRACKED_RANGES,
     };
     use crate::{
         frame::Frame,
@@ -1114,48 +1058,6 @@ mod tests {
         );
     }
 
-    #[test]
-    fn pnspaceset_default() {
-        let set = PacketNumberSpaceSet::default();
-        assert!(!set[PacketNumberSpace::Initial]);
-        assert!(!set[PacketNumberSpace::Handshake]);
-        assert!(!set[PacketNumberSpace::ApplicationData]);
-    }
-
-    #[test]
-    fn pnspaceset_from() {
-        let set = PacketNumberSpaceSet::from(&[PacketNumberSpace::Initial]);
-        assert!(set[PacketNumberSpace::Initial]);
-        assert!(!set[PacketNumberSpace::Handshake]);
-        assert!(!set[PacketNumberSpace::ApplicationData]);
-
-        let set =
-            PacketNumberSpaceSet::from(&[PacketNumberSpace::Handshake, PacketNumberSpace::Initial]);
-        assert!(set[PacketNumberSpace::Initial]);
-        assert!(set[PacketNumberSpace::Handshake]);
-        assert!(!set[PacketNumberSpace::ApplicationData]);
-
-        let set = PacketNumberSpaceSet::from(&[
-            PacketNumberSpace::ApplicationData,
-            PacketNumberSpace::ApplicationData,
-        ]);
-        assert!(!set[PacketNumberSpace::Initial]);
-        assert!(!set[PacketNumberSpace::Handshake]);
-        assert!(set[PacketNumberSpace::ApplicationData]);
-    }
-
-    #[test]
-    fn pnspaceset_copy() {
-        let set = PacketNumberSpaceSet::from(&[
-            PacketNumberSpace::Handshake,
-            PacketNumberSpace::ApplicationData,
-        ]);
-        let copy = set;
-        assert!(!copy[PacketNumberSpace::Initial]);
-        assert!(copy[PacketNumberSpace::Handshake]);
-        assert!(copy[PacketNumberSpace::ApplicationData]);
-    }
-
     #[test]
     fn from_packet_type() {
         assert_eq!(