3 files changed, 47 insertions, 13 deletions

diff --git a/compiler/rustc_middle/src/mir/interpret/pointer.rs b/compiler/rustc_middle/src/mir/interpret/pointer.rs
index 6cfd07d699c..1700b0f02ec 100644
--- a/compiler/rustc_middle/src/mir/interpret/pointer.rs
+++ b/compiler/rustc_middle/src/mir/interpret/pointer.rs
@@ -80,14 +80,23 @@ pub trait Provenance: Copy + fmt::Debug + 'static {
 }
 
 /// The type of provenance in the compile-time interpreter.
-/// This is a packed representation of an `AllocId` and an `immutable: bool`.
+/// This is a packed representation of:
+/// - an `AllocId` (non-zero)
+/// - an `immutable: bool`
+/// - a `shared_ref: bool`
+///
+/// with the extra invariant that if `immutable` is `true`, then so
+/// is `shared_ref`.
 #[derive(Copy, Clone, Eq, Hash, Ord, PartialEq, PartialOrd)]
 pub struct CtfeProvenance(NonZero<u64>);
 
 impl From<AllocId> for CtfeProvenance {
     fn from(value: AllocId) -> Self {
         let prov = CtfeProvenance(value.0);
-        assert!(!prov.immutable(), "`AllocId` with the highest bit set cannot be used in CTFE");
+        assert!(
+            prov.alloc_id() == value,
+            "`AllocId` with the highest bits set cannot be used in CTFE"
+        );
         prov
     }
 }
@@ -103,12 +112,14 @@ impl fmt::Debug for CtfeProvenance {
 }
 
 const IMMUTABLE_MASK: u64 = 1 << 63; // the highest bit
+const SHARED_REF_MASK: u64 = 1 << 62;
+const ALLOC_ID_MASK: u64 = u64::MAX & !IMMUTABLE_MASK & !SHARED_REF_MASK;
 
 impl CtfeProvenance {
     /// Returns the `AllocId` of this provenance.
     #[inline(always)]
     pub fn alloc_id(self) -> AllocId {
-        AllocId(NonZero::new(self.0.get() & !IMMUTABLE_MASK).unwrap())
+        AllocId(NonZero::new(self.0.get() & ALLOC_ID_MASK).unwrap())
     }
 
     /// Returns whether this provenance is immutable.
@@ -117,10 +128,38 @@ impl CtfeProvenance {
         self.0.get() & IMMUTABLE_MASK != 0
     }
 
+    /// Returns whether this provenance is derived from a shared reference.
+    #[inline]
+    pub fn shared_ref(self) -> bool {
+        self.0.get() & SHARED_REF_MASK != 0
+    }
+
+    pub fn into_parts(self) -> (AllocId, bool, bool) {
+        (self.alloc_id(), self.immutable(), self.shared_ref())
+    }
+
+    pub fn from_parts((alloc_id, immutable, shared_ref): (AllocId, bool, bool)) -> Self {
+        let prov = CtfeProvenance::from(alloc_id);
+        if immutable {
+            // This sets both flags, so we don't even have to check `shared_ref`.
+            prov.as_immutable()
+        } else if shared_ref {
+            prov.as_shared_ref()
+        } else {
+            prov
+        }
+    }
+
     /// Returns an immutable version of this provenance.
     #[inline]
     pub fn as_immutable(self) -> Self {
-        CtfeProvenance(self.0 | IMMUTABLE_MASK)
+        CtfeProvenance(self.0 | IMMUTABLE_MASK | SHARED_REF_MASK)
+    }
+
+    /// Returns a "shared reference" (but not necessarily immutable!) version of this provenance.
+    #[inline]
+    pub fn as_shared_ref(self) -> Self {
+        CtfeProvenance(self.0 | SHARED_REF_MASK)
     }
 }
 
diff --git a/compiler/rustc_middle/src/ty/codec.rs b/compiler/rustc_middle/src/ty/codec.rs
index 46203ee150f..7e533bc4291 100644
--- a/compiler/rustc_middle/src/ty/codec.rs
+++ b/compiler/rustc_middle/src/ty/codec.rs
@@ -165,8 +165,7 @@ impl<'tcx, E: TyEncoder<I = TyCtxt<'tcx>>> Encodable<E> for AllocId {
 
 impl<'tcx, E: TyEncoder<I = TyCtxt<'tcx>>> Encodable<E> for CtfeProvenance {
     fn encode(&self, e: &mut E) {
-        self.alloc_id().encode(e);
-        self.immutable().encode(e);
+        self.into_parts().encode(e);
     }
 }
 
@@ -295,10 +294,8 @@ impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> Decodable<D> for AllocId {
 
 impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> Decodable<D> for CtfeProvenance {
     fn decode(decoder: &mut D) -> Self {
-        let alloc_id: AllocId = Decodable::decode(decoder);
-        let prov = CtfeProvenance::from(alloc_id);
-        let immutable: bool = Decodable::decode(decoder);
-        if immutable { prov.as_immutable() } else { prov }
+        let parts = Decodable::decode(decoder);
+        CtfeProvenance::from_parts(parts)
     }
 }
 
diff --git a/compiler/rustc_middle/src/ty/impls_ty.rs b/compiler/rustc_middle/src/ty/impls_ty.rs
index b5b7b8bcfef..5f6305bb48a 100644
--- a/compiler/rustc_middle/src/ty/impls_ty.rs
+++ b/compiler/rustc_middle/src/ty/impls_ty.rs
@@ -75,11 +75,9 @@ impl<'a> HashStable<StableHashingContext<'a>> for mir::interpret::AllocId {
     }
 }
 
-// CtfeProvenance is an AllocId and a bool.
 impl<'a> HashStable<StableHashingContext<'a>> for mir::interpret::CtfeProvenance {
     fn hash_stable(&self, hcx: &mut StableHashingContext<'a>, hasher: &mut StableHasher) {
-        self.alloc_id().hash_stable(hcx, hasher);
-        self.immutable().hash_stable(hcx, hasher);
+        self.into_parts().hash_stable(hcx, hasher);
     }
 }