1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
use tower::layer::util::Identity;

use crate::prepare_behave::{
    effect_traits::{PrepareMiddlewareEffect, PrepareStateEffect},
    StateCollector,
};

/// [PrepareStateEffect] or [PrepareMiddlewareEffect] adding a new state type
pub struct AddState<S>(pub S);

impl<State: 'static, Service> PrepareMiddlewareEffect<Service> for AddState<State> {
    type Middleware = Identity;

    fn take(self, states: &mut StateCollector) -> Self::Middleware {
        self.take_state(states);
        Identity::new()
    }
}

impl<S: 'static> PrepareStateEffect for AddState<S> {
    fn take_state(self, states: &mut StateCollector) {
        states.insert(self.0)
    }
}

impl<S> AddState<S> {
    pub fn new(state: S) -> Self
    where
        S: Clone + Send + Sync + 'static,
    {
        Self(state)
    }
}