Set !captures metadata for store of &Freeze

When a `&Freeze` reference is stored (with retag), emit
`!captures !{!"address", !"read_provenance"}` metadata to indicate
that it's UB to write through the stored pointer.

Author: nikic

compiler/rustc_codegen_llvm/src/builder.rs

@@ -854,6 +854,18 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
                     self.set_metadata_node(store, llvm::MD_nontemporal, &[one]);
                 }
             }
+            if flags.contains(MemFlags::CAPTURES_READ_ONLY)
+                && crate::llvm_util::get_version() >= (22, 0, 0)
+            {
+                let args = [
+                    self.cx.create_metadata(b"address"),
+                    self.cx.create_metadata(b"read_provenance"),
+                ];
+                // FIXME: Switch this to use MD_captures once LLVM 22 is the minimum.
+                let id = self.get_md_kind_id("captures");
+                let md = llvm::LLVMMDNodeInContext2(self.cx.llcx, args.as_ptr(), args.len());
+                self.set_metadata(store, id, md);
+            }
             store
         }
     }

compiler/rustc_codegen_ssa/src/lib.rs

@@ -167,6 +167,7 @@ bitflags::bitflags! {
         const VOLATILE = 1 << 0;
         const NONTEMPORAL = 1 << 1;
         const UNALIGNED = 1 << 2;
+        const CAPTURES_READ_ONLY = 1 << 3;
     }
 }
 

compiler/rustc_codegen_ssa/src/mir/operand.rs

@@ -296,11 +296,21 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
         self,
         bx: &mut Bx,
         dest: PlaceRef<'tcx, V>,
+    ) {
+        self.store_with_annotation_and_flags(bx, dest, MemFlags::empty())
+    }
+
+    /// Same as store_with_annotation(), but also specify flags for the store.
+    pub fn store_with_annotation_and_flags<Bx: BuilderMethods<'a, 'tcx, Value = V>>(
+        self,
+        bx: &mut Bx,
+        dest: PlaceRef<'tcx, V>,
+        flags: MemFlags,
     ) {
         if let Some(instance) = self.move_annotation {
-            bx.with_move_annotation(instance, |bx| self.val.store(bx, dest))
+            bx.with_move_annotation(instance, |bx| self.val.store_with_flags(bx, dest, flags))
         } else {
-            self.val.store(bx, dest)
+            self.val.store_with_flags(bx, dest, flags)
         }
     }
 
@@ -961,7 +971,8 @@ impl<'a, 'tcx, V: CodegenObject> OperandValue<V> {
                 let llptr = bx.inbounds_ptradd(dest.val.llval, bx.const_usize(b_offset.bytes()));
                 let val = bx.from_immediate(b);
                 let align = dest.val.align.restrict_for_offset(b_offset);
-                bx.store_with_flags(val, llptr, align, flags);
+                // The CAPTURES_READ_ONLY flag only applies to the first element.
+                bx.store_with_flags(val, llptr, align, flags & !MemFlags::CAPTURES_READ_ONLY);
             }
         }
     }

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -2,7 +2,7 @@ use itertools::Itertools as _;
 use rustc_abi::{self as abi, BackendRepr, FIRST_VARIANT};
 use rustc_middle::ty::adjustment::PointerCoercion;
 use rustc_middle::ty::layout::{HasTyCtxt, HasTypingEnv, LayoutOf, TyAndLayout};
-use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
+use rustc_middle::ty::{self, Instance, Mutability, Ty, TyCtxt};
 use rustc_middle::{bug, mir, span_bug};
 use rustc_session::config::OptLevel;
 use tracing::{debug, instrument};
