use core::cell::{Ref, RefCell, RefMut};
use crate::{GetMock, Guard, SetMock};
#[derive(Debug)]
pub struct ThreadLocal<T: Send> {
tls: thread_local::ThreadLocal<ThreadLocalInner<T>>,
}
impl<T: Send> Default for ThreadLocal<T> {
fn default() -> Self {
Self {
tls: thread_local::ThreadLocal::new(),
}
}
}
#[derive(Debug)]
struct ThreadLocalInner<T> {
inner: RefCell<Option<T>>,
write_lock: RefCell<()>,
}
impl<T> Default for ThreadLocalInner<T> {
fn default() -> Self {
Self {
inner: RefCell::new(None),
write_lock: RefCell::new(()),
}
}
}
impl<'a, T: Send + 'static> GetMock<'a, T> for ThreadLocal<T> {
type Ref = Ref<'a, T>;
fn get(&'a self) -> Option<Ref<'a, T>> {
let cell = self.tls.get_or_default();
let borrow = cell.inner.borrow();
if borrow.is_some() {
Some(Ref::map(borrow, |option| option.as_ref().unwrap()))
} else {
None
}
}
}
impl<'a, T: Send + 'static> SetMock<'a, T> for ThreadLocal<T> {
type Guard = ThreadLocalGuard<'a, T>;
fn set(&self, state: T) -> ThreadLocalGuard<'_, T> {
let cell = self.tls.get_or_default();
let guard = cell.write_lock.try_borrow_mut().unwrap_or_else(|_| {
panic!("cannot set mock state while the previous state is active");
});
*cell.inner.borrow_mut() = Some(state);
ThreadLocalGuard {
mock: &cell.inner,
_guard: guard,
}
}
}
#[derive(Debug)]
pub struct ThreadLocalGuard<'a, T> {
mock: &'a RefCell<Option<T>>,
_guard: RefMut<'a, ()>,
}
impl<T> Drop for ThreadLocalGuard<'_, T> {
fn drop(&mut self) {
self.mock.borrow_mut().take();
}
}
impl<T> Guard<T> for ThreadLocalGuard<'_, T> {
fn with<R>(&mut self, action: impl FnOnce(&mut T) -> R) -> R {
action(self.mock.borrow_mut().as_mut().unwrap())
}
fn into_inner(self) -> T {
self.mock.borrow_mut().take().unwrap()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Static;
use static_assertions::assert_impl_all;
use std::cell::Cell;
assert_impl_all!(ThreadLocal<Cell<u8>>: Send, Sync);
assert_impl_all!(Static<ThreadLocal<Cell<u8>>>: Send, Sync);
}