Rollup merge of #122233 - RalfJung:custom-alloc-box, r=oli-obk

miri: do not apply aliasing restrictions to Box with custom allocator

This is the Miri side of rust-lang/rust#122018. The "intrinsics with body" made this much more pleasant. :)

Fixes #3341.
r? `@oli-obk`

Author: matthiaskrgr

src/borrow_tracker/mod.rs

@@ -5,7 +5,7 @@ use std::num::NonZero;
 use smallvec::SmallVec;
 
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
-use rustc_middle::mir::RetagKind;
+use rustc_middle::{mir::RetagKind, ty::Ty};
 use rustc_target::abi::Size;
 
 use crate::*;
@@ -291,6 +291,19 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         }
     }
 
+    fn retag_box_to_raw(
+        &mut self,
+        val: &ImmTy<'tcx, Provenance>,
+        alloc_ty: Ty<'tcx>,
+    ) -> InterpResult<'tcx, ImmTy<'tcx, Provenance>> {
+        let this = self.eval_context_mut();
+        let method = this.machine.borrow_tracker.as_ref().unwrap().borrow().borrow_tracker_method;
+        match method {
+            BorrowTrackerMethod::StackedBorrows => this.sb_retag_box_to_raw(val, alloc_ty),
+            BorrowTrackerMethod::TreeBorrows => this.tb_retag_box_to_raw(val, alloc_ty),
+        }
+    }
+
     fn retag_place_contents(
         &mut self,
         kind: RetagKind,

src/borrow_tracker/stacked_borrows/mod.rs

@@ -865,6 +865,24 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         this.sb_retag_reference(val, new_perm, RetagInfo { cause, in_field: false })
     }
 
+    fn sb_retag_box_to_raw(
+        &mut self,
+        val: &ImmTy<'tcx, Provenance>,
+        alloc_ty: Ty<'tcx>,
+    ) -> InterpResult<'tcx, ImmTy<'tcx, Provenance>> {
+        let this = self.eval_context_mut();
+        let is_global_alloc = alloc_ty.ty_adt_def().is_some_and(|adt| {
+            let global_alloc = this.tcx.require_lang_item(rustc_hir::LangItem::GlobalAlloc, None);
+            adt.did() == global_alloc
+        });
+        if is_global_alloc {
+            // Retag this as-if it was a mutable reference.
+            this.sb_retag_ptr_value(RetagKind::Raw, val)
+        } else {
+            Ok(val.clone())
+        }
+    }
+
     fn sb_retag_place_contents(
         &mut self,
         kind: RetagKind,
@@ -916,10 +934,18 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                 self.ecx
             }
 
