Add `-Zmiri-tree-borrows-no-precise-interior-mut` flag

8 files changed, 154 insertions, 35 deletions

diff --git a/src/tools/miri/README.md b/src/tools/miri/README.md
index de521393cd0..f89708e0d8f 100644
--- a/src/tools/miri/README.md
+++ b/src/tools/miri/README.md
@@ -458,6 +458,10 @@ to Miri failing to detect cases of undefined behavior in a program.
   This is much less likely with Stacked Borrows.
   Using Tree Borrows currently implies `-Zmiri-strict-provenance` because integer-to-pointer
   casts are not supported in this mode, but that may change in the future.
+* `-Zmiri-tree-borrows-no-precise-interior-mut` makes Tree Borrows
+  track interior mutable data on the level of references instead of on the
+  byte-level as is done by default.  Therefore, with this flag, Tree
+  Borrows will be more permissive.
 * `-Zmiri-force-page-size=<num>` overrides the default page size for an architecture, in multiples of 1k.
   `4` is default for most targets. This value should always be a power of 2 and nonzero.
 
diff --git a/src/tools/miri/src/bin/miri.rs b/src/tools/miri/src/bin/miri.rs
index 0121472d330..1d3742cb0e9 100644
--- a/src/tools/miri/src/bin/miri.rs
+++ b/src/tools/miri/src/bin/miri.rs
@@ -35,7 +35,7 @@ use std::sync::{Arc, Once};
 
 use miri::{
     BacktraceStyle, BorrowTrackerMethod, GenmcConfig, GenmcCtx, MiriConfig, MiriEntryFnType,
-    ProvenanceMode, RetagFields, ValidationMode,
+    ProvenanceMode, RetagFields, TreeBorrowsParams, ValidationMode,
 };
 use rustc_abi::ExternAbi;
 use rustc_data_structures::sync;
@@ -554,8 +554,21 @@ fn main() {
         } else if arg == "-Zmiri-disable-stacked-borrows" {
             miri_config.borrow_tracker = None;
         } else if arg == "-Zmiri-tree-borrows" {
-            miri_config.borrow_tracker = Some(BorrowTrackerMethod::TreeBorrows);
+            miri_config.borrow_tracker =
+                Some(BorrowTrackerMethod::TreeBorrows(TreeBorrowsParams {
+                    precise_interior_mut: true,
+                }));
             miri_config.provenance_mode = ProvenanceMode::Strict;
+        } else if arg == "-Zmiri-tree-borrows-no-precise-interior-mut" {
+            match &mut miri_config.borrow_tracker {
+                Some(BorrowTrackerMethod::TreeBorrows(params)) => {
+                    params.precise_interior_mut = false;
+                }
+                _ =>
+                    show_error!(
+                        "`-Zmiri-tree-borrows` is required before `-Zmiri-tree-borrows-no-precise-interior-mut`"
+                    ),
+            };
         } else if arg == "-Zmiri-disable-data-race-detector" {
             miri_config.data_race_detector = false;
             miri_config.weak_memory_emulation = false;
@@ -725,7 +738,7 @@ fn main() {
         }
     }
     // Tree Borrows implies strict provenance, and is not compatible with native calls.
