Sometimes return the same AllocId for a ConstAllocation

Author: saethlin

compiler/rustc_const_eval/src/interpret/machine.rs

@@ -501,6 +501,11 @@ pub trait Machine<'mir, 'tcx: 'mir>: Sized {
     ) -> InterpResult<'tcx> {
         Ok(())
     }
+
+    #[inline(always)]
+    fn const_alloc_id(ecx: &InterpCx<'mir, 'tcx, Self>, alloc: ConstAllocation<'tcx>) -> AllocId {
+        ecx.tcx.reserve_and_set_memory_alloc(alloc)
+    }
 }
 
 /// A lot of the flexibility above is just needed for `Miri`, but all "compile-time" machines

compiler/rustc_const_eval/src/interpret/operand.rs

@@ -760,7 +760,8 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             mir::ConstValue::Slice { data, meta } => {
                 // We rely on mutability being set correctly in `data` to prevent writes
                 // where none should happen.
-                let ptr = Pointer::new(self.tcx.reserve_and_set_memory_alloc(data), Size::ZERO);
+                let alloc_id = M::const_alloc_id(self, data);
+                let ptr = Pointer::new(alloc_id, Size::ZERO);
                 Immediate::new_slice(self.global_base_pointer(ptr)?.into(), meta, self)
             }
         };

src/tools/miri/src/machine.rs

@@ -10,6 +10,7 @@ use std::process;
 use either::Either;
 use rand::rngs::StdRng;
 use rand::SeedableRng;
+use rand::Rng;
 
 use rustc_ast::ast::Mutability;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
@@ -531,6 +532,8 @@ pub struct MiriMachine<'mir, 'tcx> {
     /// The spans we will use to report where an allocation was created and deallocated in
     /// diagnostics.
     pub(crate) allocation_spans: RefCell<FxHashMap<AllocId, (Span, Option<Span>)>>,
+
+    const_cache: RefCell<FxHashMap<(ConstAllocation<'tcx>, Instance<'tcx>), Vec<AllocId>>>,
 }
 
 impl<'mir, 'tcx> MiriMachine<'mir, 'tcx> {
@@ -656,6 +659,7 @@ impl<'mir, 'tcx> MiriMachine<'mir, 'tcx> {
             stack_size,
             collect_leak_backtraces: config.collect_leak_backtraces,
             allocation_spans: RefCell::new(FxHashMap::default()),
+            const_cache: RefCell::new(FxHashMap::default()),
         }
     }
 
@@ -841,6 +845,7 @@ impl VisitProvenance for MiriMachine<'_, '_> {
             stack_size: _,
             collect_leak_backtraces: _,
             allocation_spans: _,
+            const_cache: _,
         } = self;
 
         threads.visit_provenance(visit);
@@ -1450,4 +1455,20 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
         ecx.machine.allocation_spans.borrow_mut().insert(alloc_id, (span, None));
         Ok(())
     }
+
+    fn const_alloc_id(
+        ecx: &InterpCx<'mir, 'tcx, Self>,
+        alloc: ConstAllocation<'tcx>,
+    ) -> AllocId {
+        let instance = ecx.active_thread_stack().last().unwrap().instance;
+        let mut cache = ecx.machine.const_cache.borrow_mut();
+        let contents = cache.entry((alloc, instance)).or_insert_with(Vec::new);
+        if contents.len() < 16 {
+            let new = ecx.tcx.reserve_and_set_memory_alloc(alloc);
+            contents.push(new);
+            new
+        } else {
+            contents[ecx.machine.rng.borrow_mut().gen::<u8>() as usize % 16]
+        }
+    }
 }

src/tools/miri/tests/pass/const-addrs.rs

@@ -0,0 +1,19 @@
+#![feature(strict_provenance)]
+
+const EVALS: usize = 256;
+
+use std::collections::HashSet;
+fn main() {
+    let mut addrs = HashSet::new();
+    for _ in 0..EVALS {
+        addrs.insert(get_const_addr());
+    }
+    // Check that the const allocation has multiple base addresses
+    assert!(addrs.len() > 1);
+    // But also that we get a limited number of unique base addresses
+    assert!(addrs.len() < EVALS);
+}
+
+fn get_const_addr() -> usize {
+    "test".as_bytes().as_ptr().addr()
+}