From 2e29a3d5440f7e5b7b8fb291020aaa771babdeb4 Mon Sep 17 00:00:00 2001
From: Markus Unterwaditzer <markus@unterwaditzer.net>
Date: Mon, 22 Aug 2016 21:53:41 +0200
Subject: [PATCH] feat(headers): Headers::remove returns the header

Fix #891
---
 src/header/internals/cell.rs | 14 ++++++++++++++
 src/header/internals/item.rs |  8 ++++++++
 src/header/mod.rs            | 15 ++++++++++++---
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/src/header/internals/cell.rs b/src/header/internals/cell.rs
index fd15b1e9af..895ff54654 100644
--- a/src/header/internals/cell.rs
+++ b/src/header/internals/cell.rs
@@ -86,6 +86,20 @@ impl<V: ?Sized + Any + 'static> PtrMapCell<V> {
         }.map(|val| &mut **val)
     }
 
+    #[inline]
+    pub fn into_value(self, key: TypeId) -> Option<Box<V>> {
+        let mut map = unsafe { self.0.into_inner() };
+        match map {
+            PtrMap::Empty => None,
+            PtrMap::One(id, v) => if id == key {
+                Some(v)
+            } else {
+                None
+            },
+            PtrMap::Many(mut hm) => hm.remove(&key)
+        }
+    }
+
     #[inline]
     pub unsafe fn insert(&self, key: TypeId, val: Box<V>) {
         let mut map = &mut *self.0.get();
diff --git a/src/header/internals/item.rs b/src/header/internals/item.rs
index bf9ed44638..ab9ca9015d 100644
--- a/src/header/internals/item.rs
+++ b/src/header/internals/item.rs
@@ -82,6 +82,14 @@ impl Item {
         }
         self.typed.get_mut(tid).map(|typed| unsafe { typed.downcast_mut_unchecked() })
     }
+
+    pub fn into_typed<H: Header>(self) -> Option<H> {
+        let tid = TypeId::of::<H>();
+        match self.typed.into_value(tid) {
+            Some(val) => Some(val),
+            None => parse::<H>(self.raw.as_ref().expect("item.raw must exist")).ok()
+        }.map(|typed| unsafe { typed.downcast_unchecked() })
+    }
 }
 
 #[inline]
diff --git a/src/header/mod.rs b/src/header/mod.rs
index 8e36bf4f97..493d09f3c8 100644
--- a/src/header/mod.rs
+++ b/src/header/mod.rs
@@ -182,6 +182,11 @@ impl Header + Send + Sync {
     unsafe fn downcast_mut_unchecked<T: 'static>(&mut self) -> &mut T {
         &mut *(mem::transmute::<*mut _, (*mut (), *mut ())>(self).0 as *mut T)
     }
+
+    #[inline]
+    unsafe fn downcast_unchecked<T: 'static>(self: Box<Self>) -> T {
+        *(mem::transmute::<*mut _, (*mut (), *mut ())>(Box::into_raw(self)).0 as *mut T)
+    }
 }
 
 impl Clone for Box<Header + Send + Sync> {
@@ -320,10 +325,14 @@ impl Headers {
     }
 
     /// Removes a header from the map, if one existed.
-    /// Returns true if a header has been removed.
-    pub fn remove<H: Header>(&mut self) -> bool {
+    /// Returns the header, if one has been removed and could be parsed.
+    ///
+    /// Note that this function may return `None` even though a header was removed. If you want to
+    /// know whether a header exists, rather rely on `has`.
+    pub fn remove<H: Header>(&mut self) -> Option<H> {
         trace!("Headers.remove( {:?} )", header_name::<H>());
-        self.data.remove(&HeaderName(UniCase(Cow::Borrowed(header_name::<H>())))).is_some()
+        self.data.remove(&HeaderName(UniCase(Cow::Borrowed(header_name::<H>()))))
+            .and_then(Item::into_typed::<H>)
     }
 
     /// Returns an iterator over the header fields.