use std::{borrow::Cow, collections::hash_map, iter, ops};
use ahash::AHashMap;
use http::header::{HeaderName, HeaderValue};
use smallvec::{smallvec, SmallVec};
use super::AsHeaderName;
#[derive(Debug, Clone, Default)]
pub struct HeaderMap {
pub(crate) inner: AHashMap<HeaderName, Value>,
}
#[derive(Debug, Clone)]
pub(crate) struct Value {
inner: SmallVec<[HeaderValue; 4]>,
}
impl Value {
fn one(val: HeaderValue) -> Self {
Self {
inner: smallvec![val],
}
}
fn first(&self) -> &HeaderValue {
&self.inner[0]
}
fn first_mut(&mut self) -> &mut HeaderValue {
&mut self.inner[0]
}
fn append(&mut self, new_val: HeaderValue) {
self.inner.push(new_val)
}
}
impl ops::Deref for Value {
type Target = SmallVec<[HeaderValue; 4]>;
fn deref(&self) -> &Self::Target {
&self.inner
}
}
impl HeaderMap {
pub fn new() -> Self {
HeaderMap::default()
}
pub fn with_capacity(capacity: usize) -> Self {
HeaderMap {
inner: AHashMap::with_capacity(capacity),
}
}
pub(crate) fn from_drain<I>(mut drain: I) -> Self
where
I: Iterator<Item = (Option<HeaderName>, HeaderValue)>,
{
let (first_name, first_value) = match drain.next() {
None => return HeaderMap::new(),
Some((name, val)) => {
let name = name.expect("drained first item had no name");
(name, val)
}
};
let (lb, ub) = drain.size_hint();
let capacity = ub.unwrap_or(lb);
let mut map = HeaderMap::with_capacity(capacity);
map.append(first_name.clone(), first_value);
let (map, _) = drain.fold((map, first_name), |(mut map, prev_name), (name, value)| {
let name = name.unwrap_or(prev_name);
map.append(name.clone(), value);
(map, name)
});
map
}
pub fn len(&self) -> usize {
self.inner
.iter()
.fold(0, |acc, (_, values)| acc + values.len())
}
pub fn len_keys(&self) -> usize {
self.inner.len()
}
pub fn is_empty(&self) -> bool {
self.inner.len() == 0
}
pub fn clear(&mut self) {
self.inner.clear();
}
fn get_value(&self, key: impl AsHeaderName) -> Option<&Value> {
match key.try_as_name(super::as_name::Seal).ok()? {
Cow::Borrowed(name) => self.inner.get(name),
Cow::Owned(name) => self.inner.get(&name),
}
}
pub fn get(&self, key: impl AsHeaderName) -> Option<&HeaderValue> {
self.get_value(key).map(Value::first)
}
pub fn get_mut(&mut self, key: impl AsHeaderName) -> Option<&mut HeaderValue> {
match key.try_as_name(super::as_name::Seal).ok()? {
Cow::Borrowed(name) => self.inner.get_mut(name).map(Value::first_mut),
Cow::Owned(name) => self.inner.get_mut(&name).map(Value::first_mut),
}
}
pub fn get_all(&self, key: impl AsHeaderName) -> std::slice::Iter<'_, HeaderValue> {
match self.get_value(key) {
Some(value) => value.iter(),
None => (&[]).iter(),
}
}
pub fn contains_key(&self, key: impl AsHeaderName) -> bool {
match key.try_as_name(super::as_name::Seal) {
Ok(Cow::Borrowed(name)) => self.inner.contains_key(name),
Ok(Cow::Owned(name)) => self.inner.contains_key(&name),
Err(_) => false,
}
}
pub fn insert(&mut self, key: HeaderName, val: HeaderValue) -> Removed {
let value = self.inner.insert(key, Value::one(val));
Removed::new(value)
}
pub fn append(&mut self, key: HeaderName, value: HeaderValue) {
match self.inner.entry(key) {
hash_map::Entry::Occupied(mut entry) => {
entry.get_mut().append(value);
}
hash_map::Entry::Vacant(entry) => {
entry.insert(Value::one(value));
}
};
}
pub fn remove(&mut self, key: impl AsHeaderName) -> Removed {
let value = match key.try_as_name(super::as_name::Seal) {
Ok(Cow::Borrowed(name)) => self.inner.remove(name),
Ok(Cow::Owned(name)) => self.inner.remove(&name),
Err(_) => None,
};
Removed::new(value)
}
pub fn capacity(&self) -> usize {
self.inner.capacity()
}
pub fn reserve(&mut self, additional: usize) {
self.inner.reserve(additional)
}
pub fn iter(&self) -> Iter<'_> {
Iter::new(self.inner.iter())
}
pub fn keys(&self) -> Keys<'_> {
Keys(self.inner.keys())
}
pub fn drain(&mut self) -> Drain<'_> {
Drain::new(self.inner.drain())
}
}
impl IntoIterator for HeaderMap {
type Item = (HeaderName, HeaderValue);
type IntoIter = IntoIter;
#[inline]
fn into_iter(self) -> Self::IntoIter {
IntoIter::new(self.inner.into_iter())
}
}
impl<'a> IntoIterator for &'a HeaderMap {
type Item = (&'a HeaderName, &'a HeaderValue);
type IntoIter = Iter<'a>;
#[inline]
fn into_iter(self) -> Self::IntoIter {
Iter::new(self.inner.iter())
}
}
impl From<http::HeaderMap> for HeaderMap {
fn from(mut map: http::HeaderMap) -> HeaderMap {
HeaderMap::from_drain(map.drain())
}
}
#[derive(Debug)]
pub struct Removed {
inner: Option<smallvec::IntoIter<[HeaderValue; 4]>>,
}
impl Removed {
fn new(value: Option<Value>) -> Self {
let inner = value.map(|value| value.inner.into_iter());
Self { inner }
}
pub fn is_empty(&self) -> bool {
match self.inner {
Some(ref iter) => iter.size_hint().0 == 0,
None => true,
}
}
}
impl Iterator for Removed {
type Item = HeaderValue;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
self.inner.as_mut()?.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
match self.inner {
Some(ref iter) => iter.size_hint(),
None => (0, None),
}
}
}
impl ExactSizeIterator for Removed {}
impl iter::FusedIterator for Removed {}
#[derive(Debug)]
pub struct Keys<'a>(hash_map::Keys<'a, HeaderName, Value>);
impl<'a> Iterator for Keys<'a> {
type Item = &'a HeaderName;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
self.0.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.0.size_hint()
}
}
impl ExactSizeIterator for Keys<'_> {}
impl iter::FusedIterator for Keys<'_> {}
#[derive(Debug)]
pub struct Iter<'a> {
inner: hash_map::Iter<'a, HeaderName, Value>,
multi_inner: Option<(&'a HeaderName, &'a SmallVec<[HeaderValue; 4]>)>,
multi_idx: usize,
}
impl<'a> Iter<'a> {
fn new(iter: hash_map::Iter<'a, HeaderName, Value>) -> Self {
Self {
inner: iter,
multi_idx: 0,
multi_inner: None,
}
}
}
impl<'a> Iterator for Iter<'a> {
type Item = (&'a HeaderName, &'a HeaderValue);
fn next(&mut self) -> Option<Self::Item> {
if let Some((name, ref mut vals)) = self.multi_inner {
match vals.get(self.multi_idx) {
Some(val) => {
self.multi_idx += 1;
return Some((name, val));
}
None => {
self.multi_idx = 0;
self.multi_inner = None;
}
}
}
let (name, value) = self.inner.next()?;
self.multi_inner = Some((name, &value.inner));
self.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(self.inner.size_hint().0, None)
}
}
impl ExactSizeIterator for Iter<'_> {}
impl iter::FusedIterator for Iter<'_> {}
#[derive(Debug)]
pub struct Drain<'a> {
inner: hash_map::Drain<'a, HeaderName, Value>,
multi_inner: Option<(Option<HeaderName>, SmallVec<[HeaderValue; 4]>)>,
multi_idx: usize,
}
impl<'a> Drain<'a> {
fn new(iter: hash_map::Drain<'a, HeaderName, Value>) -> Self {
Self {
inner: iter,
multi_inner: None,
multi_idx: 0,
}
}
}
impl<'a> Iterator for Drain<'a> {
type Item = (Option<HeaderName>, HeaderValue);
fn next(&mut self) -> Option<Self::Item> {
if let Some((ref mut name, ref mut vals)) = self.multi_inner {
if !vals.is_empty() {
return Some((name.take(), vals.remove(0)));
} else {
self.multi_inner = None;
self.multi_idx = 0;
}
}
let (name, value) = self.inner.next()?;
self.multi_inner = Some((Some(name), value.inner));
self.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(self.inner.size_hint().0, None)
}
}
impl ExactSizeIterator for Drain<'_> {}
impl iter::FusedIterator for Drain<'_> {}
#[derive(Debug)]
pub struct IntoIter {
inner: hash_map::IntoIter<HeaderName, Value>,
multi_inner: Option<(HeaderName, smallvec::IntoIter<[HeaderValue; 4]>)>,
}
impl IntoIter {
fn new(inner: hash_map::IntoIter<HeaderName, Value>) -> Self {
Self {
inner,
multi_inner: None,
}
}
}
impl Iterator for IntoIter {
type Item = (HeaderName, HeaderValue);
fn next(&mut self) -> Option<Self::Item> {
if let Some((ref name, ref mut vals)) = self.multi_inner {
match vals.next() {
Some(val) => {
return Some((name.clone(), val));
}
None => {
self.multi_inner = None;
}
}
}
let (name, value) = self.inner.next()?;
self.multi_inner = Some((name, value.inner.into_iter()));
self.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(self.inner.size_hint().0, None)
}
}
impl ExactSizeIterator for IntoIter {}
impl iter::FusedIterator for IntoIter {}
#[cfg(test)]
mod tests {
use std::iter::FusedIterator;
use http::header;
use static_assertions::assert_impl_all;
use super::*;
assert_impl_all!(HeaderMap: IntoIterator);
assert_impl_all!(Keys<'_>: Iterator, ExactSizeIterator, FusedIterator);
assert_impl_all!(std::slice::Iter<'_, HeaderValue>: Iterator, ExactSizeIterator, FusedIterator);
assert_impl_all!(Removed: Iterator, ExactSizeIterator, FusedIterator);
assert_impl_all!(Iter<'_>: Iterator, ExactSizeIterator, FusedIterator);
assert_impl_all!(IntoIter: Iterator, ExactSizeIterator, FusedIterator);
assert_impl_all!(Drain<'_>: Iterator, ExactSizeIterator, FusedIterator);
#[test]
fn create() {
let map = HeaderMap::new();
assert_eq!(map.len(), 0);
assert_eq!(map.capacity(), 0);
let map = HeaderMap::with_capacity(16);
assert_eq!(map.len(), 0);
assert!(map.capacity() >= 16);
}
#[test]
fn insert() {
let mut map = HeaderMap::new();
map.insert(header::LOCATION, HeaderValue::from_static("/test"));
assert_eq!(map.len(), 1);
}
#[test]
fn contains() {
let mut map = HeaderMap::new();
assert!(!map.contains_key(header::LOCATION));
map.insert(header::LOCATION, HeaderValue::from_static("/test"));
assert!(map.contains_key(header::LOCATION));
assert!(map.contains_key("Location"));
assert!(map.contains_key("Location".to_owned()));
assert!(map.contains_key("location"));
}
#[test]
fn entries_iter() {
let mut map = HeaderMap::new();
map.append(header::HOST, HeaderValue::from_static("duck.com"));
map.append(header::COOKIE, HeaderValue::from_static("one=1"));
map.append(header::COOKIE, HeaderValue::from_static("two=2"));
let mut iter = map.iter();
assert!(iter.next().is_some());
assert!(iter.next().is_some());
assert!(iter.next().is_some());
assert!(iter.next().is_none());
let pairs = map.iter().collect::<Vec<_>>();
assert!(pairs.contains(&(&header::HOST, &HeaderValue::from_static("duck.com"))));
assert!(pairs.contains(&(&header::COOKIE, &HeaderValue::from_static("one=1"))));
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();
map.append(header::HOST, HeaderValue::from_static("duck.com"));
map.append(header::COOKIE, HeaderValue::from_static("one=1"));
map.append(header::COOKIE, HeaderValue::from_static("two=2"));
let mut iter = map.into_iter();
assert!(iter.next().is_some());
assert!(iter.next().is_some());
assert!(iter.next().is_some());
assert!(iter.next().is_none());
}
#[test]
fn iter_and_into_iter_same_order() {
let mut map = HeaderMap::new();
map.append(header::HOST, HeaderValue::from_static("duck.com"));
map.append(header::COOKIE, HeaderValue::from_static("one=1"));
map.append(header::COOKIE, HeaderValue::from_static("two=2"));
let mut iter = map.iter();
let mut into_iter = map.clone().into_iter();
assert_eq!(iter.next().map(owned_pair), into_iter.next());
assert_eq!(iter.next().map(owned_pair), into_iter.next());
assert_eq!(iter.next().map(owned_pair), into_iter.next());
assert_eq!(iter.next().map(owned_pair), into_iter.next());
}
#[test]
fn get_all_and_remove_same_order() {
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 = map.get_all(header::COOKIE);
let mut removed = map.clone().remove(header::COOKIE);
assert_eq!(vals.next(), removed.next().as_ref());
assert_eq!(vals.next(), removed.next().as_ref());
assert_eq!(vals.next(), removed.next().as_ref());
}
#[test]
fn get_all_iteration_order_matches_insertion_order() {
let mut map = HeaderMap::new();
let mut vals = map.get_all(header::COOKIE);
assert!(vals.next().is_none());
map.append(header::COOKIE, HeaderValue::from_static("1"));
let mut vals = map.get_all(header::COOKIE);
assert_eq!(vals.next().unwrap().as_bytes(), b"1");
assert!(vals.next().is_none());
map.append(header::COOKIE, HeaderValue::from_static("2"));
let mut vals = map.get_all(header::COOKIE);
assert_eq!(vals.next().unwrap().as_bytes(), b"1");
assert_eq!(vals.next().unwrap().as_bytes(), b"2");
assert!(vals.next().is_none());
map.append(header::COOKIE, HeaderValue::from_static("3"));
map.append(header::COOKIE, HeaderValue::from_static("4"));
map.append(header::COOKIE, HeaderValue::from_static("5"));
let mut vals = map.get_all(header::COOKIE);
assert_eq!(vals.next().unwrap().as_bytes(), b"1");
assert_eq!(vals.next().unwrap().as_bytes(), b"2");
assert_eq!(vals.next().unwrap().as_bytes(), b"3");
assert_eq!(vals.next().unwrap().as_bytes(), b"4");
assert_eq!(vals.next().unwrap().as_bytes(), b"5");
assert!(vals.next().is_none());
let _ = map.insert(header::COOKIE, HeaderValue::from_static("6"));
let mut vals = map.get_all(header::COOKIE);
assert_eq!(vals.next().unwrap().as_bytes(), b"6");
assert!(vals.next().is_none());
let _ = map.insert(header::COOKIE, HeaderValue::from_static("7"));
let _ = map.insert(header::COOKIE, HeaderValue::from_static("8"));
let mut vals = map.get_all(header::COOKIE);
assert_eq!(vals.next().unwrap().as_bytes(), b"8");
assert!(vals.next().is_none());
map.append(header::COOKIE, HeaderValue::from_static("9"));
let mut vals = map.get_all(header::COOKIE);
assert_eq!(vals.next().unwrap().as_bytes(), b"8");
assert_eq!(vals.next().unwrap().as_bytes(), b"9");
assert!(vals.next().is_none());
assert!(vals.next().is_none());
}
fn owned_pair<'a>(
(name, val): (&'a HeaderName, &'a HeaderValue),
) -> (HeaderName, HeaderValue) {
(name.clone(), val.clone())
}
}