1//! Utility for helping miri understand our exposed pointers.
2//!
3//! During normal execution, this module is equivalent to pointer casts. However, when running
4//! under miri, pointer casts are replaced with lookups in a hash map. This makes Tokio compatible
5//! with strict provenance when running under miri (which comes with a performance cost).
6
7use std::marker::PhantomData;
8#[cfg(miri)]
9use {crate::loom::sync::Mutex, std::collections::BTreeMap};
10
11pub(crate) struct PtrExposeDomain<T> {
12 #[cfg(miri)]
13 map: Mutex<BTreeMap<usize, *const T>>,
14 _phantom: PhantomData<T>,
15}
16
17// SAFETY: Actually using the pointers is unsafe, so it's sound to transfer them across threads.
18unsafe impl<T> Sync for PtrExposeDomain<T> {}
19
20impl<T> PtrExposeDomain<T> {
21 pub(crate) const fn new() -> Self {
22 Self {
23 #[cfg(miri)]
24 map: Mutex::const_new(BTreeMap::new()),
25 _phantom: PhantomData,
26 }
27 }
28
29 #[inline]
30 pub(crate) fn expose_provenance(&self, ptr: *const T) -> usize {
31 #[cfg(miri)]
32 {
33 // FIXME: Use `pointer:addr` when it is stable.
34 // SAFETY: Equivalent to `pointer::addr` which is safe.
35 let addr: usize = unsafe { std::mem::transmute(ptr) };
36 self.map.lock().insert(addr, ptr);
37 addr
38 }
39
40 #[cfg(not(miri))]
41 {
42 ptr as usize
43 }
44 }
45
46 #[inline]
47 #[allow(clippy::wrong_self_convention)] // mirrors std name
48 pub(crate) fn from_exposed_addr(&self, addr: usize) -> *const T {
49 #[cfg(miri)]
50 {
51 let maybe_ptr = self.map.lock().get(&addr).copied();
52
53 // SAFETY: Intentionally trigger a miri failure if the provenance we want is not
54 // exposed.
55 unsafe { maybe_ptr.unwrap_unchecked() }
56 }
57
58 #[cfg(not(miri))]
59 {
60 addr as *const T
61 }
62 }
63
64 #[inline]
65 pub(crate) fn unexpose_provenance(&self, _ptr: *const T) {
66 #[cfg(miri)]
67 {
68 // SAFETY: Equivalent to `pointer::addr` which is safe.
69 let addr: usize = unsafe { std::mem::transmute(_ptr) };
70 let maybe_ptr = self.map.lock().remove(&addr);
71
72 // SAFETY: Intentionally trigger a miri failure if the provenance we want is not
73 // exposed.
74 unsafe { maybe_ptr.unwrap_unchecked() };
75 }
76 }
77}
78