add draining iterator

actix-http/src/client/h1proto.rs

@@ -48,11 +48,11 @@ where
             match wrt.get_mut().split().freeze().try_into_value() {
                 Ok(value) => match head {
                     RequestHeadType::Owned(ref mut head) => {
-                        head.headers.insert(HOST, value)
+                        head.headers.insert(HOST, value);
                     }
                     RequestHeadType::Rc(_, ref mut extra_headers) => {
                         let headers = extra_headers.get_or_insert(HeaderMap::new());
-                        headers.insert(HOST, value)
+                        headers.insert(HOST, value);
                     }
                 },
                 Err(e) => log::error!("Can not set HOST header {}", e),

actix-http/src/header/map.rs

@@ -131,24 +131,6 @@ impl HeaderMap {
         self.inner.clear();
     }
 
-    /// Returns the number of single-value headers the map can hold without needing to reallocate.
-    ///
-    /// Since this is a multi-value map, the actual capacity is much larger when considering
-    /// each header name can be associated with an arbitrary number of values. The effect is that
-    /// the size of `len` may be greater than `capacity` since it counts all the values.
-    /// Conversely, [`len_keys`](Self::len_keys) will never be larger than capacity.
-    pub fn capacity(&self) -> usize {
-        self.inner.capacity()
-    }
-
-    /// Reserves capacity for at least `additional` more headers to be inserted in the map.
-    ///
-    /// The header map may reserve more space to avoid frequent reallocations. Additional capacity
-    /// only considers single-value headers.
-    pub fn reserve(&mut self, additional: usize) {
-        self.inner.reserve(additional)
-    }
-
     fn get_value(&self, key: impl AsHeaderName) -> Option<&Value> {
         match key.try_as_name().ok()? {
             Cow::Borrowed(name) => self.inner.get(name),
@@ -197,22 +179,6 @@ impl HeaderMap {
         }
     }
 
-    /// An iterator over all name-value pairs.
-    ///
-    /// Names will be yielded for each associated value. So, if a key has 3 associated values, it
-    /// will be yielded 3 times. The iteration order should be considered arbitrary.
-    pub fn iter(&self) -> Iter<'_> {
-        Iter::new(self.inner.iter())
-    }
-
-    /// An iterator over all contained header names.
-    ///
-    /// Each name will only be yielded once even if it has multiple associated values. The iteration
-    /// order should be considered arbitrary.
-    pub fn keys(&self) -> Keys<'_> {
-        Keys(self.inner.keys())
-    }
-
     /// Inserts a name-value pair into the map.
     ///
     /// If the map already contained this key, the new value is associated with the key and all
@@ -248,6 +214,51 @@ impl HeaderMap {
 
         Removed::new(value)
     }
+
+    /// Returns the number of single-value headers the map can hold without needing to reallocate.
+    ///
+    /// Since this is a multi-value map, the actual capacity is much larger when considering
+    /// each header name can be associated with an arbitrary number of values. The effect is that
+    /// the size of `len` may be greater than `capacity` since it counts all the values.
+    /// Conversely, [`len_keys`](Self::len_keys) will never be larger than capacity.
+    pub fn capacity(&self) -> usize {
+        self.inner.capacity()
+    }
+
+    /// Reserves capacity for at least `additional` more headers to be inserted in the map.
+    ///
+    /// The header map may reserve more space to avoid frequent reallocations. Additional capacity
+    /// only considers single-value headers.
+    pub fn reserve(&mut self, additional: usize) {
+        self.inner.reserve(additional)
+    }
+
+    /// An iterator over all name-value pairs.
+    ///
+    /// Names will be yielded for each associated value. So, if a key has 3 associated values, it
+    /// will be yielded 3 times. The iteration order should be considered arbitrary.
+    pub fn iter(&self) -> Iter<'_> {
+        Iter::new(self.inner.iter())
+    }
+
+    /// An iterator over all contained header names.
+    ///
+    /// Each name will only be yielded once even if it has multiple associated values. The iteration
+    /// order should be considered arbitrary.
+    pub fn keys(&self) -> Keys<'_> {
+        Keys(self.inner.keys())
+    }
+
+    /// Clears the map, returning all name-value sets as an iterator.
+    ///
+    /// Header names will only be yielded for the first value in each set. All items that are
+    /// yielded without a name and after an item with a name are associated with that same name.
+    /// The first item will always contain a name.
+    ///
+    /// Keeps the allocated memory for reuse.
+    pub fn drain(&mut self) -> Drain<'_> {
+        Drain::new(self.inner.drain())
+    }
 }
 
 /// Note that this implementation will clone a [HeaderName] for each value.