-    if matches!(miri_config.borrow_tracker, Some(BorrowTrackerMethod::TreeBorrows)) {
+    if matches!(miri_config.borrow_tracker, Some(BorrowTrackerMethod::TreeBorrows { .. })) {
         if miri_config.provenance_mode != ProvenanceMode::Strict {
             show_error!(
                 "Tree Borrows does not support integer-to-pointer casts, and hence requires strict provenance"
@@ -735,6 +748,7 @@ fn main() {
             show_error!("Tree Borrows is not compatible with calling native functions");
         }
     }
+
     // Native calls and strict provenance are not compatible.
     if miri_config.native_lib.is_some() && miri_config.provenance_mode == ProvenanceMode::Strict {
         show_error!("strict provenance is not compatible with calling native functions");
diff --git a/src/tools/miri/src/borrow_tracker/mod.rs b/src/tools/miri/src/borrow_tracker/mod.rs
index b66c561d2b8..36c61053a32 100644
--- a/src/tools/miri/src/borrow_tracker/mod.rs
+++ b/src/tools/miri/src/borrow_tracker/mod.rs
@@ -226,7 +226,13 @@ pub enum BorrowTrackerMethod {
     /// Stacked Borrows, as implemented in borrow_tracker/stacked_borrows
     StackedBorrows,
     /// Tree borrows, as implemented in borrow_tracker/tree_borrows
-    TreeBorrows,
+    TreeBorrows(TreeBorrowsParams),
+}
+
+/// Parameters that Tree Borrows can take.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct TreeBorrowsParams {
+    pub precise_interior_mut: bool,
 }
 
 impl BorrowTrackerMethod {
@@ -237,6 +243,13 @@ impl BorrowTrackerMethod {
             config.retag_fields,
         ))
     }
+
+    pub fn get_tree_borrows_params(self) -> TreeBorrowsParams {
+        match self {
+            BorrowTrackerMethod::TreeBorrows(params) => params,
+            _ => panic!("can only be called when `BorrowTrackerMethod` is `TreeBorrows`"),
+        }
+    }
 }
 
 impl GlobalStateInner {
@@ -252,7 +265,7 @@ impl GlobalStateInner {
                 AllocState::StackedBorrows(Box::new(RefCell::new(Stacks::new_allocation(
                     id, alloc_size, self, kind, machine,
                 )))),
-            BorrowTrackerMethod::TreeBorrows =>
+            BorrowTrackerMethod::TreeBorrows { .. } =>
                 AllocState::TreeBorrows(Box::new(RefCell::new(Tree::new_allocation(
                     id, alloc_size, self, kind, machine,
                 )))),
@@ -271,7 +284,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         let method = this.machine.borrow_tracker.as_ref().unwrap().borrow().borrow_tracker_method;
         match method {
             BorrowTrackerMethod::StackedBorrows => this.sb_retag_ptr_value(kind, val),
-            BorrowTrackerMethod::TreeBorrows => this.tb_retag_ptr_value(kind, val),
+            BorrowTrackerMethod::TreeBorrows { .. } => this.tb_retag_ptr_value(kind, val),
         }
     }
 
@@ -284,7 +297,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         let method = this.machine.borrow_tracker.as_ref().unwrap().borrow().borrow_tracker_method;
         match method {
             BorrowTrackerMethod::StackedBorrows => this.sb_retag_place_contents(kind, place),
-            BorrowTrackerMethod::TreeBorrows => this.tb_retag_place_contents(kind, place),
+            BorrowTrackerMethod::TreeBorrows { .. } => this.tb_retag_place_contents(kind, place),
         }
     }
 
@@ -293,7 +306,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         let method = this.machine.borrow_tracker.as_ref().unwrap().borrow().borrow_tracker_method;
         match method {
             BorrowTrackerMethod::StackedBorrows => this.sb_protect_place(place),
-            BorrowTrackerMethod::TreeBorrows => this.tb_protect_place(place),
+            BorrowTrackerMethod::TreeBorrows { .. } => this.tb_protect_place(place),
         }
     }
 
@@ -302,7 +315,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         let method = this.machine.borrow_tracker.as_ref().unwrap().borrow().borrow_tracker_method;
         match method {
             BorrowTrackerMethod::StackedBorrows => this.sb_expose_tag(alloc_id, tag),
-            BorrowTrackerMethod::TreeBorrows => this.tb_expose_tag(alloc_id, tag),
+            BorrowTrackerMethod::TreeBorrows { .. } => this.tb_expose_tag(alloc_id, tag),
         }
     }
 
@@ -319,7 +332,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 this.tcx.tcx.dcx().warn("Stacked Borrows does not support named pointers; `miri_pointer_name` is a no-op");
                 interp_ok(())
             }
-            BorrowTrackerMethod::TreeBorrows =>
+            BorrowTrackerMethod::TreeBorrows { .. } =>
                 this.tb_give_pointer_debug_name(ptr, nth_parent, name),
         }
     }
@@ -333,7 +346,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         let method = borrow_tracker.borrow().borrow_tracker_method;
         match method {
             BorrowTrackerMethod::StackedBorrows => this.print_stacks(alloc_id),
-            BorrowTrackerMethod::TreeBorrows => this.print_tree(alloc_id, show_unnamed),
+            BorrowTrackerMethod::TreeBorrows { .. } => this.print_tree(alloc_id, show_unnamed),
         }
     }
 
diff --git a/src/tools/miri/src/borrow_tracker/tree_borrows/mod.rs b/src/tools/miri/src/borrow_tracker/tree_borrows/mod.rs
index 411ae89da90..ce8fe03ee47 100644
--- a/src/tools/miri/src/borrow_tracker/tree_borrows/mod.rs
+++ b/src/tools/miri/src/borrow_tracker/tree_borrows/mod.rs
@@ -312,8 +312,6 @@ trait EvalContextPrivExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         }
 
         let span = this.machine.current_span();