-            fn visit_box(&mut self, place: &PlaceTy<'tcx, Provenance>) -> InterpResult<'tcx> {
-                // Boxes get a weak protectors, since they may be deallocated.
-                let new_perm = NewPermission::from_box_ty(place.layout.ty, self.kind, self.ecx);
-                self.retag_ptr_inplace(place, new_perm)
+            fn visit_box(
+                &mut self,
+                box_ty: Ty<'tcx>,
+                place: &PlaceTy<'tcx, Provenance>,
+            ) -> InterpResult<'tcx> {
+                // Only boxes for the global allocator get any special treatment.
+                if box_ty.is_box_global(*self.ecx.tcx) {
+                    // Boxes get a weak protectors, since they may be deallocated.
+                    let new_perm = NewPermission::from_box_ty(place.layout.ty, self.kind, self.ecx);
+                    self.retag_ptr_inplace(place, new_perm)?;
+                }
+                Ok(())
             }
 
             fn visit_value(&mut self, place: &PlaceTy<'tcx, Provenance>) -> InterpResult<'tcx> {

src/borrow_tracker/tree_borrows/mod.rs

@@ -392,6 +392,15 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         }
     }
 
+    fn tb_retag_box_to_raw(
+        &mut self,
+        val: &ImmTy<'tcx, Provenance>,
+        _alloc_ty: Ty<'tcx>,
+    ) -> InterpResult<'tcx, ImmTy<'tcx, Provenance>> {
+        // Casts to raw pointers are NOPs in Tree Borrows.
+        Ok(val.clone())
+    }
+
     /// Retag all pointers that are stored in this place.
     fn tb_retag_place_contents(
         &mut self,
@@ -441,14 +450,22 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
             /// Regardless of how `Unique` is handled, Boxes are always reborrowed.
             /// When `Unique` is also reborrowed, then it behaves exactly like `Box`
             /// except for the fact that `Box` has a non-zero-sized reborrow.
-            fn visit_box(&mut self, place: &PlaceTy<'tcx, Provenance>) -> InterpResult<'tcx> {
-                let new_perm = NewPermission::from_unique_ty(
-                    place.layout.ty,
-                    self.kind,
-                    self.ecx,
-                    /* zero_size */ false,
-                );
-                self.retag_ptr_inplace(place, new_perm)
+            fn visit_box(
+                &mut self,
+                box_ty: Ty<'tcx>,
+                place: &PlaceTy<'tcx, Provenance>,
+            ) -> InterpResult<'tcx> {
+                // Only boxes for the global allocator get any special treatment.
+                if box_ty.is_box_global(*self.ecx.tcx) {
+                    let new_perm = NewPermission::from_unique_ty(
+                        place.layout.ty,
+                        self.kind,
+                        self.ecx,
+                        /* zero_size */ false,
+                    );
+                    self.retag_ptr_inplace(place, new_perm)?;
+                }
+                Ok(())
             }
 
             fn visit_value(&mut self, place: &PlaceTy<'tcx, Provenance>) -> InterpResult<'tcx> {

src/shims/intrinsics/mod.rs

@@ -129,6 +129,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                 this.write_bytes_ptr(ptr, iter::repeat(val_byte).take(byte_count.bytes_usize()))?;
             }
 
+            // Memory model / provenance manipulation
             "ptr_mask" => {
                 let [ptr, mask] = check_arg_count(args)?;
 
@@ -139,6 +140,18 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
 
                 this.write_pointer(Pointer::new(ptr.provenance, masked_addr), dest)?;
             }
+            "retag_box_to_raw" => {
+                let [ptr] = check_arg_count(args)?;
+                let alloc_ty = generic_args[1].expect_ty();
+
+                let val = this.read_immediate(ptr)?;
+                let new_val = if this.machine.borrow_tracker.is_some() {
+                    this.retag_box_to_raw(&val, alloc_ty)?
+                } else {
+                    val
+                };
+                this.write_immediate(*new_val, dest)?;
+            }
 
             // We want to return either `true` or `false` at random, or else something like
             // ```

tests/pass/box-custom-alloc-aliasing.rs

@@ -0,0 +1,123 @@
+//! Regression test for <https://github.com/rust-lang/miri/issues/3341>:
+//! If `Box` has a local allocator, then it can't be `noalias` as the allocator
+//! may want to access allocator state based on the data pointer.
+
+//@revisions: stack tree
+//@[tree]compile-flags: -Zmiri-tree-borrows
+#![feature(allocator_api)]
+#![feature(strict_provenance)]
+
+use std::{
+    alloc::{AllocError, Allocator, Layout},
+    cell::{Cell, UnsafeCell},
+    ptr::{self, addr_of, NonNull},
+    thread::{self, ThreadId},
+    mem,
+};
+
+const BIN_SIZE: usize = 8;
+
+// A bin represents a collection of blocks of a specific layout.
+#[repr(align(128))]
+struct MyBin {
+    top: Cell<usize>,
+    thread_id: ThreadId,
+    memory: UnsafeCell<[usize; BIN_SIZE]>,
+}
+
+impl MyBin {
+    fn pop(&self) -> Option<NonNull<u8>> {
+        let top = self.top.get();
+        if top == BIN_SIZE {
+            return None;
+        }
+        // Cast the *entire* thing to a raw pointer to not restrict its provenance.
+        let bin = self as *const MyBin;
+        let base_ptr = UnsafeCell::raw_get(unsafe{ addr_of!((*bin).memory )}).cast::<usize>();
+        let ptr = unsafe { NonNull::new_unchecked(base_ptr.add(top)) };
+        self.top.set(top + 1);
+        Some(ptr.cast())
+    }
+
+    // Pretends to not be a throwaway allocation method like this. A more realistic
+    // substitute is using intrusive linked lists, which requires access to the
+    // metadata of this bin as well.
+    unsafe fn push(&self, ptr: NonNull<u8>) {
+        // For now just check that this really is in this bin.
+        let start = self.memory.get().addr();
+        let end = start + BIN_SIZE * mem::size_of::<usize>();
+        let addr = ptr.addr().get();
+        assert!((start..end).contains(&addr));
+    }
+}
+
+// A collection of bins.
+struct MyAllocator {
+    thread_id: ThreadId,
+    // Pretends to be some complex collection of bins, such as an array of linked lists.
+    bins: Box<[MyBin; 1]>,
+}
+
+impl MyAllocator {
+    fn new() -> Self {
+        let thread_id = thread::current().id();
+        MyAllocator {
+            thread_id,
+            bins: Box::new(
+                [MyBin {
+                    top: Cell::new(0),
+                    thread_id,
+                    memory: UnsafeCell::default(),
+                }; 1],
+            ),
+        }
+    }
+
+    // Pretends to be expensive finding a suitable bin for the layout.
+    fn find_bin(&self, layout: Layout) -> Option<&MyBin> {
+        if layout == Layout::new::<usize>() {
+            Some(&self.bins[0])
+        } else {
+            None
+        }
+    }
+}
+
+unsafe impl Allocator for MyAllocator {
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        // Expensive bin search.
+        let bin = self.find_bin(layout).ok_or(AllocError)?;
+        let ptr = bin.pop().ok_or(AllocError)?;
+        Ok(NonNull::slice_from_raw_parts(ptr, layout.size()))
+    }
+
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, _layout: Layout) {
+        // Since manually finding the corresponding bin of `ptr` is very expensive,
+        // doing pointer arithmetics is preferred.
+        // But this means we access `top` via `ptr` rather than `self`!
+        // That is fundamentally the source of the aliasing trouble in this example.
+        let their_bin = ptr.as_ptr().map_addr(|addr| addr & !127).cast::<MyBin>();
+        let thread_id = ptr::read(ptr::addr_of!((*their_bin).thread_id));
+        if self.thread_id == thread_id {
+            unsafe { (*their_bin).push(ptr) };
+        } else {
+            todo!("Deallocating from another thread")
+        }
+    }
+}
+
+// Make sure to involve `Box` in allocating these,
+// as that's where `noalias` may come from.
+fn v<T, A: Allocator>(t: T, a: A) -> Vec<T, A> {
+    (Box::new_in([t], a) as Box<[T], A>).into_vec()
+}
+
+fn main() {
+    assert!(mem::size_of::<MyBin>() <= 128); // if it grows bigger, the trick to access the "header" no longer works
+    let my_alloc = MyAllocator::new();
+    let a = v(1usize, &my_alloc);
+    let b = v(2usize, &my_alloc);
+    assert_eq!(a[0] + 1, b[0]);
+    assert_eq!(addr_of!(a[0]).wrapping_add(1), addr_of!(b[0]));
+    drop((a, b));
+}