move lazy_sync helper methods to be with InterpCx

4 files changed, 103 insertions, 79 deletions

diff --git a/src/tools/miri/src/concurrency/sync.rs b/src/tools/miri/src/concurrency/sync.rs
index 2c6a7bf0f58..e7e6c100cfe 100644
--- a/src/tools/miri/src/concurrency/sync.rs
+++ b/src/tools/miri/src/concurrency/sync.rs
@@ -193,75 +193,104 @@ impl<'tcx> AllocExtra<'tcx> {
 /// If `init` is set to this, we consider the primitive initialized.
 pub const LAZY_INIT_COOKIE: u32 = 0xcafe_affe;
 
-/// Helper for lazily initialized `alloc_extra.sync` data:
-/// this forces an immediate init.
-pub fn lazy_sync_init<'tcx, T: 'static + Copy>(
-    ecx: &mut MiriInterpCx<'tcx>,
-    primitive: &MPlaceTy<'tcx>,
-    init_offset: Size,
-    data: T,
-) -> InterpResult<'tcx> {
-    let (alloc, offset, _) = ecx.ptr_get_alloc_id(primitive.ptr(), 0)?;
-    let (alloc_extra, _machine) = ecx.get_alloc_extra_mut(alloc)?;
-    alloc_extra.sync.insert(offset, Box::new(data));
-    // Mark this as "initialized".
-    let init_field = primitive.offset(init_offset, ecx.machine.layouts.u32, ecx)?;
-    ecx.write_scalar_atomic(
-        Scalar::from_u32(LAZY_INIT_COOKIE),
-        &init_field,
-        AtomicWriteOrd::Relaxed,
-    )?;
-    interp_ok(())
-}
-
-/// Helper for lazily initialized `alloc_extra.sync` data:
-/// Checks if the primitive is initialized, and return its associated data if so.
-/// Otherwise, calls `new_data` to initialize the primitive.
-pub fn lazy_sync_get_data<'tcx, T: 'static + Copy>(
-    ecx: &mut MiriInterpCx<'tcx>,
-    primitive: &MPlaceTy<'tcx>,
-    init_offset: Size,
-    name: &str,
-    new_data: impl FnOnce(&mut MiriInterpCx<'tcx>) -> InterpResult<'tcx, T>,
-) -> InterpResult<'tcx, T> {
-    // Check if this is already initialized. Needs to be atomic because we can race with another
-    // thread initializing. Needs to be an RMW operation to ensure we read the *latest* value.
-    // So we just try to replace MUTEX_INIT_COOKIE with itself.
-    let init_cookie = Scalar::from_u32(LAZY_INIT_COOKIE);
-    let init_field = primitive.offset(init_offset, ecx.machine.layouts.u32, ecx)?;
-    let (_init, success) = ecx
-        .atomic_compare_exchange_scalar(
-            &init_field,
-            &ImmTy::from_scalar(init_cookie, ecx.machine.layouts.u32),
-            init_cookie,
-            AtomicRwOrd::Relaxed,
-            AtomicReadOrd::Relaxed,
-            /* can_fail_spuriously */ false,
-        )?
-        .to_scalar_pair();
-
-    if success.to_bool()? {
-        // If it is initialized, it must be found in the "sync primitive" table,
-        // or else it has been moved illegally.
-        let (alloc, offset, _) = ecx.ptr_get_alloc_id(primitive.ptr(), 0)?;
-        let alloc_extra = ecx.get_alloc_extra(alloc)?;
-        let data = alloc_extra
-            .get_sync::<T>(offset)
-            .ok_or_else(|| err_ub_format!("`{name}` can't be moved after first use"))?;
-        interp_ok(*data)
-    } else {
-        let data = new_data(ecx)?;
-        lazy_sync_init(ecx, primitive, init_offset, data)?;
-        interp_ok(data)
-    }
-}
-
 // Public interface to synchronization primitives. Please note that in most
 // cases, the function calls are infallible and it is the client's (shim
 // implementation's) responsibility to detect and deal with erroneous
 // situations.
 impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
 pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
