refactor some unidiomatic code

quinn-proto/src/config.rs

@@ -449,6 +449,10 @@ pub struct AckTimestampsConfig {
 }
 
 impl AckTimestampsConfig {
+    pub fn enabled(&self) -> bool {
+        self.max_timestamps_per_ack.is_some()
+    }
+
     /// Sets the maximum number of timestamp entries per ACK frame.
     pub fn max_timestamps_per_ack(&mut self, value: VarInt) -> &mut Self {
         self.max_timestamps_per_ack = Some(value);

quinn-proto/src/connection/mod.rs

@@ -1397,30 +1397,26 @@ impl Connection {
             }
         }
 
-        let mut timestamp_iter = if self
-            .config
-            .ack_timestamps_config
-            .max_timestamps_per_ack
-            .is_some()
-        {
-            let decoder = ack
-                .timestamp_iter(self.epoch, self.config.ack_timestamps_config.exponent.0)
-                .unwrap();
-            let mut v: tinyvec::TinyVec<[PacketTimestamp; 10]> = tinyvec::TinyVec::new();
-            decoder.for_each(|elt| v.push(elt));
-            v.reverse();
-            Some(v.into_iter().peekable())
-        } else {
-            None
-        };
+        let timestamp_iter =
+            ack.timestamps_iter(self.epoch, self.config.ack_timestamps_config.exponent.0);
+        if self.config.ack_timestamps_config.enabled() && timestamp_iter.is_some() {
+            // Safety: checked by is_some()
+            let iter = timestamp_iter.unwrap();
+            let packet_space = &mut self.spaces[space];
+            for pkt in iter {
+                if let Some(sent_packet) = packet_space.get_mut_sent_packet(pkt.packet_number) {
+                    sent_packet.time_received = Some(pkt.timestamp);
+                }
+            }
+        }
 
         if newly_acked.is_empty() {
             return Ok(());
         }
 
         let mut ack_eliciting_acked = false;
         for packet in newly_acked.elts() {
-            if let Some(mut info) = self.spaces[space].take(packet) {
+            if let Some(info) = self.spaces[space].take(packet) {
                 if let Some(acked) = info.largest_acked {
                     // Assume ACKs for all packets below the largest acknowledged in `packet` have
                     // been received. This can cause the peer to spuriously retransmit if some of
@@ -1442,25 +1438,6 @@ impl Connection {
                 // Notify ack frequency that a packet was acked, because it might contain an ACK_FREQUENCY frame
                 self.ack_frequency.on_acked(packet);
 
-                if let Some(timestamp_iter) = timestamp_iter.as_mut() {
-                    while let Some(v) = timestamp_iter.peek() {
-                        match v.packet_number.cmp(&packet) {
-                            cmp::Ordering::Less => {
-                                let _ = timestamp_iter.next();
-                            }
-                            cmp::Ordering::Equal => {
-                                // Unwrap safety is guaranteed because a value was validated
-                                // to exist using peek
-                                let ts = timestamp_iter.next().unwrap();
-                                info.time_received = Some(ts.timestamp);
-                            }
-                            cmp::Ordering::Greater => {
-                                break;
-                            }
-                        }
-                    }
-                }
-
                 self.on_packet_acked(now, packet, info);
             }
         }

quinn-proto/src/connection/spaces.rs

@@ -221,6 +221,10 @@ impl PacketSpace {
         Some(packet)
     }
 
+    pub(super) fn get_mut_sent_packet(&mut self, number: u64) -> Option<&mut SentPacket> {
+        self.sent_packets.get_mut(&number)
+    }
+
     /// Returns the number of bytes to *remove* from the connection's in-flight count
     pub(super) fn sent(&mut self, number: u64, packet: SentPacket) -> u64 {
         // Retain state for at most this many non-ACK-eliciting packets sent after the most recently

quinn-proto/src/frame.rs

@@ -364,7 +364,7 @@ impl fmt::Debug for Ack {
         ranges.push(']');
 
         let timestamp_count = self
-            .timestamp_iter(Instant::now(), 0)
+            .timestamps_iter(Instant::now(), 0)
             .map(|iter| iter.count());
 
         f.debug_struct("Ack")
@@ -432,7 +432,7 @@ impl Ack {
         self.into_iter()
     }
 
-    pub fn timestamp_iter(&self, epoch: Instant, exponent: u64) -> Option<AckTimestampIter> {
+    pub fn timestamps_iter(&self, epoch: Instant, exponent: u64) -> Option<AckTimestampIter> {
         self.timestamps
             .as_ref()
             .map(|v| AckTimestampIter::new(self.largest, epoch, exponent, &v[..]))