Skip to main content

rainier_support/
extensions.rs

1//! A type-keyed bag of values — [`Extensions`].
2//!
3//! The framework needs several "attach an arbitrary value and get it back with
4//! its type intact" maps: per-request attributes (the authenticated user, the
5//! matched route, resolved route-model bindings), and the container's binding
6//! table. Rust has no reflection, so the key is the value's [`TypeId`] and the
7//! value is a `Box<dyn Any>` — which is exactly enough, because the caller
8//! always knows the type it wants back.
9
10use std::any::{Any, TypeId};
11use std::collections::HashMap;
12use std::fmt;
13
14/// A map from a type to one value of that type.
15///
16/// Only one value per type can be stored; inserting again replaces it. Wrap in
17/// a newtype when you need two values of the same underlying type with
18/// different meanings (e.g. `struct RequestId(String)` vs `struct TraceId(String)`).
19#[derive(Default)]
20pub struct Extensions {
21    map: HashMap<TypeId, Box<dyn Any + Send + Sync>>,
22}
23
24impl Extensions {
25    /// An empty bag.
26    pub fn new() -> Self {
27        Self::default()
28    }
29
30    /// Store `value`, returning whatever value of the same type it replaced.
31    pub fn insert<T: Send + Sync + 'static>(&mut self, value: T) -> Option<T> {
32        self.map
33            .insert(TypeId::of::<T>(), Box::new(value))
34            .and_then(|prev| prev.downcast().ok().map(|boxed| *boxed))
35    }
36
37    /// Borrow the stored `T`, if there is one.
38    pub fn get<T: Send + Sync + 'static>(&self) -> Option<&T> {
39        self.map.get(&TypeId::of::<T>()).and_then(|v| v.downcast_ref())
40    }
41
42    /// Mutably borrow the stored `T`, if there is one.
43    pub fn get_mut<T: Send + Sync + 'static>(&mut self) -> Option<&mut T> {
44        self.map.get_mut(&TypeId::of::<T>()).and_then(|v| v.downcast_mut())
45    }
46
47    /// Remove and return the stored `T`.
48    pub fn remove<T: Send + Sync + 'static>(&mut self) -> Option<T> {
49        self.map.remove(&TypeId::of::<T>()).and_then(|v| v.downcast().ok().map(|boxed| *boxed))
50    }
51
52    /// Whether a value of type `T` is present.
53    pub fn contains<T: Send + Sync + 'static>(&self) -> bool {
54        self.map.contains_key(&TypeId::of::<T>())
55    }
56
57    /// Borrow the stored `T`, inserting the value from `default` first if it
58    /// is missing.
59    pub fn get_or_insert_with<T: Send + Sync + 'static>(
60        &mut self,
61        default: impl FnOnce() -> T,
62    ) -> &mut T {
63        self.map
64            .entry(TypeId::of::<T>())
65            .or_insert_with(|| Box::new(default()))
66            .downcast_mut()
67            .expect("extensions entry is keyed by its own TypeId")
68    }
69
70    /// How many values are stored.
71    pub fn len(&self) -> usize {
72        self.map.len()
73    }
74
75    /// Whether the bag holds nothing.
76    pub fn is_empty(&self) -> bool {
77        self.map.is_empty()
78    }
79
80    /// Drop everything.
81    pub fn clear(&mut self) {
82        self.map.clear();
83    }
84}
85
86impl fmt::Debug for Extensions {
87    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
88        // The values are opaque `dyn Any`, so the count is all we can show.
89        f.debug_struct("Extensions").field("len", &self.map.len()).finish()
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    #[derive(Debug, PartialEq)]
98    struct UserId(u64);
99    #[derive(Debug, PartialEq)]
100    struct TraceId(String);
101
102    #[test]
103    fn round_trips_by_type() {
104        let mut ext = Extensions::new();
105        ext.insert(UserId(7));
106        ext.insert(TraceId("abc".into()));
107
108        assert_eq!(ext.get::<UserId>(), Some(&UserId(7)));
109        assert_eq!(ext.get::<TraceId>(), Some(&TraceId("abc".into())));
110        assert_eq!(ext.len(), 2);
111    }
112
113    #[test]
114    fn inserting_the_same_type_replaces_and_returns_the_old_value() {
115        let mut ext = Extensions::new();
116        ext.insert(UserId(1));
117        assert_eq!(ext.insert(UserId(2)), Some(UserId(1)));
118        assert_eq!(ext.get::<UserId>(), Some(&UserId(2)));
119    }
120
121    #[test]
122    fn missing_types_read_as_none() {
123        let ext = Extensions::new();
124        assert!(ext.get::<UserId>().is_none());
125        assert!(!ext.contains::<UserId>());
126        assert!(ext.is_empty());
127    }
128
129    #[test]
130    fn remove_hands_the_value_back() {
131        let mut ext = Extensions::new();
132        ext.insert(UserId(3));
133        assert_eq!(ext.remove::<UserId>(), Some(UserId(3)));
134        assert!(ext.get::<UserId>().is_none());
135    }
136
137    #[test]
138    fn get_or_insert_with_only_builds_once() {
139        let mut ext = Extensions::new();
140        assert_eq!(*ext.get_or_insert_with(|| UserId(1)), UserId(1));
141        assert_eq!(*ext.get_or_insert_with(|| UserId(99)), UserId(1));
142    }
143}