feat: Use EnumMap for AckTracker

neqo-transport/src/recovery/mod.rs

@@ -16,10 +16,10 @@ use std::{
     time::{Duration, Instant},
 };
 
+use enum_map::{enum_map, EnumMap};
 use neqo_common::{qdebug, qinfo, qlog::NeqoQlog, qtrace, qwarn};
 pub use sent::SentPacket;
 use sent::SentPackets;
-use smallvec::{smallvec, SmallVec};
 pub use token::{RecoveryToken, StreamRecoveryToken};
 
 use crate::{
@@ -361,66 +361,53 @@ impl LossRecoverySpace {
 
 #[derive(Debug)]
 pub struct LossRecoverySpaces {
-    /// When we have all of the loss recovery spaces, this will use a separate
-    /// allocation, but this is reduced once the handshake is done.
-    spaces: SmallVec<[LossRecoverySpace; 1]>,
+    spaces: EnumMap<PacketNumberSpace, Option<LossRecoverySpace>>,
 }
 
 impl LossRecoverySpaces {
-    const fn idx(space: PacketNumberSpace) -> usize {
-        match space {
-            PacketNumberSpace::ApplicationData => 0,
-            PacketNumberSpace::Handshake => 1,
-            PacketNumberSpace::Initial => 2,
-        }
-    }
-
     /// Drop a packet number space and return all the packets that were
     /// outstanding, so that those can be marked as lost.
     ///
     /// # Panics
     ///
     /// If the space has already been removed.
     pub fn drop_space(&mut self, space: PacketNumberSpace) -> impl IntoIterator<Item = SentPacket> {
-        let sp = match space {
-            PacketNumberSpace::Initial => self.spaces.pop(),
-            PacketNumberSpace::Handshake => {
-                let sp = self.spaces.pop();
-                self.spaces.shrink_to_fit();
-                sp
-            }
-            PacketNumberSpace::ApplicationData => panic!("discarding application space"),
-        };
-        let mut sp = sp.unwrap();
-        assert_eq!(sp.space(), space, "dropping spaces out of order");
-        sp.remove_ignored()
+        let sp = self.spaces[space].take();
+        assert_ne!(
+            space,
+            PacketNumberSpace::ApplicationData,
+            "discarding application space"
+        );
+        sp.unwrap().remove_ignored()
     }
 
     pub fn get(&self, space: PacketNumberSpace) -> Option<&LossRecoverySpace> {
-        self.spaces.get(Self::idx(space))
+        self.spaces[space].as_ref()
     }
 
     pub fn get_mut(&mut self, space: PacketNumberSpace) -> Option<&mut LossRecoverySpace> {
-        self.spaces.get_mut(Self::idx(space))
+        self.spaces[space].as_mut()
     }
 
     fn iter(&self) -> impl Iterator<Item = &LossRecoverySpace> {
-        self.spaces.iter()
+        self.spaces.iter().filter_map(|(_, recvd)| recvd.as_ref())
     }
 
     fn iter_mut(&mut self) -> impl Iterator<Item = &mut LossRecoverySpace> {
-        self.spaces.iter_mut()
+        self.spaces
+            .iter_mut()
+            .filter_map(|(_, recvd)| recvd.as_mut())
     }
 }
 
 impl Default for LossRecoverySpaces {
     fn default() -> Self {
         Self {
-            spaces: smallvec![
-                LossRecoverySpace::new(PacketNumberSpace::ApplicationData),
-                LossRecoverySpace::new(PacketNumberSpace::Handshake),
-                LossRecoverySpace::new(PacketNumberSpace::Initial),
-            ],
+            spaces: enum_map! {
+                PacketNumberSpace::Initial => Some(LossRecoverySpace::new(PacketNumberSpace::Initial)),
+                PacketNumberSpace::Handshake => Some(LossRecoverySpace::new(PacketNumberSpace::Handshake)),
+                PacketNumberSpace::ApplicationData =>Some(LossRecoverySpace::new(PacketNumberSpace::ApplicationData)),
+            },
         }
     }
 }
@@ -1382,13 +1369,6 @@ mod tests {
         lr.discard(PacketNumberSpace::ApplicationData, now());
     }
 
-    #[test]
-    #[should_panic(expected = "dropping spaces out of order")]
-    fn drop_out_of_order() {
-        let mut lr = Fixture::default();
-        lr.discard(PacketNumberSpace::Handshake, now());
-    }
-
     #[test]
     fn ack_after_drop() {
         let mut lr = Fixture::default();

neqo-transport/src/tracking.rs

@@ -13,10 +13,9 @@
     time::{Duration, Instant},
 };
 
-use enum_map::Enum;
+use enum_map::{enum_map, Enum, EnumMap};
 use neqo_common::{qdebug, qinfo, qtrace, qwarn, IpTosEcn};
 use neqo_crypto::{Epoch, TLS_EPOCH_HANDSHAKE, TLS_EPOCH_INITIAL};
-use smallvec::{smallvec, SmallVec};
 
 use crate::{
     ecn::EcnCount,
@@ -26,7 +25,6 @@
     stats::FrameStats,
 };
 
-// TODO(mt) look at enabling EnumMap for this: https://stackoverflow.com/a/44905797/1375574
 #[derive(Clone, Copy, Debug, PartialEq, PartialOrd, Ord, Eq, Enum)]
 pub enum PacketNumberSpace {
     Initial,
@@ -70,17 +68,17 @@
 
 #[derive(Clone, Copy, Default)]
 pub struct PacketNumberSpaceSet {
-    initial: bool,
-    handshake: bool,
-    application_data: bool,
+    spaces: EnumMap<PacketNumberSpace, bool>,
 }
 
 impl PacketNumberSpaceSet {
-    pub const fn all() -> Self {
+    pub fn all() -> Self {
         Self {
-            initial: true,
-            handshake: true,
-            application_data: true,
+            spaces: enum_map! {
+                PacketNumberSpace::Initial => true,
+                PacketNumberSpace::Handshake => true,
+                PacketNumberSpace::ApplicationData => true,
+            },
         }
     }
 }
@@ -89,21 +87,13 @@
     type Output = bool;
 
     fn index(&self, space: PacketNumberSpace) -> &Self::Output {
-        match space {
-            PacketNumberSpace::Initial => &self.initial,
-            PacketNumberSpace::Handshake => &self.handshake,
-            PacketNumberSpace::ApplicationData => &self.application_data,
-        }
+        &self.spaces[space]
     }
 }
 
 impl IndexMut<PacketNumberSpace> for PacketNumberSpaceSet {
     fn index_mut(&mut self, space: PacketNumberSpace) -> &mut Self::Output {
-        match space {
-            PacketNumberSpace::Initial => &mut self.initial,
-            PacketNumberSpace::Handshake => &mut self.handshake,
-            PacketNumberSpace::ApplicationData => &mut self.application_data,
-        }
+        &mut self.spaces[space]
     }
 }
 
@@ -550,34 +540,25 @@
     }
 }
 
-#[derive(Debug)]
 pub struct AckTracker {
-    /// This stores information about received packets in *reverse* order
-    /// by spaces.  Why reverse?  Because we ultimately only want to keep
-    /// `ApplicationData` and this allows us to drop other spaces easily.
-    spaces: SmallVec<[RecvdPackets; 1]>,
+    spaces: EnumMap<PacketNumberSpace, Option<RecvdPackets>>,
 }
 
 impl AckTracker {
     pub fn drop_space(&mut self, space: PacketNumberSpace) {
-        let sp = match space {
-            PacketNumberSpace::Initial => self.spaces.pop(),
-            PacketNumberSpace::Handshake => {
-                let sp = self.spaces.pop();
-                self.spaces.shrink_to_fit();
-                sp
-            }
-            PacketNumberSpace::ApplicationData => panic!("discarding application space"),
-        };
-        assert_eq!(sp.unwrap().space, space, "dropping spaces out of order");
+        assert_ne!(
+            space,
+            PacketNumberSpace::ApplicationData,
+            "discarding application space"
+        );
+        if space == PacketNumberSpace::Handshake {
+            assert!(self.spaces[PacketNumberSpace::Initial].is_none());
+        }
+        self.spaces[space].take();
     }
 
     pub fn get_mut(&mut self, space: PacketNumberSpace) -> Option<&mut RecvdPackets> {
-        self.spaces.get_mut(match space {
-            PacketNumberSpace::ApplicationData => 0,
-            PacketNumberSpace::Handshake => 1,
-            PacketNumberSpace::Initial => 2,
-        })
+        self.spaces[space].as_mut()
     }
 
     pub fn ack_freq(
@@ -602,23 +583,31 @@
 
     /// Determine the earliest time that an ACK might be needed.
     pub fn ack_time(&self, now: Instant) -> Option<Instant> {
-        for recvd in &self.spaces {
-            qtrace!("ack_time for {} = {:?}", recvd.space, recvd.ack_time());
+        #[cfg(debug_assertions)]
+        for (space, recvd) in &self.spaces {
+            if let Some(recvd) = recvd {
+                qtrace!("ack_time for {} = {:?}", space, recvd.ack_time());
+            }
         }
 
-        if self.spaces.len() == 1 {
-            self.spaces[0].ack_time()
-        } else {
-            // Ignore any time that is in the past relative to `now`.
-            // That is something of a hack, but there are cases where we can't send ACK
-            // frames for all spaces, which can mean that one space is stuck in the past.
-            // That isn't a problem because we guarantee that earlier spaces will always
-            // be able to send ACK frames.
-            self.spaces
-                .iter()
-                .filter_map(|recvd| recvd.ack_time().filter(|t| *t > now))
-                .min()
+        if self.spaces[PacketNumberSpace::Initial].is_none()
+            && self.spaces[PacketNumberSpace::Handshake].is_none()
+        {
+            if let Some(recvd) = &self.spaces[PacketNumberSpace::ApplicationData] {
+                return recvd.ack_time();
+            }
         }
+
+        // Ignore any time that is in the past relative to `now`.
+        // That is something of a hack, but there are cases where we can't send ACK
+        // frames for all spaces, which can mean that one space is stuck in the past.
+        // That isn't a problem because we guarantee that earlier spaces will always
+        // be able to send ACK frames.
+        self.spaces
+            .values()
+            .flatten()
+            .filter_map(|recvd| recvd.ack_time().filter(|t| *t > now))
+            .min()
     }
 
     pub fn acked(&mut self, token: &AckToken) {
@@ -645,11 +634,11 @@
 impl Default for AckTracker {
     fn default() -> Self {
         Self {
-            spaces: smallvec![
-                RecvdPackets::new(PacketNumberSpace::ApplicationData),
-                RecvdPackets::new(PacketNumberSpace::Handshake),
-                RecvdPackets::new(PacketNumberSpace::Initial),
-            ],
+            spaces: enum_map! {
+                PacketNumberSpace::Initial => Some(RecvdPackets::new(PacketNumberSpace::Initial)),
+                PacketNumberSpace::Handshake => Some(RecvdPackets::new(PacketNumberSpace::Handshake)),
+                PacketNumberSpace::ApplicationData => Some(RecvdPackets::new(PacketNumberSpace::ApplicationData)),
+            },
         }
     }
 }
@@ -667,7 +656,7 @@
     };
     use crate::{
         frame::Frame,
-        packet::{PacketBuilder, PacketNumber},
+        packet::{PacketBuilder, PacketNumber, PacketType},
         stats::FrameStats,
     };
 
@@ -942,13 +931,6 @@
         tracker.drop_space(PacketNumberSpace::ApplicationData);
     }
 
-    #[test]
-    #[should_panic(expected = "dropping spaces out of order")]
-    fn drop_out_of_order() {
-        let mut tracker = AckTracker::default();
-        tracker.drop_space(PacketNumberSpace::Handshake);
-    }
-
     #[test]
     fn drop_spaces() {
         let mut tracker = AckTracker::default();
@@ -1136,4 +1118,28 @@
         assert!(copy[PacketNumberSpace::Handshake]);
         assert!(copy[PacketNumberSpace::ApplicationData]);
     }
+
+    #[test]
+    fn from_packet_type() {
+        assert_eq!(
+            PacketNumberSpace::from(PacketType::Initial),
+            PacketNumberSpace::Initial
+        );
+        assert_eq!(
+            PacketNumberSpace::from(PacketType::Handshake),
+            PacketNumberSpace::Handshake
+        );
+        assert_eq!(
+            PacketNumberSpace::from(PacketType::ZeroRtt),
+            PacketNumberSpace::ApplicationData
+        );
+        assert_eq!(
+            PacketNumberSpace::from(PacketType::Short),
+            PacketNumberSpace::ApplicationData
+        );
+        assert!(std::panic::catch_unwind(|| {
+            PacketNumberSpace::from(PacketType::VersionNegotiation)
+        })
+        .is_err());
+    }
 }