Auto merge of #2613 - RalfJung:scalar-only-field-retagging, r=RalfJung

add scalar-abi-only field retagging option

`@saethlin` has requested a Stacked Borrows field retagging mode that matches when we do in terms of emitting `noalias` for LLVM. So here you go! Types with `Scalar` and `ScalarPair` ABI are being recursively retagged, everything else is not.

Given that many types can get `ScalarPair` ABI, I don't think this actually helps to make many things sound that are unsound under full field retagging -- but it might still be useful for some experimentation.

9 files changed, 128 insertions, 10 deletions

diff --git a/src/tools/miri/README.md b/src/tools/miri/README.md
index 4f5d406288f..81c4f5ffef4 100644
--- a/src/tools/miri/README.md
+++ b/src/tools/miri/README.md
@@ -377,6 +377,11 @@ to Miri failing to detect cases of undefined behavior in a program.
 * `-Zmiri-retag-fields` changes Stacked Borrows retagging to recurse into fields.
   This means that references in fields of structs/enums/tuples/arrays/... are retagged,
   and in particular, they are protected when passed as function arguments.
+* `-Zmiri-retag-fields=<all|none|scalar>` controls when Stacked Borrows retagging recurses into
+  fields. `all` means it always recurses (like `-Zmiri-retag-fields`), `none` means it never
+  recurses (the default), `scalar` means it only recurses for types where we would also emit
+  `noalias` annotations in the generated LLVM IR (types passed as indivudal scalars or pairs of
+  scalars).
 * `-Zmiri-tag-gc=<blocks>` configures how often the pointer tag garbage collector runs. The default
   is to search for and remove unreachable tags once every `10000` basic blocks. Setting this to
   `0` disables the garbage collector, which causes some programs to have explosive memory usage
diff --git a/src/tools/miri/src/bin/miri.rs b/src/tools/miri/src/bin/miri.rs
index 5b16fc2948c..2e114c71d66 100644
--- a/src/tools/miri/src/bin/miri.rs
+++ b/src/tools/miri/src/bin/miri.rs
@@ -32,7 +32,7 @@ use rustc_middle::{
 };
 use rustc_session::{config::CrateType, search_paths::PathKind, CtfeBacktrace};
 