+    /// Helper for lazily initialized `alloc_extra.sync` data:
+    /// this forces an immediate init.
+    fn lazy_sync_init<T: 'static + Copy>(
+        &mut self,
+        primitive: &MPlaceTy<'tcx>,
+        init_offset: Size,
+        data: T,
+    ) -> InterpResult<'tcx> {
+        let this = self.eval_context_mut();
+
+        let (alloc, offset, _) = this.ptr_get_alloc_id(primitive.ptr(), 0)?;
+        let (alloc_extra, _machine) = this.get_alloc_extra_mut(alloc)?;
+        alloc_extra.sync.insert(offset, Box::new(data));
+        // Mark this as "initialized".
+        let init_field = primitive.offset(init_offset, this.machine.layouts.u32, this)?;
+        this.write_scalar_atomic(
+            Scalar::from_u32(LAZY_INIT_COOKIE),
+            &init_field,
+            AtomicWriteOrd::Relaxed,
+        )?;
+        interp_ok(())
+    }
+
+    /// Helper for lazily initialized `alloc_extra.sync` data:
+    /// Checks if the primitive is initialized, and return its associated data if so.
+    /// Otherwise, calls `new_data` to initialize the primitive.
+    fn lazy_sync_get_data<T: 'static + Copy>(
+        &mut self,
+        primitive: &MPlaceTy<'tcx>,
+        init_offset: Size,
+        name: &str,
+        new_data: impl FnOnce(&mut MiriInterpCx<'tcx>) -> InterpResult<'tcx, T>,
+    ) -> InterpResult<'tcx, T> {
+        let this = self.eval_context_mut();
+
+        // Check if this is already initialized. Needs to be atomic because we can race with another
+        // thread initializing. Needs to be an RMW operation to ensure we read the *latest* value.
+        // So we just try to replace MUTEX_INIT_COOKIE with itself.
+        let init_cookie = Scalar::from_u32(LAZY_INIT_COOKIE);
+        let init_field = primitive.offset(init_offset, this.machine.layouts.u32, this)?;
+        let (_init, success) = this
+            .atomic_compare_exchange_scalar(
+                &init_field,
+                &ImmTy::from_scalar(init_cookie, this.machine.layouts.u32),
+                init_cookie,
+                AtomicRwOrd::Relaxed,
+                AtomicReadOrd::Relaxed,
+                /* can_fail_spuriously */ false,
+            )?
+            .to_scalar_pair();
+
+        if success.to_bool()? {
+            // If it is initialized, it must be found in the "sync primitive" table,
+            // or else it has been moved illegally.
+            let (alloc, offset, _) = this.ptr_get_alloc_id(primitive.ptr(), 0)?;
+            let alloc_extra = this.get_alloc_extra(alloc)?;
+            let data = alloc_extra
+                .get_sync::<T>(offset)
+                .ok_or_else(|| err_ub_format!("`{name}` can't be moved after first use"))?;
+            interp_ok(*data)
+        } else {
+            let data = new_data(this)?;
+            this.lazy_sync_init(primitive, init_offset, data)?;
+            interp_ok(data)
+        }
+    }
+
+    /// Get the synchronization primitive associated with the given pointer,
+    /// or initialize a new one.
+    fn get_sync_or_init<'a, T: 'static>(
+        &'a mut self,
+        ptr: Pointer,
+        new: impl FnOnce(&'a mut MiriMachine<'tcx>) -> InterpResult<'tcx, T>,
+    ) -> InterpResult<'tcx, &'a T>
+    where
+        'tcx: 'a,
+    {
+        let this = self.eval_context_mut();
+        // Ensure there is memory behind this pointer, so that this allocation
+        // is truly the only place where the data could be stored.
+        this.check_ptr_access(ptr, Size::from_bytes(1), CheckInAllocMsg::InboundsTest)?;
+
+        let (alloc, offset, _) = this.ptr_get_alloc_id(ptr, 0)?;
+        let (alloc_extra, machine) = this.get_alloc_extra_mut(alloc)?;
+        // Due to borrow checker reasons, we have to do the lookup twice.
+        if alloc_extra.get_sync::<T>(offset).is_none() {
+            let new = new(machine)?;
+            alloc_extra.sync.insert(offset, Box::new(new));
+        }
+        interp_ok(alloc_extra.get_sync::<T>(offset).unwrap())
+    }
+
     #[inline]
     /// Get the id of the thread that currently owns this lock.
     fn mutex_get_owner(&mut self, id: MutexId) -> ThreadId {
diff --git a/src/tools/miri/src/shims/unix/macos/sync.rs b/src/tools/miri/src/shims/unix/macos/sync.rs
index cd5b0ed1d07..25e1ff42fb9 100644
--- a/src/tools/miri/src/shims/unix/macos/sync.rs
+++ b/src/tools/miri/src/shims/unix/macos/sync.rs
@@ -23,15 +23,11 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
         let lock = this.deref_pointer(lock_ptr)?;
         // We store the mutex ID in the `sync` metadata. This means that when the lock is moved,
         // that's just implicitly creating a new lock at the new location.
-        let (alloc, offset, _) = this.ptr_get_alloc_id(lock.ptr(), 0)?;
-        let (alloc_extra, machine) = this.get_alloc_extra_mut(alloc)?;
-        if let Some(data) = alloc_extra.get_sync::<MacOsUnfairLock>(offset) {
-            interp_ok(data.id)
-        } else {
+        let data = this.get_sync_or_init(lock.ptr(), |machine| {
             let id = machine.sync.mutex_create();
-            alloc_extra.sync.insert(offset, Box::new(MacOsUnfairLock { id }));
-            interp_ok(id)
-        }
+            interp_ok(MacOsUnfairLock { id })
+        })?;
+        interp_ok(data.id)
     }
 }
 
