cat_dev/net/

ext_map.rs

//! A keeper of "extensions" to a request, this is a copy of hyper extensions.
//!
//! Hyperium/HTTP was dual licensed under MIT/APACHE-2 when we copied this on
//! 9/15/2022
//!
//! <https://github.com/hyperium/http/blob/master/LICENSE-MIT>
//! <https://github.com/hyperium/http/blob/master/LICENSE-APACHE>
//!
//! This was to keep the same interface as our axum extensions.

use std::any::{Any, TypeId};
use std::collections::HashMap;
use std::fmt;
use std::hash::{BuildHasherDefault, Hasher};

type AnyMap = HashMap<TypeId, Box<dyn Any + Send + Sync>, BuildHasherDefault<IdHasher>>;

// With TypeIds as keys, there's no need to hash them. They are already hashes
// themselves, coming from the compiler. The IdHasher just holds the u64 of
// the TypeId, and then returns it, instead of doing any bit fiddling.
#[derive(Default)]
struct IdHasher(u64);

impl Hasher for IdHasher {
	fn write(&mut self, _: &[u8]) {
		unreachable!("TypeId calls write_u64");
	}

	#[inline]
	fn write_u64(&mut self, id: u64) {
		self.0 = id;
	}

	#[inline]
	fn finish(&self) -> u64 {
		self.0
	}
}

/// A type map of protocol extensions.
///
/// `Extensions` can be used by `Request` and `Response` to store
/// extra data derived from the underlying protocol.
#[derive(Default)]
pub struct Extensions {
	// If extensions are never used, no need to carry around an empty HashMap.
	// That's 3 words. Instead, this is only 1 word.
	map: Option<Box<AnyMap>>,
}

impl Extensions {
	/// Create an empty `Extensions`.
	#[inline]
	#[must_use]
	pub const fn new() -> Extensions {
		Extensions { map: None }
	}

	/// Insert a type into this `Extensions`.
	///
	/// If a extension of this type already existed, it will
	/// be returned.
	///
	/// # Example
	///
	/// ```
	/// # use cat_dev::net::Extensions;
	/// let mut ext = Extensions::new();
	/// assert!(ext.insert(5i32).is_none());
	/// assert!(ext.insert(4u8).is_none());
	/// assert_eq!(ext.insert(9i32), Some(5i32));
	/// ```
	pub fn insert<T: Send + Sync + 'static>(&mut self, val: T) -> Option<T> {
		self.map
			.get_or_insert_with(Box::default)
			.insert(TypeId::of::<T>(), Box::new(val))
			.and_then(|boxed| {
				(boxed as Box<dyn Any + 'static>)
					.downcast()
					.ok()
					.map(|boxed| *boxed)
			})
	}

