stk/
tls.rs

1//! Utilities for storing and accessing the virtual machine from thread-local
2//! storage.
3//!
4//! **Warning:** This is potentially very unsafe, and maybe even unsound.
5//!
6//! The serde implementation of `VirtualPtr` relies on being called inside of
7//! [with_vm].
8
9use crate::vm::Vm;
10use std::cell::RefCell;
11use std::future::Future;
12use std::pin::Pin;
13use std::ptr::NonNull;
14use std::task::{Context, Poll};
15
16thread_local!(static VM: RefCell<Option<NonNull<Vm>>> = RefCell::new(None));
17
18/// Guard that restored the old VM in the threadlocal when dropped.
19struct Guard<'a>(&'a RefCell<Option<NonNull<Vm>>>, Option<NonNull<Vm>>);
20
21impl Drop for Guard<'_> {
22    fn drop(&mut self) {
23        if let Some(vm) = self.1.take() {
24            *self.0.borrow_mut() = Some(vm);
25        }
26    }
27}
28
29/// Inject the vm into TLS while running the given closure.
30pub fn inject_vm<F, O>(vm: &mut Vm, f: F) -> O
31where
32    F: FnOnce() -> O,
33{
34    let vm = unsafe { NonNull::new_unchecked(vm) };
35
36    VM.with(|storage| {
37        let old_vm = storage.borrow_mut().replace(vm);
38        let _guard = Guard(&storage, old_vm);
39        f()
40    })
41}
42
43/// Run the given closure with access to the vm.
44pub fn with_vm<F, O>(f: F) -> O
45where
46    F: FnOnce(&mut Vm) -> O,
47{
48    VM.with(|storage| {
49        let mut b = storage.borrow_mut().expect("vm must be available");
50        f(unsafe { b.as_mut() })
51    })
52}
53
54/// A future which wraps polls to have access to the TLS virtual machine.
55pub struct InjectVm<'vm, T> {
56    vm: &'vm mut Vm,
57    future: T,
58}
59
60impl<'vm, T> InjectVm<'vm, T> {
61    /// Construct a future that gives the called future thread-local access to
62    /// the virtual machine.
63    ///
64    /// # Safety
65    ///
66    /// Caller must ensure that `InjectVm` is correctly pinned w.r.t. its inner
67    /// future.
68    pub unsafe fn new(vm: &'vm mut Vm, future: T) -> Self {
69        Self { vm, future }
70    }
71}
72
73impl<'vm, T> Future for InjectVm<'vm, T>
74where
75    T: Future,
76{
77    type Output = T::Output;
78
79    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
80        // Safety: This future can only be constructed unsafely, and relies on
81        // being embedded into an already pinned future and called directly.
82        unsafe {
83            let this = Pin::into_inner_unchecked(self);
84            let future = Pin::new_unchecked(&mut this.future);
85            inject_vm(this.vm, || future.poll(cx))
86        }
87    }
88}
stk/tls.rs

stk/
tls.rs