diff --git a/src/tools/miri/src/shims/unix/sync.rs b/src/tools/miri/src/shims/unix/sync.rs
index 913f53adbbb..8eb874a4565 100644
--- a/src/tools/miri/src/shims/unix/sync.rs
+++ b/src/tools/miri/src/shims/unix/sync.rs
@@ -2,7 +2,7 @@ use std::sync::atomic::{AtomicBool, Ordering};
 
 use rustc_target::abi::Size;
 
-use crate::concurrency::sync::{LAZY_INIT_COOKIE, lazy_sync_get_data, lazy_sync_init};
+use crate::concurrency::sync::LAZY_INIT_COOKIE;
 use crate::*;
 
 /// Do a bytewise comparison of the two places, using relaxed atomic reads. This is used to check if
@@ -176,7 +176,7 @@ fn mutex_create<'tcx>(
     let mutex = ecx.deref_pointer(mutex_ptr)?;
     let id = ecx.machine.sync.mutex_create();
     let data = PthreadMutex { id, kind };
-    lazy_sync_init(ecx, &mutex, mutex_init_offset(ecx)?, data)?;
+    ecx.lazy_sync_init(&mutex, mutex_init_offset(ecx)?, data)?;
     interp_ok(data)
 }
 
@@ -189,7 +189,7 @@ fn mutex_get_data<'tcx, 'a>(
     mutex_ptr: &OpTy<'tcx>,
 ) -> InterpResult<'tcx, PthreadMutex> {
     let mutex = ecx.deref_pointer(mutex_ptr)?;
-    lazy_sync_get_data(ecx, &mutex, mutex_init_offset(ecx)?, "pthread_mutex_t", |ecx| {
+    ecx.lazy_sync_get_data(&mutex, mutex_init_offset(ecx)?, "pthread_mutex_t", |ecx| {
         let kind = mutex_kind_from_static_initializer(ecx, &mutex)?;
         let id = ecx.machine.sync.mutex_create();
         interp_ok(PthreadMutex { id, kind })
@@ -261,7 +261,7 @@ fn rwlock_get_data<'tcx>(
     rwlock_ptr: &OpTy<'tcx>,
 ) -> InterpResult<'tcx, PthreadRwLock> {
     let rwlock = ecx.deref_pointer(rwlock_ptr)?;
-    lazy_sync_get_data(ecx, &rwlock, rwlock_init_offset(ecx)?, "pthread_rwlock_t", |ecx| {
+    ecx.lazy_sync_get_data(&rwlock, rwlock_init_offset(ecx)?, "pthread_rwlock_t", |ecx| {
         if !bytewise_equal_atomic_relaxed(
             ecx,
             &rwlock,
@@ -377,7 +377,7 @@ fn cond_create<'tcx>(
     let cond = ecx.deref_pointer(cond_ptr)?;
     let id = ecx.machine.sync.condvar_create();
     let data = PthreadCondvar { id, clock };
-    lazy_sync_init(ecx, &cond, cond_init_offset(ecx)?, data)?;
+    ecx.lazy_sync_init(&cond, cond_init_offset(ecx)?, data)?;
     interp_ok(data)
 }
 
@@ -386,7 +386,7 @@ fn cond_get_data<'tcx>(
     cond_ptr: &OpTy<'tcx>,
 ) -> InterpResult<'tcx, PthreadCondvar> {
     let cond = ecx.deref_pointer(cond_ptr)?;
-    lazy_sync_get_data(ecx, &cond, cond_init_offset(ecx)?, "pthread_cond_t", |ecx| {
+    ecx.lazy_sync_get_data(&cond, cond_init_offset(ecx)?, "pthread_cond_t", |ecx| {
         if !bytewise_equal_atomic_relaxed(
             ecx,
             &cond,
diff --git a/src/tools/miri/src/shims/windows/sync.rs b/src/tools/miri/src/shims/windows/sync.rs
index 8771bb4a8ac..3701f479e5c 100644
--- a/src/tools/miri/src/shims/windows/sync.rs
+++ b/src/tools/miri/src/shims/windows/sync.rs
@@ -3,7 +3,6 @@ use std::time::Duration;
 use rustc_target::abi::Size;
 
 use crate::concurrency::init_once::InitOnceStatus;
-use crate::concurrency::sync::lazy_sync_get_data;
 use crate::*;
 
 #[derive(Copy, Clone)]
@@ -25,7 +24,7 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
         let init_once = this.deref_pointer(init_once_ptr)?;
         let init_offset = Size::ZERO;
 
-        lazy_sync_get_data(this, &init_once, init_offset, "INIT_ONCE", |this| {
+        this.lazy_sync_get_data(&init_once, init_offset, "INIT_ONCE", |this| {
             // TODO: check that this is still all-zero.
             let id = this.machine.sync.init_once_create();
             interp_ok(WindowsInitOnce { id })