From a1c7f3b180cd25e79913e70e3661757272f5a174 Mon Sep 17 00:00:00 2001
From: Ben Kimock <kimockb@gmail.com>
Date: Sun, 26 Nov 2023 17:06:13 -0500
Subject: [PATCH] Sometimes return the same AllocId for a ConstAllocation

---
 .../rustc_const_eval/src/interpret/machine.rs |  5 +++++
 .../rustc_const_eval/src/interpret/operand.rs |  3 ++-
 src/tools/miri/src/machine.rs                 | 21 +++++++++++++++++++
 src/tools/miri/tests/pass/const-addrs.rs      | 19 +++++++++++++++++
 4 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 src/tools/miri/tests/pass/const-addrs.rs

diff --git a/compiler/rustc_const_eval/src/interpret/machine.rs b/compiler/rustc_const_eval/src/interpret/machine.rs
index 6617f6f2ffb68..176723301e214 100644
--- a/compiler/rustc_const_eval/src/interpret/machine.rs
+++ b/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
diff --git a/compiler/rustc_const_eval/src/interpret/operand.rs b/compiler/rustc_const_eval/src/interpret/operand.rs
index 255dd1eba97f2..313cebe0bffc8 100644
--- a/compiler/rustc_const_eval/src/interpret/operand.rs
+++ b/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)
             }
         };
diff --git a/src/tools/miri/src/machine.rs b/src/tools/miri/src/machine.rs
index 66d7dfcf3a108..56cc6279a90dc 100644
--- a/src/tools/miri/src/machine.rs
+++ b/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]
+        }
+    }
 }
diff --git a/src/tools/miri/tests/pass/const-addrs.rs b/src/tools/miri/tests/pass/const-addrs.rs
new file mode 100644
index 0000000000000..dbc2203b6351d
--- /dev/null
+++ b/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()
+}