Repository: tokio-rs/tokio
File: tokio/src/util/ptr_expose.rs (Rust)

Code before:
//! Utility for helping miri understand our exposed pointers.
//!
//! During normal execution, this module is equivalent to pointer casts. However, when running
//! under miri, pointer casts are replaced with lookups in a hash map. This makes Tokio compatible
//! with strict provenance when running under miri (which comes with a perf cost).

use std::marker::PhantomData;
#[cfg(miri)]
use {std::collections::HashMap, crate::loom::sync::Mutex};

pub(crate) struct PtrExposeDomain<T> {
    #[cfg(miri)]
    map: Mutex<HashMap<usize, *mut T>>,
    _phantom: PhantomData<T>,
}

impl<T> PtrExposeDomain<T> {
    pub(crate) const fn new() -> Self {
        Self {
            #[cfg(miri)]
            map: Mutex::new(HashMap::new()),
            _phantom: PhantomData,
        }
    }

    #[inline]
    pub(crate) fn expose_provenance(&self, ptr: *const T) -> usize {
        #[cfg(miri)]
        {
            // SAFETY: Equivalent to `pointer::addr` which is safe.
            let addr: usize = unsafe { std::mem::transmute(ptr) };
            self.map.lock().insert(addr, ptr);
            addr
        }

        #[cfg(not(miri))]
        {
            ptr as usize
        }
    }

    #[inline]
    pub(crate) fn from_exposed_addr(&self, addr: usize) -> *const T {
        #[cfg(miri)]
        {
            self.map.lock().get(&addr).expect("Provided address is not exposed.")
        }

        #[cfg(not(miri))]
        {
            addr as *const T
        }
    }

    #[inline]
    pub(crate) fn unexpose_provenance(&self, _ptr: *const T) {
        #[cfg(miri)]
        {
            // SAFETY: Equivalent to `pointer::addr` which is safe.
            let addr: usize = unsafe { std::mem::transmute(_ptr) };
            self.map.lock().remove(addr).expect("Provided address is not exposed.");
        }
    }
}


Review comment (suggestion):
`pointer::addr` has the special status of being a libcore function so it can rely on unstable details of our semantics -- like what happens on a ptr-to-int transmute. `addr` will be adjusted if those semantics ever change, and then your code will be not equivalent to `addr` any more. User code can't just copy code from libcore and assume it can get away with the same tricks.

So this should have a FIXME to use `addr()` once possible, to ensure compatibility with future Rust versions.

Code after:
        }
    }

    #[inline]
    #[allow(clippy::wrong_self_convention)] // mirrors std name
    pub(crate) fn from_exposed_addr(&self, addr: usize) -> *const T {
        #[cfg(miri)]
        {
            let maybe_ptr = self.map.lock().get(&addr).copied();

            // SAFETY: Intentionally trigger a miri failure if the provenance we want is not
            // exposed.
            unsafe { maybe_ptr.unwrap_unchecked() }
        }

        #[cfg(not(miri))]
        {
            addr as *const T
        }
    }

    #[inline]
    pub(crate) fn unexpose_provenance(&self, _ptr: *const T) {
        #[cfg(miri)]
        {
            // SAFETY: Equivalent to `pointer::addr` which is safe.
            let addr: usize = unsafe { std::mem::transmute(_ptr) };
            let maybe_ptr = self.map.lock().remove(&addr);

            // SAFETY: Intentionally trigger a miri failure if the provenance we want is not
            // exposed.
            unsafe { maybe_ptr.unwrap_unchecked() };
        }
    }
}