From 3d0363b46a3ccc34032c7a4ec08dc1973aff7f6b Mon Sep 17 00:00:00 2001
From: Martin Thomson <mt@lowentropy.net>
Date: Tue, 7 May 2024 15:20:03 +1000
Subject: [PATCH] More simplification/performance

---
 neqo-transport/src/recovery/mod.rs  | 20 +++++++++++---------
 neqo-transport/src/recovery/sent.rs | 16 ++++++++--------
 2 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/neqo-transport/src/recovery/mod.rs b/neqo-transport/src/recovery/mod.rs
index 431bc93256..b181bd88a0 100644
--- a/neqo-transport/src/recovery/mod.rs
+++ b/neqo-transport/src/recovery/mod.rs
@@ -237,13 +237,17 @@ impl LossRecoverySpace {
             .map_or(false, |t| now > t + (pto * n_pto))
     }
 
+    fn remove_outstanding(&mut self, count: usize) {
+        debug_assert!(self.in_flight_outstanding >= count);
+        self.in_flight_outstanding -= count;
+        if self.in_flight_outstanding == 0 {
+            qtrace!("remove_packet outstanding == 0 for space {}", self.space);
+        }
+    }
+
     fn remove_packet(&mut self, p: &SentPacket) {
         if p.ack_eliciting() {
-            debug_assert!(self.in_flight_outstanding > 0);
-            self.in_flight_outstanding -= 1;
-            if self.in_flight_outstanding == 0 {
-                qtrace!("remove_packet outstanding == 0 for space {}", self.space);
-            }
+            self.remove_outstanding(1);
         }
     }
 
@@ -291,11 +295,9 @@ impl LossRecoverySpace {
     /// We try to keep these around until a probe is sent for them, so it is
     /// important that `cd` is set to at least the current PTO time; otherwise we
     /// might remove all in-flight packets and stop sending probes.
-    #[allow(clippy::option_if_let_else)] // Hard enough to read as-is.
     fn remove_old_lost(&mut self, now: Instant, cd: Duration) {
-        for p in self.sent_packets.remove_expired(now, cd) {
-            self.remove_packet(&p);
-        }
+        let removed = self.sent_packets.remove_expired(now, cd);
+        self.remove_outstanding(removed);
     }
 
     /// Detect lost packets.
diff --git a/neqo-transport/src/recovery/sent.rs b/neqo-transport/src/recovery/sent.rs
index a52344c0c6..b6abc35b1b 100644
--- a/neqo-transport/src/recovery/sent.rs
+++ b/neqo-transport/src/recovery/sent.rs
@@ -218,13 +218,10 @@ impl SentPackets {
     }
 
     /// See `LossRecoverySpace::remove_old_lost` for details on `now` and `cd`.
-    pub fn remove_expired(
-        &mut self,
-        now: Instant,
-        cd: Duration,
-    ) -> impl Iterator<Item = SentPacket> {
+    /// Returns the number of ack-eliciting packets removed.
+    pub fn remove_expired(&mut self, now: Instant, cd: Duration) -> usize {
         let mut it = self.packets.iter();
-        // If the first item is not expired, do nothing.
+        // If the first item is not expired, do nothing (the most common case).
         if it.next().map_or(false, |(_, p)| p.expired(now, cd)) {
             // Find the index of the first unexpired packet.
             let to_remove = if let Some(first_keep) =
@@ -237,9 +234,12 @@ impl SentPackets {
                 // All packets are expired.
                 std::mem::take(&mut self.packets)
             };
-            to_remove.into_values()
+            to_remove
+                .into_values()
+                .filter(SentPacket::ack_eliciting)
+                .count()
         } else {
-            BTreeMap::new().into_values()
+            0
         }
     }
 }