-        let alloc_extra = this.get_alloc_extra(alloc_id)?;
-        let mut tree_borrows = alloc_extra.borrow_tracker_tb().borrow_mut();
 
         // Store initial permissions and their corresponding range.
         let mut perms_map: RangeMap<LocationState> = RangeMap::new(
@@ -342,36 +340,83 @@ trait EvalContextPrivExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         assert!(new_perm.freeze_access);
 
         let protected = new_perm.protector.is_some();
-        this.visit_freeze_sensitive(place, ptr_size, |range, frozen| {
-            has_unsafe_cell = has_unsafe_cell || !frozen;
-
-            // We are only ever `Frozen` inside the frozen bits.
-            let (perm, access) = if frozen {
+        let precise_interior_mut = this
+            .machine
+            .borrow_tracker
+            .as_mut()
+            .unwrap()
+            .get_mut()
+            .borrow_tracker_method
+            .get_tree_borrows_params()
+            .precise_interior_mut;
+
+        let default_perm = if !precise_interior_mut {
+            // NOTE: Using `ty_is_freeze` doesn't give the same result as going through the range
+            // and computing `has_unsafe_cell`.  This is because of zero-sized `UnsafeCell`, for which
+            // `has_unsafe_cell` is false, but `!ty_is_freeze` is true.
+            let ty_is_freeze = place.layout.ty.is_freeze(*this.tcx, this.typing_env());
+            let (perm, access) = if ty_is_freeze {
                 (new_perm.freeze_perm, new_perm.freeze_access)
             } else {
                 (new_perm.nonfreeze_perm, new_perm.nonfreeze_access)
             };
+            let sifa = perm.strongest_idempotent_foreign_access(protected);
+            let new_loc = if access {
+                LocationState::new_accessed(perm, sifa)
+            } else {
+                LocationState::new_non_accessed(perm, sifa)
+            };
+
+            for (_loc_range, loc) in perms_map.iter_mut_all() {
+                *loc = new_loc;
+            }
+
+            perm
+        } else {
+            this.visit_freeze_sensitive(place, ptr_size, |range, frozen| {
+                has_unsafe_cell = has_unsafe_cell || !frozen;
 
-            // Store initial permissions.
-            for (_loc_range, loc) in perms_map.iter_mut(range.start, range.size) {
+                // We are only ever `Frozen` inside the frozen bits.
+                let (perm, access) = if frozen {
+                    (new_perm.freeze_perm, new_perm.freeze_access)
+                } else {
+                    (new_perm.nonfreeze_perm, new_perm.nonfreeze_access)
+                };
                 let sifa = perm.strongest_idempotent_foreign_access(protected);
                 // NOTE: Currently, `access` is false if and only if `perm` is Cell, so this `if`
                 // doesn't not change whether any code is UB or not. We could just always use
                 // `new_accessed` and everything would stay the same. But that seems conceptually
                 // odd, so we keep the initial "accessed" bit of the `LocationState` in sync with whether
                 // a read access is performed below.
-                if access {
-                    *loc = LocationState::new_accessed(perm, sifa);
+                let new_loc = if access {
+                    LocationState::new_accessed(perm, sifa)
                 } else {
-                    *loc = LocationState::new_non_accessed(perm, sifa);
+                    LocationState::new_non_accessed(perm, sifa)
+                };
+
+                // Store initial permissions.
+                for (_loc_range, loc) in perms_map.iter_mut(range.start, range.size) {
+                    *loc = new_loc;
                 }
-            }
 
-            // Some reborrows incur a read access to the parent.
-            if access {
+                interp_ok(())
+            })?;
+
+            // Allow lazily writing to surrounding data if we found an `UnsafeCell`.
+            if has_unsafe_cell { new_perm.nonfreeze_perm } else { new_perm.freeze_perm }
+        };
+
+        let alloc_extra = this.get_alloc_extra(alloc_id)?;
+        let mut tree_borrows = alloc_extra.borrow_tracker_tb().borrow_mut();
+
+        for (perm_range, perm) in perms_map.iter_mut_all() {
+            if perm.is_accessed() {
+                // Some reborrows incur a read access to the parent.
                 // Adjust range to be relative to allocation start (rather than to `place`).
-                let mut range_in_alloc = range;
-                range_in_alloc.start += base_offset;
+                let range_in_alloc = AllocRange {
+                    start: Size::from_bytes(perm_range.start) + base_offset,
+                    size: Size::from_bytes(perm_range.end - perm_range.start),
+                };
 
                 tree_borrows.perform_access(
                     orig_tag,
@@ -382,7 +427,7 @@ trait EvalContextPrivExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 )?;
 
                 // Also inform the data race model (but only if any bytes are actually affected).
-                if range.size.bytes() > 0 {
+                if range_in_alloc.size.bytes() > 0 {
                     if let Some(data_race) = alloc_extra.data_race.as_vclocks_ref() {
                         data_race.read(
                             alloc_id,
@@ -394,8 +439,7 @@ trait EvalContextPrivExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                     }
                 }
             }
-            interp_ok(())
-        })?;
+        }
 
         // Record the parent-child pair in the tree.
         tree_borrows.new_child(
@@ -403,8 +447,7 @@ trait EvalContextPrivExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             orig_tag,
             new_tag,
             perms_map,
-            // Allow lazily writing to surrounding data if we found an `UnsafeCell`.
-            if has_unsafe_cell { new_perm.nonfreeze_perm } else { new_perm.freeze_perm },
+            default_perm,
             protected,
             span,
         )?;
diff --git a/src/tools/miri/src/diagnostics.rs b/src/tools/miri/src/diagnostics.rs
index 1728a9cfd6d..58bf163bfad 100644
--- a/src/tools/miri/src/diagnostics.rs
+++ b/src/tools/miri/src/diagnostics.rs
@@ -682,7 +682,10 @@ impl<'tcx> MiriMachine<'tcx> {
                     ),
                 ];
                 if self.borrow_tracker.as_ref().is_some_and(|b| {
-                    matches!(b.borrow().borrow_tracker_method(), BorrowTrackerMethod::TreeBorrows)
+                    matches!(
+                        b.borrow().borrow_tracker_method(),
+                        BorrowTrackerMethod::TreeBorrows { .. }
+                    )
                 }) {
                     v.push(
                         note!("Tree Borrows does not support integer-to-pointer casts, so the program is likely to go wrong when this pointer gets used")
diff --git a/src/tools/miri/src/lib.rs b/src/tools/miri/src/lib.rs
index 8802216448f..8b457b98017 100644
--- a/src/tools/miri/src/lib.rs
+++ b/src/tools/miri/src/lib.rs
@@ -112,7 +112,9 @@ pub use crate::borrow_tracker::stacked_borrows::{
     EvalContextExt as _, Item, Permission, Stack, Stacks,
 };
 pub use crate::borrow_tracker::tree_borrows::{EvalContextExt as _, Tree};
-pub use crate::borrow_tracker::{BorTag, BorrowTrackerMethod, EvalContextExt as _, RetagFields};
+pub use crate::borrow_tracker::{
+    BorTag, BorrowTrackerMethod, EvalContextExt as _, RetagFields, TreeBorrowsParams,
+};
 pub use crate::clock::{Instant, MonotonicClock};
 pub use crate::concurrency::cpu_affinity::MAX_CPUS;
 pub use crate::concurrency::data_race::{
diff --git a/src/tools/miri/tests/pass/tree_borrows/cell-inside-struct.rs b/src/tools/miri/tests/pass/tree_borrows/cell-inside-struct.rs
new file mode 100644
index 00000000000..fd68685a2f4
--- /dev/null
+++ b/src/tools/miri/tests/pass/tree_borrows/cell-inside-struct.rs
@@ -0,0 +1,34 @@
+//! The same as `tests/fail/tree-borrows/cell-inside-struct` but with
+//! precise tracking of interior mutability disabled.
+//@compile-flags: -Zmiri-tree-borrows -Zmiri-tree-borrows-no-precise-interior-mut
+#[path = "../../utils/mod.rs"]
+#[macro_use]
+mod utils;
+
+use std::cell::Cell;
+
+struct Foo {
+    field1: u32,
+    field2: Cell<u32>,
+}
+
+pub fn main() {
+    let root = Foo { field1: 42, field2: Cell::new(88) };
+    unsafe {
+        let a = &root;
+
+        name!(a as *const Foo, "a");
+
+        let a: *const Foo = a as *const Foo;
+        let a: *mut Foo = a as *mut Foo;
+
+        let alloc_id = alloc_id!(a);
+        print_state!(alloc_id);
+
+        // Writing to `field2`, which is interior mutable, should be allowed.
+        (*a).field2.set(10);
+
+        // Writing to `field1` should be allowed because it also has the `Cell` permission.
+        (*a).field1 = 88;
+    }
+}
diff --git a/src/tools/miri/tests/pass/tree_borrows/cell-inside-struct.stderr b/src/tools/miri/tests/pass/tree_borrows/cell-inside-struct.stderr
new file mode 100644
index 00000000000..1d939329040
--- /dev/null
+++ b/src/tools/miri/tests/pass/tree_borrows/cell-inside-struct.stderr
@@ -0,0 +1,6 @@
+──────────────────────────────────────────────────
+Warning: this tree is indicative only. Some tags may have been hidden.
+0..   8
+| Act |    └─┬──<TAG=root of the allocation>
+|?Cel |      └────<TAG=a>
+──────────────────────────────────────────────────