	/// Get a reference to a type previously inserted on this `Extensions`.
	///
	/// # Example
	///
	/// ```
	/// # use cat_dev::net::Extensions;
	/// let mut ext = Extensions::new();
	/// assert!(ext.get::<i32>().is_none());
	/// ext.insert(5i32);
	///
	/// assert_eq!(ext.get::<i32>(), Some(&5i32));
	/// ```
	#[must_use]
	pub fn get<T: Send + Sync + 'static>(&self) -> Option<&T> {
		self.map
			.as_ref()
			.and_then(|map| map.get(&TypeId::of::<T>()))
			.and_then(|boxed| (&**boxed as &(dyn Any + 'static)).downcast_ref())
	}

	/// Get a mutable reference to a type previously inserted on this `Extensions`.
	///
	/// # Example
	///
	/// ```
	/// # use cat_dev::net::Extensions;
	/// let mut ext = Extensions::new();
	/// ext.insert(String::from("Hello"));
	/// ext.get_mut::<String>().unwrap().push_str(" World");
	///
	/// assert_eq!(ext.get::<String>().unwrap(), "Hello World");
	/// ```
	#[must_use]
	pub fn get_mut<T: Send + Sync + 'static>(&mut self) -> Option<&mut T> {
		self.map
			.as_mut()
			.and_then(|map| map.get_mut(&TypeId::of::<T>()))
			.and_then(|boxed| (&mut **boxed as &mut (dyn Any + 'static)).downcast_mut())
	}

	/// Remove a type from this `Extensions`.
	///
	/// If a extension of this type existed, it will be returned.
	///
	/// # Example
	///
	/// ```
	/// # use cat_dev::net::Extensions;
	/// let mut ext = Extensions::new();
	/// ext.insert(5i32);
	/// assert_eq!(ext.remove::<i32>(), Some(5i32));
	/// assert!(ext.get::<i32>().is_none());
	/// ```
	#[must_use]
	pub fn remove<T: Send + Sync + 'static>(&mut self) -> Option<T> {
		self.map
			.as_mut()
			.and_then(|map| map.remove(&TypeId::of::<T>()))
			.and_then(|boxed| {
				(boxed as Box<dyn Any + 'static>)
					.downcast()
					.ok()
					.map(|boxed| *boxed)
			})
	}

	/// Clear the `Extensions` of all inserted extensions.
	///
	/// # Example
	///
	/// ```
	/// # use cat_dev::net::Extensions;
	/// let mut ext = Extensions::new();
	/// ext.insert(5i32);
	/// ext.clear();
	///
	/// assert!(ext.get::<i32>().is_none());
	/// ```
	#[inline]
	pub fn clear(&mut self) {
		if let Some(ref mut map) = self.map {
			map.clear();
		}
	}

	/// Check whether the extension set is empty or not.
	///
	/// # Example
	///
	/// ```
	/// # use cat_dev::net::Extensions;
	/// let mut ext = Extensions::new();
	/// assert!(ext.is_empty());
	/// ext.insert(5i32);
	/// assert!(!ext.is_empty());
	/// ```
	#[inline]
	#[must_use]
	pub fn is_empty(&self) -> bool {
		self.map.as_ref().is_none_or(|map| map.is_empty())
	}

	/// Get the numer of extensions available.
	///
	/// # Example
	///
	/// ```
	/// # use cat_dev::net::Extensions;
	/// let mut ext = Extensions::new();
	/// assert_eq!(ext.len(), 0);
	/// ext.insert(5i32);
	/// assert_eq!(ext.len(), 1);
	/// ```
	#[inline]
	#[must_use]
	pub fn len(&self) -> usize {
		self.map.as_ref().map_or(0, |map| map.len())
	}

	/// Extends `self` with another `Extensions`.
	///
	/// If an instance of a specific type exists in both, the one in `self` is overwritten with the
	/// one from `other`.
	///
	/// # Example
	///
	/// ```
	/// # use cat_dev::net::Extensions;
	/// let mut ext_a = Extensions::new();
	/// ext_a.insert(8u8);
	/// ext_a.insert(16u16);
	///
	/// let mut ext_b = Extensions::new();
	/// ext_b.insert(4u8);
	/// ext_b.insert("hello");
	///
	/// ext_a.extend(ext_b);
	/// assert_eq!(ext_a.len(), 3);
	/// assert_eq!(ext_a.get::<u8>(), Some(&4u8));
	/// assert_eq!(ext_a.get::<u16>(), Some(&16u16));
	/// assert_eq!(ext_a.get::<&'static str>().copied(), Some("hello"));
	/// ```
	pub fn extend(&mut self, other: Self) {
		if let Some(other) = other.map {
			if let Some(map) = &mut self.map {
				map.extend(*other);
			} else {
				self.map = Some(other);
			}
		}
	}
}

impl fmt::Debug for Extensions {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		f.debug_struct("Extensions").finish()
	}
}

#[cfg(test)]
#[test]
fn test_extensions() {
	#[derive(Debug, PartialEq)]
	struct MyType(i32);

	let mut extensions = Extensions::new();

	extensions.insert(5i32);
	extensions.insert(MyType(10));

	assert_eq!(extensions.get(), Some(&5i32));
	assert_eq!(extensions.get_mut(), Some(&mut 5i32));

	assert_eq!(extensions.remove::<i32>(), Some(5i32));
	assert!(extensions.get::<i32>().is_none());

	assert_eq!(extensions.get::<bool>(), None);
	assert_eq!(extensions.get(), Some(&MyType(10)));
}