@@ -23,7 +23,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         rvalue: &mir::Rvalue<'tcx>,
     ) {
         match *rvalue {
-            mir::Rvalue::Use(ref operand, _) => {
+            mir::Rvalue::Use(ref operand, with_retag) => {
                 if let mir::Operand::Constant(const_op) = operand {
                     let val = self.eval_mir_constant(&const_op);
                     if val.all_bytes_uninit(self.cx.tcx()) {
@@ -40,9 +40,20 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 ) {
                     debug_assert!(!matches!(cg_operand.val, OperandValue::Ref(..)));
                 }
+                // If this is storing a &Freeze reference with a retag, record that it's not
+                // possible to perform writes through the stored pointer.
+                let flags = if let ty::Ref(_, pointee_ty, Mutability::Not) =
+                    operand.ty(self.mir, self.cx.tcx()).kind()
+                    && with_retag.yes()
+                    && pointee_ty.is_freeze(self.cx.tcx(), self.cx.typing_env())
+                {
+                    MemFlags::CAPTURES_READ_ONLY
+                } else {
+                    MemFlags::empty()
+                };
                 // FIXME: consider not copying constants through stack. (Fixable by codegen'ing
                 // constants into `OperandValue::Ref`; why don’t we do that yet if we don’t?)
-                cg_operand.store_with_annotation(bx, dest);
+                cg_operand.store_with_annotation_and_flags(bx, dest, flags);
             }
 
             mir::Rvalue::Cast(

tests/codegen-llvm/store-captures.rs

@@ -0,0 +1,68 @@
+//@ min-llvm-version: 22
+//@ compile-flags: -O -C no-prepopulate-passes
+
+#![crate_type = "lib"]
+
+use std::cell::UnsafeCell;
+
+#[unsafe(no_mangle)]
+pub fn store_ptr<'a>(x: *const i32, y: &mut *const i32) {
+    // CHECK-LABEL: define void @store_ptr
+    // CHECK-NOT: !captures
+    *y = x;
+}
+
+#[unsafe(no_mangle)]
+pub fn store_ref_freeze<'a>(x: &'a i32, y: &mut &'a i32) {
+    // CHECK-LABEL: define void @store_ref_freeze
+    // CHECK: store ptr %x, ptr %y, {{.*}}, !captures ![[CAPTURES:[0-9]+]]
+    *y = x;
+}
+
+#[unsafe(no_mangle)]
+pub fn store_ref_not_freeze<'a>(x: &'a UnsafeCell<i32>, y: &mut &'a UnsafeCell<i32>) {
+    // CHECK-LABEL: define void @store_ref_not_freeze
+    // CHECK-NOT: !captures
+    *y = x;
+}
+
+#[unsafe(no_mangle)]
+pub fn store_mut_ref<'a>(x: &'a mut i32, y: &mut &'a mut i32) {
+    // CHECK-LABEL: define void @store_mut_ref
+    // CHECK-NOT: !captures
+    *y = x;
+}
+
+#[unsafe(no_mangle)]
+pub fn store_mut_ref_as_shared_ref<'a>(x: &'a mut i32, y: &mut &'a i32) {
+    // CHECK-LABEL: define void @store_mut_ref_as_shared_ref
+    // CHECK: store ptr %x, ptr %y, {{.*}}, !captures ![[CAPTURES:[0-9]+]]
+    *y = x;
+}
+
+#[unsafe(no_mangle)]
+pub fn store_mut_ref_as_ptr(x: &mut i32, y: &mut *const i32) {
+    // CHECK-LABEL: define void @store_mut_ref_as_ptr
+    // CHECK-NOT: !captures
+    *y = x;
+}
+
+#[unsafe(no_mangle)]
+pub fn store_slice<'a>(x: &'a [i32], y: &mut &'a [i32]) {
+    // CHECK-LABEL: define void @store_slice
+    // CHECK: store ptr %x.0, ptr %y, {{.*}}, !captures ![[CAPTURES:[0-9]+]]
+    // Second store is slice size, can't have !captures
+    // CHECK-NOT: !captures
+    *y = x;
+}
+
+#[unsafe(no_mangle)]
+pub fn store_dyn<'a>(x: &'a dyn Drop, y: &mut &'a dyn Drop) {
+    // dyn trait is not known Freeze. The vtable could use !captures,
+    // but that's probably not particularly useful.
+    // CHECK-LABEL: define void @store_dyn
+    // CHECK-NOT: !captures
+    *y = x;
+}
+
+// CHECK: ![[CAPTURES]] = !{!"address", !"read_provenance"}