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),
}
}
}
#[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 {
pub fn builder() -> AroStateBuilder {
AroStateBuilder {
services: HashMap::new(),
}
}
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()
}
#[must_use]
pub fn merge(mut self, other: Self) -> Self {
self.services.extend(other.services);
self
}
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>()))
}
}
}
pub struct AroStateBuilder {
services: HashMap<TypeId, Entry>,
}
impl AroStateBuilder {
#[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
}
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);
assert_eq!(merged.get::<dyn Greeter>().unwrap().greet(), "hello");
assert_eq!(merged.get::<dyn Counter>().unwrap().count(), 99);
}
}