-use miri::{BacktraceStyle, ProvenanceMode};
+use miri::{BacktraceStyle, ProvenanceMode, RetagFields};
 
 struct MiriCompilerCalls {
     miri_config: miri::MiriConfig,
@@ -426,7 +426,14 @@ fn main() {
         } else if arg == "-Zmiri-mute-stdout-stderr" {
             miri_config.mute_stdout_stderr = true;
         } else if arg == "-Zmiri-retag-fields" {
-            miri_config.retag_fields = true;
+            miri_config.retag_fields = RetagFields::Yes;
+        } else if let Some(retag_fields) = arg.strip_prefix("-Zmiri-retag-fields=") {
+            miri_config.retag_fields = match retag_fields {
+                "all" => RetagFields::Yes,
+                "none" => RetagFields::No,
+                "scalar" => RetagFields::OnlyScalar,
+                _ => show_error!("`-Zmiri-retag-fields` can only be `all`, `none`, or `scalar`"),
+            };
         } else if arg == "-Zmiri-track-raw-pointers" {
             eprintln!(
                 "WARNING: `-Zmiri-track-raw-pointers` has no effect; it is enabled by default"
diff --git a/src/tools/miri/src/eval.rs b/src/tools/miri/src/eval.rs
index b211f3c5f71..a3fc343f8b6 100644
--- a/src/tools/miri/src/eval.rs
+++ b/src/tools/miri/src/eval.rs
@@ -126,7 +126,7 @@ pub struct MiriConfig {
     /// Report the current instruction being executed every N basic blocks.
     pub report_progress: Option<u32>,
     /// Whether Stacked Borrows retagging should recurse into fields of datatypes.
-    pub retag_fields: bool,
+    pub retag_fields: RetagFields,
     /// The location of a shared object file to load when calling external functions
     /// FIXME! consider allowing users to specify paths to multiple SO files, or to a directory
     pub external_so_file: Option<PathBuf>,
@@ -163,7 +163,7 @@ impl Default for MiriConfig {
             mute_stdout_stderr: false,
             preemption_rate: 0.01, // 1%
             report_progress: None,
-            retag_fields: false,
+            retag_fields: RetagFields::No,
             external_so_file: None,
             gc_interval: 10_000,
             num_cpus: 1,
diff --git a/src/tools/miri/src/lib.rs b/src/tools/miri/src/lib.rs
index 97271c33a2e..21e65bb1b70 100644
--- a/src/tools/miri/src/lib.rs
+++ b/src/tools/miri/src/lib.rs
@@ -105,7 +105,7 @@ pub use crate::mono_hash_map::MonoHashMap;
 pub use crate::operator::EvalContextExt as _;
 pub use crate::range_map::RangeMap;
 pub use crate::stacked_borrows::{
-    CallId, EvalContextExt as _, Item, Permission, SbTag, Stack, Stacks,
+    CallId, EvalContextExt as _, Item, Permission, RetagFields, SbTag, Stack, Stacks,
 };
 pub use crate::tag_gc::{EvalContextExt as _, VisitTags};
 
diff --git a/src/tools/miri/src/stacked_borrows/mod.rs b/src/tools/miri/src/stacked_borrows/mod.rs
index f1dd38e5fc1..959e351d1a1 100644
--- a/src/tools/miri/src/stacked_borrows/mod.rs
+++ b/src/tools/miri/src/stacked_borrows/mod.rs
@@ -17,6 +17,7 @@ use rustc_middle::ty::{
     Ty,
 };
 use rustc_span::DUMMY_SP;
+use rustc_target::abi::Abi;
 use rustc_target::abi::Size;
 use smallvec::SmallVec;
 
@@ -114,7 +115,18 @@ pub struct GlobalStateInner {
     /// The call ids to trace
     tracked_call_ids: FxHashSet<CallId>,
     /// Whether to recurse into datatypes when searching for pointers to retag.
-    retag_fields: bool,
+    retag_fields: RetagFields,
+}
+
+#[derive(Copy, Clone, Debug)]
+pub enum RetagFields {
+    /// Don't retag any fields.
+    No,
+    /// Retag all fields.
+    Yes,
+    /// Only retag fields of types with Scalar and ScalarPair layout,
+    /// to match the LLVM `noalias` we generate.
+    OnlyScalar,
 }
 
 impl VisitTags for GlobalStateInner {
@@ -173,7 +185,7 @@ impl GlobalStateInner {
     pub fn new(
         tracked_pointer_tags: FxHashSet<SbTag>,
         tracked_call_ids: FxHashSet<CallId>,
-        retag_fields: bool,
+        retag_fields: RetagFields,
     ) -> Self {
         GlobalStateInner {
             next_ptr_tag: SbTag(NonZeroU64::new(1).unwrap()),
@@ -999,7 +1011,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
             ecx: &'ecx mut MiriInterpCx<'mir, 'tcx>,
             kind: RetagKind,
             retag_cause: RetagCause,
-            retag_fields: bool,
+            retag_fields: RetagFields,
         }
         impl<'ecx, 'mir, 'tcx> RetagVisitor<'ecx, 'mir, 'tcx> {
             #[inline(always)] // yes this helps in our benchmarks
@@ -1046,6 +1058,18 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                     return Ok(());
                 }
 
+                let recurse_for_fields = || {
+                    match self.retag_fields {
+                        RetagFields::No => false,
+                        RetagFields::Yes => true,
+                        RetagFields::OnlyScalar => {
+                            // Matching `ArgAbi::new` at the time of writing, only fields of
+                            // `Scalar` and `ScalarPair` ABI are considered.
+                            matches!(place.layout.abi, Abi::Scalar(..) | Abi::ScalarPair(..))
+                        }
+                    }
+                };
+
                 if let Some((ref_kind, protector)) = qualify(place.layout.ty, self.kind) {
                     self.retag_place(place, ref_kind, self.retag_cause, protector)?;
                 } else if matches!(place.layout.ty.kind(), ty::RawPtr(..)) {
@@ -1054,7 +1078,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                     // Do *not* recurse into them.
                     // (No need to worry about wide references, those always "qualify". And Boxes
                     // are handles specially by the visitor anyway.)
-                } else if self.retag_fields
+                } else if recurse_for_fields()
                     || place.layout.ty.ty_adt_def().is_some_and(|adt| adt.is_box())
                 {
                     // Recurse deeper. Need to always recurse for `Box` to even hit `visit_box`.
diff --git a/src/tools/miri/tests/fail/stacked_borrows/newtype_pair_retagging.rs b/src/tools/miri/tests/fail/stacked_borrows/newtype_pair_retagging.rs
new file mode 100644
index 00000000000..5cefdb08e78
--- /dev/null
+++ b/src/tools/miri/tests/fail/stacked_borrows/newtype_pair_retagging.rs
@@ -0,0 +1,19 @@
+//@compile-flags: -Zmiri-retag-fields=scalar
+//@error-pattern: which is protected
+struct Newtype<'a>(&'a mut i32, i32);
+
+fn dealloc_while_running(_n: Newtype<'_>, dealloc: impl FnOnce()) {
+    dealloc();
+}
+
+// Make sure that we protect references inside structs that are passed as ScalarPair.
+fn main() {
+    let ptr = Box::into_raw(Box::new(0i32));
+    #[rustfmt::skip] // I like my newlines
+    unsafe {
+        dealloc_while_running(
+            Newtype(&mut *ptr, 0),
+            || drop(Box::from_raw(ptr)),
+        )
+    };
+}
diff --git a/src/tools/miri/tests/fail/stacked_borrows/newtype_pair_retagging.stderr b/src/tools/miri/tests/fail/stacked_borrows/newtype_pair_retagging.stderr
new file mode 100644
index 00000000000..60a8b2a6260
--- /dev/null
+++ b/src/tools/miri/tests/fail/stacked_borrows/newtype_pair_retagging.stderr
@@ -0,0 +1,44 @@
+error: Undefined Behavior: not granting access to tag <TAG> because that would remove [Unique for <TAG>] which is protected because it is an argument of call ID
+  --> RUSTLIB/alloc/src/boxed.rs:LL:CC
+   |
+LL |         Box(unsafe { Unique::new_unchecked(raw) }, alloc)
+   |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ not granting access to tag <TAG> because that would remove [Unique for <TAG>] which is protected because it is an argument of call ID
+   |
+   = help: this indicates a potential bug in the program: it performed an invalid operation, but the Stacked Borrows rules it violated are still experimental
+   = help: see https://github.com/rust-lang/unsafe-code-guidelines/blob/master/wip/stacked-borrows.md for further information
+help: <TAG> was created by a SharedReadWrite retag at offsets [0x0..0x4]
+  --> $DIR/newtype_pair_retagging.rs:LL:CC
+   |
+LL |     let ptr = Box::into_raw(Box::new(0i32));
+   |               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+help: <TAG> is this argument
+  --> $DIR/newtype_pair_retagging.rs:LL:CC
+   |
+LL | fn dealloc_while_running(_n: Newtype<'_>, dealloc: impl FnOnce()) {
+   |                          ^^
+   = note: BACKTRACE:
+   = note: inside `std::boxed::Box::<i32>::from_raw_in` at RUSTLIB/alloc/src/boxed.rs:LL:CC
+   = note: inside `std::boxed::Box::<i32>::from_raw` at RUSTLIB/alloc/src/boxed.rs:LL:CC
+note: inside closure at $DIR/newtype_pair_retagging.rs:LL:CC
+  --> $DIR/newtype_pair_retagging.rs:LL:CC
+   |
+LL |             || drop(Box::from_raw(ptr)),
+   |                     ^^^^^^^^^^^^^^^^^^
+note: inside `dealloc_while_running::<[closure@$DIR/newtype_pair_retagging.rs:LL:CC]>` at $DIR/newtype_pair_retagging.rs:LL:CC
+  --> $DIR/newtype_pair_retagging.rs:LL:CC
+   |
+LL |     dealloc();
+   |     ^^^^^^^^^
+note: inside `main` at $DIR/newtype_pair_retagging.rs:LL:CC
+  --> $DIR/newtype_pair_retagging.rs:LL:CC
+   |
+LL | /         dealloc_while_running(
+LL | |             Newtype(&mut *ptr, 0),
+LL | |             || drop(Box::from_raw(ptr)),
+LL | |         )
+   | |_________^
+
+note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace
+
+error: aborting due to previous error
+
diff --git a/src/tools/miri/tests/fail/stacked_borrows/newtype_retagging.rs b/src/tools/miri/tests/fail/stacked_borrows/newtype_retagging.rs
index 6e7413cff5d..bc3883575c3 100644
--- a/src/tools/miri/tests/fail/stacked_borrows/newtype_retagging.rs
+++ b/src/tools/miri/tests/fail/stacked_borrows/newtype_retagging.rs
@@ -1,4 +1,4 @@
-//@compile-flags: -Zmiri-retag-fields
+//@compile-flags: -Zmiri-retag-fields=scalar
 //@error-pattern: which is protected
 struct Newtype<'a>(&'a mut i32);
 
diff --git a/src/tools/miri/tests/pass/stacked-borrows/non_scalar_field_retagging.rs b/src/tools/miri/tests/pass/stacked-borrows/non_scalar_field_retagging.rs
new file mode 100644
index 00000000000..ddedc19c999
--- /dev/null
+++ b/src/tools/miri/tests/pass/stacked-borrows/non_scalar_field_retagging.rs
@@ -0,0 +1,19 @@
+//@compile-flags: -Zmiri-retag-fields=scalar
+
+struct Newtype<'a>(&'a mut i32, i32, i32);
+
+fn dealloc_while_running(_n: Newtype<'_>, dealloc: impl FnOnce()) {
+    dealloc();
+}
+
+// Make sure that with -Zmiri-retag-fields=scalar, we do *not* retag the fields of `Newtype`.
+fn main() {
+    let ptr = Box::into_raw(Box::new(0i32));
+    #[rustfmt::skip] // I like my newlines
+    unsafe {
+        dealloc_while_running(
+            Newtype(&mut *ptr, 0, 0),
+            || drop(Box::from_raw(ptr)),
+        )
+    };
+}