Modify fairings

core/lib/src/fairing/fairings.rs

@@ -1,5 +1,3 @@
-use std::collections::HashSet;
-
 use crate::{Rocket, Request, Response, Data, Build, Orbit};
 use crate::fairing::{Fairing, Info, Kind};
 
@@ -21,14 +19,15 @@ pub struct Fairings {
 
 macro_rules! iter {
     ($_self:ident . $kind:ident) => ({
-        iter!($_self, $_self.$kind.iter()).map(|v| v.1)
+        iter!($_self, $_self.$kind.iter().copied()).map(|v| v.1)
     });
     ($_self:ident, $indices:expr) => ({
         let all_fairings = &$_self.all_fairings;
         $indices.filter_map(move |i| {
-            debug_assert!(all_fairings.get(*i).is_some());
-            let f = all_fairings.get(*i).map(|f| &**f)?;
-            Some((*i, f))
+            let i = i.clone();
+            debug_assert!(all_fairings.get(i).is_some());
+            let f = all_fairings.get(i).map(|f| &**f)?;
+            Some((i, f))
         })
     })
 }
@@ -47,10 +46,19 @@ impl Fairings {
             .chain(self.shutdown.iter())
     }
 
+    pub fn unique_active(&self) -> impl Iterator<Item = usize> {
+        let mut bitmap = vec![false; self.all_fairings.len()];
+        for i in self.active() {
+            bitmap.get_mut(*i).map(|active| *active = true);
+        }
+
+        bitmap.into_iter()
+            .enumerate()
+            .filter_map(|(i, active)| active.then_some(i))
+    }
+
     pub fn unique_set(&self) -> Vec<&dyn Fairing> {
-        iter!(self, self.active().collect::<HashSet<_>>().into_iter())
-            .map(|v| v.1)
-            .collect()
+        iter!(self, self.unique_active()).map(|v| v.1).collect()
     }
 
     pub fn add(&mut self, fairing: Box<dyn Fairing>) {
@@ -83,7 +91,7 @@ impl Fairings {
             };
 
             // Collect all of the active duplicates.
-            let mut dups: Vec<usize> = iter!(self, self.active())
+            let mut dups: Vec<usize> = iter!(self, self.unique_active())
                 .filter(|(_, f)| f.type_id() == this.type_id())
                 .map(|(i, _)| i)
                 .collect();
@@ -167,10 +175,31 @@ impl Fairings {
     }
 
     pub fn audit(&self) -> Result<(), &[Info]> {
-        match self.failures.is_empty() {
-            true => Ok(()),
-            false => Err(&self.failures)
+        match &self.failures[..] {
+            [] => Ok(()),
+            failures => Err(failures)
+        }
+    }
+
+    pub fn filter<F: Fairing>(&self) -> impl Iterator<Item = &F> {
+        iter!(self, self.unique_active())
+            .filter_map(|v| v.1.downcast_ref::<F>())
+    }
+
+    pub fn filter_mut<F: Fairing>(&mut self) -> impl Iterator<Item = &mut F> {
+        let mut bitmap = vec![false; self.all_fairings.len()];
+        for &i in self.active() {
+            let is_target = self.all_fairings.get(i)
+                .and_then(|f| f.downcast_ref::<F>())
+                .is_some();
+
+            bitmap.get_mut(i).map(|target| *target = is_target);
         }
+
+        self.all_fairings.iter_mut()
+            .enumerate()
+            .filter(move |(i, _)| *bitmap.get(*i).unwrap_or(&false))
+            .filter_map(|(_, f)| f.downcast_mut::<F>())
     }
 }
 

core/lib/src/fairing/mod.rs

@@ -425,7 +425,7 @@ pub type Result<T = Rocket<Build>, E = Rocket<Build>> = std::result::Result<T, E
 ///
 /// [request-local state]: https://rocket.rs/master/guide/state/#request-local-state
 #[crate::async_trait]
-pub trait Fairing: Send + Sync + Any + 'static {
+pub trait Fairing: Send + Sync + AsAny + 'static {
     /// Returns an [`Info`] structure containing the `name` and [`Kind`] of this
     /// fairing. The `name` can be any arbitrary string. `Kind` must be an `or`d
     /// set of `Kind` variants.
@@ -533,6 +533,11 @@ pub trait Fairing: Send + Sync + Any + 'static {
     async fn on_shutdown(&self, _rocket: &Rocket<Orbit>) { }
 }
 
+pub trait AsAny: Any {
+    fn as_any_ref(&self) -> &dyn Any;
+    fn as_any_mut(&mut self) -> &mut dyn Any;
+}
+
 #[crate::async_trait]
 impl<T: Fairing + ?Sized> Fairing for std::sync::Arc<T> {
     #[inline]
@@ -565,3 +570,18 @@ impl<T: Fairing + ?Sized> Fairing for std::sync::Arc<T> {
         (self as &T).on_shutdown(rocket).await
     }
 }
+
+impl<T: Any> AsAny for T {
+    fn as_any_ref(&self) -> &dyn Any { self }
+    fn as_any_mut(&mut self) -> &mut dyn Any { self }
+}
+
+impl dyn Fairing {
+    fn downcast_ref<T: Any>(&self) -> Option<&T> {
+        self.as_any_ref().downcast_ref()
+    }
+
+    fn downcast_mut<T: Any>(&mut self) -> Option<&mut T> {
+        self.as_any_mut().downcast_mut()
+    }
+}

core/lib/src/phase.rs

@@ -14,7 +14,8 @@ mod private {
 #[doc(hidden)]
 pub trait Stateful: private::Sealed {
     fn into_state(self) -> State;
-    fn as_state_ref(&self) -> StateRef<'_>;
+    fn as_ref(&self) -> StateRef<'_>;
+    fn as_mut(&mut self) -> StateRefMut<'_>;
 }
 
 /// A marker trait for Rocket's launch phases.
@@ -48,7 +49,8 @@ macro_rules! phase {
 
         impl Stateful for $S {
             fn into_state(self) -> State { State::$P(self) }
-            fn as_state_ref(&self) -> StateRef<'_> { StateRef::$P(self) }
+            fn as_ref(&self) -> StateRef<'_> { StateRef::$P(self) }
+            fn as_mut(&mut self) -> StateRefMut<'_> { StateRefMut::$P(self) }
         }
 
         #[doc(hidden)]
@@ -70,6 +72,9 @@ macro_rules! phases {
         #[doc(hidden)]
         pub enum StateRef<'a> { $($P(&'a $S)),* }
 
+        #[doc(hidden)]
+        pub enum StateRefMut<'a> { $($P(&'a mut $S)),* }
+
         $(phase!($(#[$o])* $P ($(#[$i])* $S) { $($fields)* });)*
     )
 }