@@ -335,7 +346,7 @@ impl<'a> Iterator for GetAll<'a> {
     type Item = &'a HeaderValue;
 
     #[inline]
-    fn next(&mut self) -> Option<&'a HeaderValue> {
+    fn next(&mut self) -> Option<Self::Item> {
         let val = self.value?;
 
         match val {
@@ -382,7 +393,7 @@ impl Iterator for Removed {
     type Item = HeaderValue;
 
     #[inline]
-    fn next(&mut self) -> Option<HeaderValue> {
+    fn next(&mut self) -> Option<Self::Item> {
         self.inner.next()
     }
 }
@@ -395,14 +406,14 @@ impl<'a> Iterator for Keys<'a> {
     type Item = &'a HeaderName;
 
     #[inline]
-    fn next(&mut self) -> Option<&'a HeaderName> {
+    fn next(&mut self) -> Option<Self::Item> {
         self.0.next()
     }
 }
 
 #[derive(Debug)]
 pub struct Iter<'a> {
-    iter: hash_map::Iter<'a, HeaderName, Value>,
+    inner: hash_map::Iter<'a, HeaderName, Value>,
     multi_inner: Option<(&'a HeaderName, &'a SmallVec<[HeaderValue; 4]>)>,
     multi_idx: usize,
 }
@@ -410,7 +421,7 @@ pub struct Iter<'a> {
 impl<'a> Iter<'a> {
     fn new(iter: hash_map::Iter<'a, HeaderName, Value>) -> Self {
         Self {
-            iter,
+            inner: iter,
             multi_idx: 0,
             multi_inner: None,
         }
@@ -421,7 +432,7 @@ impl<'a> Iterator for Iter<'a> {
     type Item = (&'a HeaderName, &'a HeaderValue);
 
     #[inline]
-    fn next(&mut self) -> Option<(&'a HeaderName, &'a HeaderValue)> {
+    fn next(&mut self) -> Option<Self::Item> {
         // handle in-progress multi value lists first
         if let Some((ref name, ref mut vals)) = self.multi_inner {
             match vals.get(self.multi_idx) {
@@ -437,7 +448,7 @@ impl<'a> Iterator for Iter<'a> {
             }
         }
 
-        let (name, value) = self.iter.next()?;
+        let (name, value) = self.inner.next()?;
 
         match value {
             Value::One(ref value) => Some((name, value)),
@@ -450,6 +461,56 @@ impl<'a> Iterator for Iter<'a> {
     }
 }
 
+/// Iterator over drained name-value pairs.
+///
+/// Iterator items are `(Option<HeaderName>, HeaderValue)` to avoid cloning.
+#[derive(Debug)]
+pub struct Drain<'a> {
+    inner: hash_map::Drain<'a, HeaderName, Value>,
+    multi_inner: Option<(Option<HeaderName>, SmallVec<[HeaderValue; 4]>)>,
+    multi_inner_idx: usize,
+}
+
+impl<'a> Drain<'a> {
+    fn new(iter: hash_map::Drain<'a, HeaderName, Value>) -> Self {
+        Self {
+            inner: iter,
+            multi_inner: None,
+            multi_inner_idx: 0,
+        }
+    }
+}
+
+impl<'a> Iterator for Drain<'a> {
+    type Item = (Option<HeaderName>, HeaderValue);
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        // handle in-progress multi value iterators first
+        if let Some((ref mut name, ref mut vals)) = self.multi_inner {
+            if vals.len() > 0 {
+                // OPTIMISE: array removals
+                return Some((name.take(), vals.remove(0)));
+            } else {
+                // no more items in value iterator; reset state
+                self.multi_inner = None;
+                self.multi_inner_idx = 0;
+            }
+        }
+
+        let (name, mut value) = self.inner.next()?;
+
+        match value {
+            Value::One(value) => Some((Some(name), value)),
+            Value::Multi(ref mut vals) => {
+                // set up new multi value inner iter and recurse into it
+                self.multi_inner = Some((Some(name), mem::take(vals)));
+                self.next()
+            }
+        }
+    }
+}
+
 /// Iterator over owned name-value pairs.
 ///
 /// Implementation necessarily clones header names for each value.
@@ -472,7 +533,7 @@ impl Iterator for IntoIter {
     type Item = (HeaderName, HeaderValue);
 
     #[inline]
-    fn next(&mut self) -> Option<(HeaderName, HeaderValue)> {
+    fn next(&mut self) -> Option<Self::Item> {
         // handle in-progress multi value iterators first
         if let Some((ref name, ref mut vals)) = self.multi_inner {
             match vals.next() {
@@ -556,6 +617,33 @@ mod tests {
         assert!(pairs.contains(&(&header::COOKIE, &HeaderValue::from_static("two=2"))));
     }
 
+    #[test]
+    fn drain_iter() {
+        let mut map = HeaderMap::new();
+
+        map.append(header::COOKIE, HeaderValue::from_static("one=1"));
+        map.append(header::COOKIE, HeaderValue::from_static("two=2"));
+
+        let mut vals = vec![];
+        let mut iter = map.drain();
+
+        let (name, val) = iter.next().unwrap();
+        assert_eq!(name, Some(header::COOKIE));
+        vals.push(val);
+
+        let (name, val) = iter.next().unwrap();
+        assert!(name.is_none());
+        vals.push(val);
+
+        assert!(vals.contains(&HeaderValue::from_static("one=1")));
+        assert!(vals.contains(&HeaderValue::from_static("two=2")));
+
+        assert!(iter.next().is_none());
+        drop(iter);
+        
+        assert!(map.is_empty());
+    }
+
     #[test]
     fn entries_into_iter() {
         let mut map = HeaderMap::new();

actix-http/src/response.rs

@@ -363,7 +363,9 @@ impl ResponseBuilder {
     {
         if let Some(parts) = parts(&mut self.head, &self.err) {
             match header.try_into_header_pair() {
-                Ok((key, value)) => parts.headers.insert(key, value),
+                Ok((key, value)) => {
+                    parts.headers.insert(key, value);
+                }
                 Err(e) => self.err = Some(e.into()),
             };
         }

awc/src/frozen.rs

@@ -145,7 +145,9 @@ impl FrozenSendBuilder {
     {
         match HeaderName::try_from(key) {
             Ok(key) => match value.try_into_value() {
-                Ok(value) => self.extra_headers.insert(key, value),
+                Ok(value) => {
+                    self.extra_headers.insert(key, value);
+                }
                 Err(e) => self.err = Some(e.into()),
             },
             Err(e) => self.err = Some(e.into()),

awc/src/request.rs

@@ -159,7 +159,9 @@ impl ClientRequest {
         H: IntoHeaderPair,
     {
         match header.try_into_header_pair() {
-            Ok((key, value)) => self.head.headers.insert(key, value),
+            Ok((key, value)) => {
+                self.head.headers.insert(key, value);
+            }
             Err(e) => self.err = Some(e.into()),
         };
 
@@ -232,7 +234,9 @@ impl ClientRequest {
         <HeaderValue as TryFrom<V>>::Error: Into<HttpError>,
     {
         match HeaderValue::try_from(value) {
-            Ok(value) => self.head.headers.insert(header::CONTENT_TYPE, value),
+            Ok(value) => {
+                self.head.headers.insert(header::CONTENT_TYPE, value);
+            }
             Err(e) => self.err = Some(e.into()),
         }
         self

awc/src/sender.rs

@@ -304,7 +304,9 @@ impl RequestSender {
             RequestSender::Owned(head) => {
                 if !head.headers.contains_key(&key) {
                     match value.try_into_value() {
-                        Ok(value) => head.headers.insert(key, value),
+                        Ok(value) => {
+                            head.headers.insert(key, value);
+                        }
                         Err(e) => return Err(e.into()),
                     }
                 }