split up prepare_for_native_access + code review

Author: nia-e

compiler/rustc_middle/src/mir/interpret/allocation.rs

@@ -796,8 +796,8 @@ impl<Prov: Provenance, Extra, Bytes: AllocBytes> Allocation<Prov, Extra, Bytes>
         Ok(())
     }
 
-    /// Initialize all previously uninitialized bytes in the entire allocation, and set
-    /// provenance of everything to `Wildcard`. Before calling this, make sure all
+    /// Initialize all previously uninitialized bytes in the entire allocation, but
+    /// do not actually mark them as init. Before calling this, make sure all
     /// provenance in this allocation is exposed!
     pub fn prepare_for_native_access(&mut self) {
         let full_range = AllocRange { start: Size::ZERO, size: Size::from_bytes(self.len()) };
@@ -809,16 +809,19 @@ impl<Prov: Provenance, Extra, Bytes: AllocBytes> Allocation<Prov, Extra, Bytes>
                 uninit_bytes.fill(0);
             }
         }
-        self.mark_foreign_write(full_range);
     }
 
     /// Mark all bytes in the given range as initialised and reset the provenance
     /// to wildcards. This entirely breaks the normal mechanisms for tracking
     /// initialisation and is only provided for Miri operating in native-lib
     /// mode. UB will be missed if the underlying bytes were not actually initialized.
-    pub fn mark_foreign_write(&mut self, range: AllocRange) {
+    ///
+    /// If `range` is `None`, defaults to performing this on the whole allocation.
+    pub fn process_native_write(&mut self, cx: &impl HasDataLayout, range: Option<AllocRange>) {
+        let range =
+            range.unwrap_or(AllocRange { start: Size::ZERO, size: Size::from_bytes(self.len()) });
         self.mark_init(range, true);
-        self.provenance.write_wildcards(range);
+        self.provenance.write_wildcards(cx, range);
     }
 
     /// Remove all provenance in the given memory range.

compiler/rustc_middle/src/mir/interpret/allocation/provenance_map.rs

@@ -213,21 +213,35 @@ impl<Prov: Provenance> ProvenanceMap<Prov> {
     }
 
     /// Overwrites all provenance in the given range with wildcard provenance.
+    /// Pointers partially overwritten will have their provenances
     ///
     /// Provided for usage in Miri and panics otherwise.
-    pub fn write_wildcards(&mut self, range: AllocRange) {
+    pub fn write_wildcards(&mut self, cx: &impl HasDataLayout, range: AllocRange) {
         assert!(
             Prov::OFFSET_IS_ADDR,
             "writing wildcard provenance is not supported when `OFFSET_IS_ADDR` is false"
         );
         let wildcard = Prov::WILDCARD.unwrap();
 
-        // Remove all pointer provenances, then write wildcards into the whole byte range.
-        self.ptrs.clear();
+        // Get pointer provenances that overlap with the range, then remove them.
+        let ptr_range = Self::adjusted_range_ptrs(range, cx);
+        let removed = ptr_range.filter_map(|i| self.ptrs.remove(&i).map(|p| (i, p)));
+        // Overwrite bytewise provenance.
         let bytes = self.bytes.get_or_insert_with(Box::default);
-        for offset in range.start..range.start + range.size {
+        for offset in range.start..range.end() {
             bytes.insert(offset, wildcard);
         }
+        // Now readd pointer provenances that were only partly overwritten. Always
+        // check both the start and end to allow for a write to happen in the middle
+        // of a pointer.
+        for (rm_ofs, prov) in removed {
+            for head_ofs in rm_ofs..range.start {
+                bytes.insert(head_ofs, prov);
+            }
+            for tail_ofs in range.end()..rm_ofs + cx.data_layout().pointer_size() {
+                bytes.insert(tail_ofs, prov);
+            }
+        }
     }
 }
 

src/tools/miri/src/shims/native_lib/mod.rs

@@ -251,8 +251,9 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
 
             // Prepare for possible write from native code if mutable.
             if info.mutbl.is_mut() {
-                let alloc = &mut this.get_alloc_raw_mut(alloc_id)?.0;
+                let (alloc, cx) = this.get_alloc_raw_mut(alloc_id)?;
                 alloc.prepare_for_native_access();
+                alloc.process_native_write(&cx.tcx, None);
                 // Also expose *mutable* provenance for the interpreter-level allocation.
                 std::hint::black_box(alloc.get_bytes_unchecked_raw_mut().expose_provenance());
             }