use std::any::{type_name, Any, TypeId};
use std::cell::RefCell;
use std::collections::HashMap;
use std::sync::{Arc, Mutex, RwLock};
use rainier_support::{Error, Result};
type Shared = Arc<dyn Any + Send + Sync>;
type Factory = Arc<dyn Fn(&Container) -> Result<Shared> + Send + Sync>;
#[derive(Clone)]
enum Binding {
Transient(Factory),
Shared { factory: Factory, instance: Arc<Mutex<Option<Shared>>>, scoped: bool },
}
#[derive(Default)]
pub struct Container {
bindings: RwLock<HashMap<TypeId, Binding>>,
}
thread_local! {
static RESOLVING: RefCell<Vec<(TypeId, &'static str)>> = const { RefCell::new(Vec::new()) };
}
impl Container {
pub fn new() -> Self {
Self::default()
}
pub fn bind<T, F>(&self, factory: F)
where
T: Send + Sync + 'static,
F: Fn(&Container) -> Result<T> + Send + Sync + 'static,
{
let factory: Factory = Arc::new(move |c| Ok(Arc::new(factory(c)?) as Shared));
self.insert::<T>(Binding::Transient(factory));
}
pub fn singleton<T, F>(&self, factory: F)
where
T: Send + Sync + 'static,
F: Fn(&Container) -> Result<T> + Send + Sync + 'static,
{
self.shared::<T, F>(factory, false);
}
pub fn scoped<T, F>(&self, factory: F)
where
T: Send + Sync + 'static,
F: Fn(&Container) -> Result<T> + Send + Sync + 'static,
{
self.shared::<T, F>(factory, true);
}
fn shared<T, F>(&self, factory: F, scoped: bool)
where
T: Send + Sync + 'static,
F: Fn(&Container) -> Result<T> + Send + Sync + 'static,
{
let factory: Factory = Arc::new(move |c| Ok(Arc::new(factory(c)?) as Shared));
self.insert::<T>(Binding::Shared { factory, instance: Arc::new(Mutex::new(None)), scoped });
}
pub fn instance<T: Send + Sync + 'static>(&self, value: T) {
self.instance_arc(Arc::new(value));
}
pub fn instance_arc<T: Send + Sync + 'static>(&self, value: Arc<T>) {
let instance: Shared = value;
self.insert::<T>(Binding::Shared {
factory: Arc::new(|_| Err(Error::internal("instance binding has no factory"))),
instance: Arc::new(Mutex::new(Some(instance))),
scoped: false,
});
}
fn insert<T: 'static>(&self, binding: Binding) {
let mut bindings = self.bindings.write().expect("container lock poisoned");
bindings.insert(TypeId::of::<T>(), binding);
}
pub fn resolve<T: Send + Sync + 'static>(&self) -> Result<Arc<T>> {
let binding = {
let bindings = self.bindings.read().expect("container lock poisoned");
bindings.get(&TypeId::of::<T>()).cloned()
};
let Some(binding) = binding else {
return Err(Error::internal(format!(
"nothing is bound for `{}` — register it in a service provider before resolving",
type_name::<T>()
)));
};
let value = self.build::<T>(binding)?;
value.downcast::<T>().map_err(|_| {
Error::internal(format!("binding for `{}` produced the wrong type", type_name::<T>()))
})
}
fn build<T: Send + Sync + 'static>(&self, binding: Binding) -> Result<Shared> {
let _guard = ResolutionGuard::enter::<T>()?;
match binding {
Binding::Transient(factory) => factory(self),
Binding::Shared { factory, instance, .. } => {
let mut slot = instance.lock().expect("container lock poisoned");
if let Some(existing) = slot.clone() {
return Ok(existing);
}
let built = factory(self)?;
*slot = Some(built.clone());
Ok(built)
}
}
}
pub fn try_resolve<T: Send + Sync + 'static>(&self) -> Option<Arc<T>> {
self.resolve::<T>().ok()
}
pub fn expect_resolve<T: Send + Sync + 'static>(&self) -> Arc<T> {
match self.resolve::<T>() {
Ok(value) => value,
Err(e) => panic!("{e}"),
}
}
pub fn bound<T: 'static>(&self) -> bool {
self.bindings.read().expect("container lock poisoned").contains_key(&TypeId::of::<T>())
}
pub fn forget<T: 'static>(&self) {
self.bindings.write().expect("container lock poisoned").remove(&TypeId::of::<T>());
}
pub fn flush_scoped(&self) {
let bindings = self.bindings.read().expect("container lock poisoned");
for binding in bindings.values() {
if let Binding::Shared { instance, scoped: true, .. } = binding {
*instance.lock().expect("container lock poisoned") = None;
}
}
}
pub fn len(&self) -> usize {
self.bindings.read().expect("container lock poisoned").len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
impl std::fmt::Debug for Container {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Container").field("bindings", &self.len()).finish()
}
}
struct ResolutionGuard;
impl ResolutionGuard {
fn enter<T: 'static>() -> Result<Self> {
let id = TypeId::of::<T>();
let name = type_name::<T>();
RESOLVING.with(|stack| {
let mut stack = stack.borrow_mut();
if stack.iter().any(|(seen, _)| *seen == id) {
let mut chain: Vec<&str> = stack.iter().map(|(_, n)| *n).collect();
chain.push(name);
return Err(Error::internal(format!(
"circular dependency while resolving: {}",
chain.join(" -> ")
)));
}
stack.push((id, name));
Ok(())
})?;
Ok(Self)
}
}
impl Drop for ResolutionGuard {
fn drop(&mut self) {
RESOLVING.with(|stack| {
stack.borrow_mut().pop();
});
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
#[derive(Debug)]
struct Counter(usize);
#[derive(Debug)]
struct Dependent(usize);
#[test]
fn instance_bindings_come_back_as_given() {
let c = Container::new();
c.instance(Counter(42));
assert_eq!(c.resolve::<Counter>().unwrap().0, 42);
}
#[test]
fn unbound_types_report_their_name() {
let c = Container::new();
let err = c.resolve::<Counter>().unwrap_err();
assert!(err.message().contains("Counter"), "{}", err.message());
}
#[test]
fn bind_rebuilds_every_time() {
static BUILDS: AtomicUsize = AtomicUsize::new(0);
let c = Container::new();
c.bind(|_| Ok(Counter(BUILDS.fetch_add(1, Ordering::SeqCst))));
assert_eq!(c.resolve::<Counter>().unwrap().0, 0);
assert_eq!(c.resolve::<Counter>().unwrap().0, 1);
assert_eq!(BUILDS.load(Ordering::SeqCst), 2);
}
#[test]
fn singleton_builds_once_and_shares() {
static BUILDS: AtomicUsize = AtomicUsize::new(0);
let c = Container::new();
c.singleton(|_| {
BUILDS.fetch_add(1, Ordering::SeqCst);
Ok(Counter(7))
});
let a = c.resolve::<Counter>().unwrap();
let b = c.resolve::<Counter>().unwrap();
assert!(Arc::ptr_eq(&a, &b));
assert_eq!(BUILDS.load(Ordering::SeqCst), 1);
}
#[test]
fn factories_can_resolve_their_dependencies() {
let c = Container::new();
c.instance(Counter(5));
c.singleton(|c: &Container| Ok(Dependent(c.resolve::<Counter>()?.0 * 2)));
assert_eq!(c.resolve::<Dependent>().unwrap().0, 10);
}
#[test]
fn a_failing_factory_surfaces_its_error() {
let c = Container::new();
c.singleton(|_| Err::<Counter, _>(Error::internal("no database")));
assert_eq!(c.resolve::<Counter>().unwrap_err().message(), "no database");
}
#[test]
fn dependency_cycles_error_instead_of_deadlocking() {
let c = Container::new();
c.singleton(|c: &Container| Ok(Counter(c.resolve::<Dependent>()?.0)));
c.singleton(|c: &Container| Ok(Dependent(c.resolve::<Counter>()?.0)));
let err = c.resolve::<Counter>().unwrap_err();
assert!(err.message().contains("circular dependency"), "{}", err.message());
}
#[test]
fn a_singleton_that_resolves_itself_errors_rather_than_self_deadlocking() {
let c = Container::new();
c.singleton(|c: &Container| Ok(Counter(c.resolve::<Counter>()?.0)));
let err = c.resolve::<Counter>().unwrap_err();
assert!(err.message().contains("circular dependency"), "{}", err.message());
}
#[test]
fn a_panicking_factory_does_not_poison_later_resolutions() {
let c = Arc::new(Container::new());
c.bind(|_| -> Result<Counter> { panic!("boom") });
let panicked = {
let c = Arc::clone(&c);
std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || {
let _ = c.resolve::<Counter>();
}))
};
assert!(panicked.is_err());
c.forget::<Counter>();
c.instance(Counter(1));
assert_eq!(c.resolve::<Counter>().unwrap().0, 1);
}
#[test]
fn scoped_bindings_rebuild_after_a_flush() {
static BUILDS: AtomicUsize = AtomicUsize::new(0);
let c = Container::new();
c.scoped(|_| Ok(Counter(BUILDS.fetch_add(1, Ordering::SeqCst))));
assert_eq!(c.resolve::<Counter>().unwrap().0, 0);
assert_eq!(c.resolve::<Counter>().unwrap().0, 0);
c.flush_scoped();
assert_eq!(c.resolve::<Counter>().unwrap().0, 1);
}
#[test]
fn flush_scoped_leaves_singletons_alone() {
static BUILDS: AtomicUsize = AtomicUsize::new(0);
let c = Container::new();
c.singleton(|_| Ok(Counter(BUILDS.fetch_add(1, Ordering::SeqCst))));
assert_eq!(c.resolve::<Counter>().unwrap().0, 0);
c.flush_scoped();
assert_eq!(c.resolve::<Counter>().unwrap().0, 0);
}
#[test]
fn bound_and_forget_track_the_registry() {
let c = Container::new();
assert!(!c.bound::<Counter>());
c.instance(Counter(1));
assert!(c.bound::<Counter>());
c.forget::<Counter>();
assert!(!c.bound::<Counter>());
}
#[test]
fn resolution_is_safe_across_threads() {
let c = Arc::new(Container::new());
c.singleton(|_| Ok(Counter(99)));
let handles: Vec<_> = (0..8)
.map(|_| {
let c = Arc::clone(&c);
std::thread::spawn(move || c.resolve::<Counter>().unwrap())
})
.collect();
let first = c.resolve::<Counter>().unwrap();
for handle in handles {
assert!(Arc::ptr_eq(&first, &handle.join().unwrap()));
}
}
}