aro_core/state/

mod.rs

//! Type-map service container for dependency injection.

use std::any::{Any, TypeId, type_name};
use std::collections::HashMap;
use std::sync::Arc;

type CloneFn = Arc<dyn Fn(&dyn Any) -> Box<dyn Any + Send + Sync> + Send + Sync>;

struct Entry {
    value: Box<dyn Any + Send + Sync>,
    clone_fn: CloneFn,
}

impl Clone for Entry {
    fn clone(&self) -> Self {
        Self {
            value: (self.clone_fn)(self.value.as_ref()),
            clone_fn: Arc::clone(&self.clone_fn),
        }
    }
}

/// A type-map service container for dependency injection.
#[derive(Clone)]
pub struct AroState {
    services: HashMap<TypeId, Entry>,
}

impl std::fmt::Debug for AroState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AroState")
            .field("services", &format!("{} registered", self.services.len()))
            .finish()
    }
}

impl AroState {
    /// Create a new builder for constructing an `AroState`.
    pub fn builder() -> AroStateBuilder {
        AroStateBuilder {
            services: HashMap::new(),
        }
    }

    /// Retrieve a registered service by trait type.
    pub fn get<T: ?Sized + 'static>(&self) -> Option<Arc<T>>
    where
        Arc<T>: 'static,
    {
        self.services
            .get(&TypeId::of::<Arc<T>>())
            .and_then(|entry| entry.value.downcast_ref::<Arc<T>>())
            .cloned()
    }

    /// Merge another `AroState` into this one.
    #[must_use]
    pub fn merge(mut self, other: Self) -> Self {
        self.services.extend(other.services);
        self
    }

    /// Validate that an expected dependency is registered.
    pub fn validate_dep<T: ?Sized + 'static>(&self) -> Result<(), String>
    where
        Arc<T>: 'static,
    {
        if self.get::<T>().is_some() {
            Ok(())
        } else {
            Err(format!("missing dependency: {}", type_name::<T>()))
        }
    }
}

/// Builder for constructing an [`AroState`] with registered services.
pub struct AroStateBuilder {
    services: HashMap<TypeId, Entry>,
}

impl AroStateBuilder {
    /// Register a service as `Arc<dyn Trait>`.
    #[must_use]
    #[expect(
        clippy::expect_used,
        reason = "type mismatch in clone_fn indicates a bug in AroState internals"
    )]
    pub fn register<T: ?Sized + 'static>(mut self, service: Arc<T>) -> Self
    where
        Arc<T>: Send + Sync + Clone + 'static,
    {
        let clone_fn: CloneFn = Arc::new(|any: &dyn Any| {
            let arc = any
                .downcast_ref::<Arc<T>>()
                .expect("type mismatch in AroState clone");
            Box::new(arc.clone())
        });
        self.services.insert(
            TypeId::of::<Arc<T>>(),
            Entry {
                value: Box::new(service),
                clone_fn,
            },
        );
        self
    }

    /// Build the `AroState` from registered services.
    pub fn build(self) -> AroState {
        AroState {
            services: self.services,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    trait Greeter: Send + Sync {
        fn greet(&self) -> &'static str;
    }

    struct HelloGreeter;
    impl Greeter for HelloGreeter {
        fn greet(&self) -> &'static str {
            "hello"
        }
    }

    trait Counter: Send + Sync {
        fn count(&self) -> u32;
    }

    struct FixedCounter(u32);
    impl Counter for FixedCounter {
        fn count(&self) -> u32 {
            self.0
        }
    }

    #[test]
    fn register_and_retrieve_service() {
        let state = AroState::builder()
            .register::<dyn Greeter>(Arc::new(HelloGreeter))
            .build();
        let greeter = state.get::<dyn Greeter>().unwrap();
        assert_eq!(greeter.greet(), "hello");
    }

    #[test]
    fn missing_service_returns_none() {
        let state = AroState::builder().build();
        assert!(state.get::<dyn Greeter>().is_none());
    }

    #[test]
    fn multiple_services() {
        let state = AroState::builder()
            .register::<dyn Greeter>(Arc::new(HelloGreeter))
            .register::<dyn Counter>(Arc::new(FixedCounter(42)))
            .build();
        assert_eq!(state.get::<dyn Greeter>().unwrap().greet(), "hello");
        assert_eq!(state.get::<dyn Counter>().unwrap().count(), 42);
    }

    #[test]
    fn validate_dep_succeeds_when_registered() {
        let state = AroState::builder()
            .register::<dyn Greeter>(Arc::new(HelloGreeter))
            .build();
        assert!(state.validate_dep::<dyn Greeter>().is_ok());
    }

    #[test]
    fn validate_dep_fails_when_missing() {
        let state = AroState::builder().build();
        let err = state.validate_dep::<dyn Greeter>().unwrap_err();
        assert!(err.contains("missing dependency"));
    }

    #[test]
    fn state_is_clone() {
        let state = AroState::builder()
            .register::<dyn Greeter>(Arc::new(HelloGreeter))
            .build();
        let cloned = state.clone();
        assert_eq!(cloned.get::<dyn Greeter>().unwrap().greet(), "hello");
    }

    #[test]
    fn merge_preserves_both_types() {
        let state_a = AroState::builder()
            .register::<dyn Greeter>(Arc::new(HelloGreeter))
            .build();
        let state_b = AroState::builder()
            .register::<dyn Counter>(Arc::new(FixedCounter(7)))
            .build();

        let merged = state_a.merge(state_b);
        assert_eq!(merged.get::<dyn Greeter>().unwrap().greet(), "hello");
        assert_eq!(merged.get::<dyn Counter>().unwrap().count(), 7);
    }

    #[test]
    fn merge_overwrites_existing_type() {
        struct AltGreeter;
        impl Greeter for AltGreeter {
            fn greet(&self) -> &'static str {
                "hi"
            }
        }

        let state_a = AroState::builder()
            .register::<dyn Greeter>(Arc::new(HelloGreeter))
            .build();
        let state_b = AroState::builder()
            .register::<dyn Greeter>(Arc::new(AltGreeter))
            .build();

        let merged = state_a.merge(state_b);
        assert_eq!(merged.get::<dyn Greeter>().unwrap().greet(), "hi");
    }

    #[test]
    fn merge_preserves_unrelated_types() {
        let state_a = AroState::builder()
            .register::<dyn Greeter>(Arc::new(HelloGreeter))
            .register::<dyn Counter>(Arc::new(FixedCounter(1)))
            .build();
        let state_b = AroState::builder()
            .register::<dyn Counter>(Arc::new(FixedCounter(99)))
            .build();

        let merged = state_a.merge(state_b);
        // Greeter was only in state_a — must still be present
        assert_eq!(merged.get::<dyn Greeter>().unwrap().greet(), "hello");
        // Counter was overwritten by state_b
        assert_eq!(merged.get::<dyn Counter>().unwrap().count(), 99);
    }
}