pyodide_webassembly_runtime_layer/
store.rs

1use std::{
2    fmt,
3    marker::PhantomData,
4    sync::{Arc, Weak},
5};
6
7use wasm_runtime_layer::backend::{
8    AsContext, AsContextMut, WasmStore, WasmStoreContext, WasmStoreContextMut,
9};
10use wobbly::sync::Wobbly;
11
12use crate::{func::PyHostFuncFn, Engine};
13
14/// A store for the [`Engine`], which stores host-defined data `T` and internal
15/// state.
16pub struct Store<T> {
17    /// The internal store is kept behind a pointer.
18    ///
19    /// This is to allow referencing and reconstructing a calling context in
20    /// exported functions, where it is not possible to prove the correct
21    /// lifetime and borrowing rules statically nor dynamically using
22    /// [`std::cell::RefCell`]s. This is because functions can be re-entrant
23    /// with exclusive but stacked calling contexts. [`std::cell::RefCell`]
24    /// and [`std::cell::RefMut`] do not allow for recursive usage by design
25    /// (and it would be nigh impossible and quite expensive to enforce at
26    /// runtime).
27    ///
28    /// The store is stored through a raw pointer, as using a `Pin<Box<T>>`
29    /// would not be possible, despite the memory location of the Box
30    /// contents technically being pinned in memory. This is because of the
31    /// stacked borrows model.
32    ///
33    /// When the outer box is moved, it invalidates all tags in its borrow
34    /// stack, even though the memory location remains. This invalidates all
35    /// references and raw pointers to `T` created from the Box.
36    ///
37    /// See: <https://blog.nilstrieb.dev/posts/box-is-a-unique-type/> for more details.
38    ///
39    /// By using a box here, we would leave invalid pointers with revoked access
40    /// permissions to the memory location of `T`.
41    ///
42    /// This creates undefined behavior as the Rust compiler will incorrectly
43    /// optimize register accesses and memory loading and incorrect no-alias
44    /// attributes.
45    ///
46    /// To circumvent this we can use a raw pointer obtained from unwrapping a
47    /// Box.
48    ///
49    /// # Playground
50    ///
51    /// - `Pin<Box<T>>` solution (UB): <https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=685c984584bc0ca1faa780ca292f406c>
52    /// - raw pointer solution (sound): <https://play.rust-lang.org/?version=stable&mode=release&edition=2021&gist=257841cb1675106d55c756ad59fde2fb>
53    ///
54    /// You can use `Tools > Miri` to test the validity
55    inner: Arc<StoreProof>,
56    /// Marker to strongly type the [`StoreProof`]
57    _marker: PhantomData<T>,
58}
59
60/// The inner state of the store, which is pinned in heap memory
61struct StoreInner<T> {
62    /// The engine used
63    engine: Engine,
64    /// The user data
65    data: T,
66    /// The user host functions, which must live in Rust and not JS to avoid a
67    /// cross-language reference cycle
68    host_funcs: Vec<Wobbly<PyHostFuncFn>>,
69}
70
71impl<T> WasmStore<T, Engine> for Store<T> {
72    fn new(engine: &Engine, data: T) -> Self {
73        #[cfg(feature = "tracing")]
74        tracing::debug!("Store::new");
75
76        Self {
77            inner: Arc::new(StoreProof::from_ptr(Box::into_raw(Box::new(StoreInner {
78                engine: engine.clone(),
79                data,
80                host_funcs: Vec::new(),
81            })))),
82            _marker: PhantomData::<T>,
83        }
84    }
85
86    fn engine(&self) -> &Engine {
87        &self.as_inner().engine
88    }
89
90    fn data(&self) -> &T {
91        &self.as_inner().data
92    }
93
94    fn data_mut(&mut self) -> &mut T {
95        &mut self.as_inner_mut().data
96    }
97
98    fn into_data(self) -> T {
99        let this = std::mem::ManuallyDrop::new(self);
100
101        // Safety:
102        //
103        // This is the only read from self, which will not be dropped, so no duplication
104        // occurs
105        let inner = Arc::into_inner(unsafe { std::ptr::read(&this.inner) })
106            .expect("Store owns the only strong reference to StoreInner");
107
108        // Safety:
109        //
110        // Ownership of `self` signifies that no guest stack is currently active
111        let inner = unsafe { Box::from_raw(inner.as_ptr()) };
112        inner.data
113    }
114}
115
116impl<T> AsContext<Engine> for Store<T> {
117    type UserState = T;
118
119    fn as_context(&self) -> StoreContext<'_, T> {
120        StoreContext {
121            store: self.as_inner(),
122            proof: &self.inner,
123        }
124    }
125}
126
127impl<T> AsContextMut<Engine> for Store<T> {
128    fn as_context_mut(&mut self) -> StoreContextMut<'_, T> {
129        // Safety:
130        //
131        // A mutable reference to the store signifies mutable ownership, and is thus
132        // safe.
133        let store = unsafe { &mut *self.inner.as_ptr() };
134
135        StoreContextMut {
136            store,
137            proof: &mut self.inner,
138        }
139    }
140}
141
142impl<T: fmt::Debug> fmt::Debug for Store<T> {
143    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
144        let store = self.as_inner();
145
146        fmt.debug_struct("Store")
147            .field("engine", &store.engine)
148            .field("data", &store.data)
149            .finish_non_exhaustive()
150    }
151}
152
153impl<T: Default> Default for Store<T> {
154    fn default() -> Self {
155        Self::new(&Engine::default(), T::default())
156    }
157}
158
159impl<T: Clone> Clone for Store<T> {
160    fn clone(&self) -> Self {
161        Self::new(self.engine(), self.data().clone())
162    }
163}
164
165impl<T> Drop for Store<T> {
166    fn drop(&mut self) {
167        std::mem::drop(unsafe { Box::from_raw(self.inner.as_ptr::<T>()) });
168
169        #[cfg(feature = "tracing")]
170        tracing::debug!("Store::drop");
171    }
172}
173
174impl<T> Store<T> {
175    fn as_inner(&self) -> &StoreInner<T> {
176        // Safety:
177        //
178        // A shared reference to the store signifies a non-mutable ownership, and is
179        // thus safe.
180        unsafe { &*self.inner.as_ptr() }
181    }
182
183    fn as_inner_mut(&mut self) -> &mut StoreInner<T> {
184        // Safety:
185        //
186        // A mutable reference to the store signifies mutable ownership, and is thus
187        // safe.
188        unsafe { &mut *self.inner.as_ptr() }
189    }
190}
191
192#[allow(clippy::module_name_repetitions)]
193/// Immutable context to the store
194pub struct StoreContext<'a, T: 'a> {
195    /// The store
196    store: &'a StoreInner<T>,
197    /// Proof that the store is being kept alive
198    proof: &'a Arc<StoreProof>,
199}
200
201#[allow(clippy::module_name_repetitions)]
202/// Mutable context to the store
203pub struct StoreContextMut<'a, T: 'a> {
204    /// The store
205    store: &'a mut StoreInner<T>,
206    /// Proof that the store is being kept alive
207    proof: &'a mut Arc<StoreProof>,
208}
209
210impl<'a, T: 'a> StoreContextMut<'a, T> {
211    #[allow(clippy::needless_pass_by_ref_mut)]
212    /// Returns a weak proof for having a mutable borrow of the inner store
213    ///
214    /// Since the inner store provides no API surface, this weak pointer can
215    /// only be used with [`Self::from_proof_unchecked`].
216    pub(crate) fn as_weak_proof(&mut self) -> Weak<StoreProof> {
217        Arc::downgrade(self.proof)
218    }
219
220    /// Reconstructs the [`StoreContextMut`] from a strong proof of having
221    /// a mutable borrow of the inner store.
222    ///
223    /// # Safety
224    ///
225    /// The `proof` must have been constructed from a [`StoreContextMut`] with
226    /// the same generic type `T` as the one that is now created.
227    ///
228    /// The caller must be allowed to obtain a mutable (re-)borrow to the inner
229    /// store for the lifetime `'a`.
230    pub(crate) unsafe fn from_proof_unchecked(proof: &'a mut Arc<StoreProof>) -> Self {
231        Self {
232            store: unsafe { &mut *(proof.as_ptr()) },
233            proof,
234        }
235    }
236
237    pub(crate) fn register_host_func(&mut self, func: Arc<PyHostFuncFn>) -> Wobbly<PyHostFuncFn> {
238        let func = Wobbly::new(func);
239        self.store.host_funcs.push(func.clone());
240        func
241    }
242}
243
244impl<'a, T: 'a> WasmStoreContext<'a, T, Engine> for StoreContext<'a, T> {
245    fn engine(&self) -> &Engine {
246        &self.store.engine
247    }
248
249    fn data(&self) -> &T {
250        &self.store.data
251    }
252}
253
254impl<'a, T: 'a> AsContext<Engine> for StoreContext<'a, T> {
255    type UserState = T;
256
257    fn as_context(&self) -> StoreContext<'_, T> {
258        StoreContext {
259            store: self.store,
260            proof: self.proof,
261        }
262    }
263}
264
265impl<'a, T: 'a> WasmStoreContext<'a, T, Engine> for StoreContextMut<'a, T> {
266    fn engine(&self) -> &Engine {
267        &self.store.engine
268    }
269
270    fn data(&self) -> &T {
271        &self.store.data
272    }
273}
274
275impl<'a, T: 'a> WasmStoreContextMut<'a, T, Engine> for StoreContextMut<'a, T> {
276    fn data_mut(&mut self) -> &mut T {
277        &mut self.store.data
278    }
279}
280
281impl<'a, T: 'a> AsContext<Engine> for StoreContextMut<'a, T> {
282    type UserState = T;
283
284    fn as_context(&self) -> StoreContext<'_, T> {
285        StoreContext {
286            store: self.store,
287            proof: self.proof,
288        }
289    }
290}
291
292impl<'a, T: 'a> AsContextMut<Engine> for StoreContextMut<'a, T> {
293    fn as_context_mut(&mut self) -> StoreContextMut<'_, T> {
294        StoreContextMut {
295            store: self.store,
296            proof: self.proof,
297        }
298    }
299}
300
301#[allow(clippy::module_name_repetitions)]
302/// Helper type to transfer an opaque pointer to a [`StoreInner`]
303pub struct StoreProof(*mut ());
304
305unsafe impl Send for StoreProof {}
306unsafe impl Sync for StoreProof {}
307
308impl StoreProof {
309    const fn from_ptr<T>(ptr: *mut StoreInner<T>) -> Self {
310        Self(ptr.cast())
311    }
312
313    const fn as_ptr<T>(&self) -> *mut StoreInner<T> {
314        self.0.cast()
315    }
316}
pyodide_webassembly_runtime_layer/store.rs

pyodide_webassembly_runtime_layer/
store.rs
