Emit retags in codegen to support BorrowSanitizer.

compiler/rustc_codegen_gcc/src/builder.rs

@@ -562,6 +562,14 @@ impl<'a, 'gcc, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'gcc, 'tcx> {
         self.llbb().end_with_conditional(self.location, cond, then_block, else_block)
     }
 
+    fn phi(
+        &mut self,
+        _typ: Type<'gcc>,
+        _cases: impl ExactSizeIterator<Item = (Self::BasicBlock, Self::Value)>,
+    ) -> Self::Value {
+        unimplemented!()
+    }
+
     fn switch(
         &mut self,
         value: RValue<'gcc>,

compiler/rustc_codegen_gcc/src/intrinsic/mod.rs

@@ -8,7 +8,6 @@ use std::iter;
 use gccjit::Type;
 use gccjit::{ComparisonOp, Function, FunctionType, RValue, ToRValue, UnaryOp};
 use rustc_abi::{BackendRepr, HasDataLayout, WrappingRange};
-use rustc_codegen_ssa::MemFlags;
 use rustc_codegen_ssa::base::wants_msvc_seh;
 use rustc_codegen_ssa::common::IntPredicate;
 use rustc_codegen_ssa::errors::InvalidMonomorphization;
@@ -20,6 +19,7 @@ use rustc_codegen_ssa::traits::{
     ArgAbiBuilderMethods, BaseTypeCodegenMethods, BuilderMethods, ConstCodegenMethods,
     IntrinsicCallBuilderMethods, LayoutTypeCodegenMethods,
 };
+use rustc_codegen_ssa::{MemFlags, RetagInfo};
 use rustc_data_structures::fx::FxHashSet;
 #[cfg(feature = "master")]
 use rustc_middle::ty::layout::FnAbiOf;
@@ -729,6 +729,14 @@ impl<'a, 'gcc, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'a, 'gcc, 'tc
         // FIXME(antoyo): implement.
         self.context.new_rvalue_from_int(self.int_type, 0)
     }
+
+    fn retag_reg(&mut self, _ptr: Self::Value, _info: RetagInfo<Self::Value>) -> Self::Value {
+        unimplemented!();
+    }
+
+    fn retag_mem(&mut self, _place: Self::Value, _info: RetagInfo<Self::Value>) {
+        unimplemented!();
+    }
 }
 
 impl<'a, 'gcc, 'tcx> ArgAbiBuilderMethods<'tcx> for Builder<'a, 'gcc, 'tcx> {

compiler/rustc_codegen_llvm/src/builder.rs

@@ -347,6 +347,18 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         }
     }
 
+    fn phi(
+        &mut self,
+        ty: &'ll Type,
+        cases: impl ExactSizeIterator<Item = (Self::BasicBlock, Self::Value)>,
+    ) -> Self::Value {
+        let phi = unsafe { llvm::LLVMBuildPhi(self.llbuilder, ty, UNNAMED) };
+        for (bb, value) in cases {
+            self.add_incoming_to_phi(phi, value, bb);
+        }
+        phi
+    }
+
     fn switch(
         &mut self,
         v: &'ll Value,

compiler/rustc_codegen_llvm/src/context.rs

@@ -962,19 +962,36 @@ impl<'ll> CodegenCx<'ll, '_> {
         // This isn't an "LLVM intrinsic", but LLVM's optimization passes
         // recognize it like one (including turning it into `bcmp` sometimes)
         // and we use it to implement intrinsics like `raw_eq` and `compare_bytes`
-        if base_name == "memcmp" {
-            let fn_ty = self
-                .type_func(&[self.type_ptr(), self.type_ptr(), self.type_isize()], self.type_int());
-            let f = self.declare_cfn("memcmp", llvm::UnnamedAddr::No, fn_ty);
-
-            return (fn_ty, f);
+        match base_name {
+            "memcmp" => {
+                let fn_ty = self.type_func(
+                    &[self.type_ptr(), self.type_ptr(), self.type_isize()],
+                    self.type_int(),
+                );
+                let f = self.declare_cfn("memcmp", llvm::UnnamedAddr::No, fn_ty);
+                (fn_ty, f)
+            }
+            "__rust_retag_reg" => {
+                let fn_ty = self.type_func(type_params, self.type_ptr());
+                let llfn = self.declare_cfn(base_name, llvm::UnnamedAddr::No, fn_ty);
+                let nounwind = llvm::AttributeKind::NoUnwind.create_attr(self.llcx);
+                attributes::apply_to_llfn(llfn, llvm::AttributePlace::Function, &[nounwind]);
+                (fn_ty, llfn)
+            }
+            "__rust_retag_mem" => {
+                let fn_ty = self.type_func(type_params, self.type_void());
+                let llfn = self.declare_cfn(base_name, llvm::UnnamedAddr::No, fn_ty);
+                let nounwind = llvm::AttributeKind::NoUnwind.create_attr(self.llcx);
+                attributes::apply_to_llfn(llfn, llvm::AttributePlace::Function, &[nounwind]);
+                (fn_ty, llfn)
+            }
+            _ => {
+                let intrinsic = llvm::Intrinsic::lookup(base_name.as_bytes())
+                    .unwrap_or_else(|| bug!("Unknown intrinsic: `{base_name}`"));
+                let f = intrinsic.get_declaration(self.llmod, &type_params);
+                (self.get_type_of_global(f), f)
+            }
         }
-
-        let intrinsic = llvm::Intrinsic::lookup(base_name.as_bytes())
-            .unwrap_or_else(|| bug!("Unknown intrinsic: `{base_name}`"));
-        let f = intrinsic.get_declaration(self.llmod, &type_params);
-
-        (self.get_type_of_global(f), f)
     }
 
     pub(crate) fn eh_catch_typeinfo(&self) -> &'ll Value {

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -6,6 +6,7 @@ use rustc_abi::{
     AddressSpace, Align, BackendRepr, Float, HasDataLayout, Integer, NumScalableVectors, Primitive,
     Size, WrappingRange,
 };
+use rustc_codegen_ssa::RetagInfo;
 use rustc_codegen_ssa::base::{compare_simd_types, wants_msvc_seh, wants_wasm_eh};
 use rustc_codegen_ssa::common::{IntPredicate, TypeKind};
 use rustc_codegen_ssa::errors::{ExpectedPointerMutability, InvalidMonomorphization};
@@ -1015,6 +1016,29 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
     fn va_end(&mut self, va_list: &'ll Value) -> &'ll Value {
         self.call_intrinsic("llvm.va_end", &[self.val_ty(va_list)], &[va_list])
     }
+
+    fn retag_reg(&mut self, ptr: Self::Value, info: RetagInfo<Self::Value>) -> Self::Value {
+        codegen_retag_inner(self, "__rust_retag_reg", ptr, info)
+    }
+
+    fn retag_mem(&mut self, ptr: Self::Value, info: RetagInfo<Self::Value>) {
+        codegen_retag_inner(self, "__rust_retag_mem", ptr, info);
+    }
+}
+
+fn codegen_retag_inner<'ll, 'tcx>(
+    bx: &mut Builder<'_, 'll, 'tcx>,
+    name: &'static str,
+    ptr: &'ll Value,
+    info: RetagInfo<&'ll Value>,
+) -> &'ll Value {
+    let size = bx.const_usize(info.size.bytes());
+    let perms = bx.const_u8(info.flags.bits());
+    bx.call_intrinsic(
+        name,
+        &[bx.type_ptr(), bx.val_ty(size), bx.type_i8(), bx.type_ptr(), bx.type_ptr()],
+        &[ptr, size, perms, info.im_layout, info.pin_layout],
+    )
 }
 
 fn llvm_arch_for(rust_arch: &Arch) -> Option<&'static str> {

compiler/rustc_codegen_ssa/src/lib.rs

@@ -17,6 +17,7 @@ use std::io;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
+use rustc_abi::Size;
 use rustc_data_structures::fx::{FxHashSet, FxIndexMap};
 use rustc_data_structures::unord::UnordMap;
 use rustc_hir::CRATE_HIR_ID;
@@ -172,6 +173,37 @@ bitflags::bitflags! {
     }
 }
 
+#[derive(Debug, Copy, Clone)]
+pub struct RetagInfo<V> {
+    /// The size of the initial range within the allocation that is
+    /// associated with the permission created by the retag.
+    pub size: Size,
+    /// Encoded type information used to determine the kind of permission
+    /// created by the retag.
+    pub flags: RetagFlags,
+    /// A constant array of (offset, size) pairs describing
+    /// the ranges covered by `UnsafeCell` within the pointee type.
+    pub im_layout: V,
+    /// A constant array of (offset, size) pairs describing
+    /// the ranges covered by `UnsafePinned` within the pointee type.
+    pub pin_layout: V,
+}
+
+bitflags::bitflags! {
+    #[repr(C)]
+    #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
+    pub struct RetagFlags: u8 {
+        /// If this is a function-entry retag.
+        const IS_PROTECTED = 1 << 0;
+        /// If this is a mutable reference or a `Box`.
+        const IS_MUTABLE = 1 << 1;
+        /// If this is a `Box`.
+        const IS_BOX = 1 << 2;
+        /// If the pointee type is `Freeze`
+        const IS_FREEZE = 1 << 3;
+    }
+}
+
 // This is the same as `rustc_session::cstore::NativeLib`, except:
 // - (important) the `foreign_module` field is missing, because it contains a `DefId`, which can't
 //   be encoded with `FileEncoder`.

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -264,6 +264,12 @@ impl<'a, 'tcx> TerminatorCodegenHelper<'tcx> {
                     bx.lifetime_end(tmp, size);
                 }
                 fx.store_return(bx, ret_dest, &fn_abi.ret, invokeret);
+
+                // If the return value has variants that needed to be retagged,
+                // then we might be in a different basic block now.
+                // Update the cached block for `target` to point to this new
+                // block, where codegen will continue.
+                fx.cached_llbbs[target] = CachedLlbb::Some(bx.llbb());
             }
             MergingSucc::False
         } else {
@@ -1055,21 +1061,23 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             let result_layout =
                 self.cx.layout_of(self.monomorphized_place_ty(destination.as_ref()));
 
+            let needs_retag = !destination.is_indirect_first_projection();
+
             let return_dest = if result_layout.is_zst() {
                 ReturnDest::Nothing
             } else if let Some(index) = destination.as_local() {
                 match self.locals[index] {
-                    LocalRef::Place(dest) => ReturnDest::Store(dest),
+                    LocalRef::Place(dest) => ReturnDest::Store(dest, needs_retag),
                     LocalRef::UnsizedPlace(_) => bug!("return type must be sized"),
                     LocalRef::PendingOperand => {
                         // Handle temporary places, specifically `Operand` ones, as
                         // they don't have `alloca`s.
-                        ReturnDest::DirectOperand(index)
+                        ReturnDest::DirectOperand(index, needs_retag)
                     }
                     LocalRef::Operand(_) => bug!("place local already assigned to"),
                 }
             } else {
-                ReturnDest::Store(self.codegen_place(bx, destination.as_ref()))
+                ReturnDest::Store(self.codegen_place(bx, destination.as_ref()), needs_retag)
             };
 
             let args =
@@ -2039,6 +2047,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         if fn_ret.is_ignore() {
             return ReturnDest::Nothing;
         }
+        let needs_retag = !dest.is_indirect_first_projection();
         let dest = if let Some(index) = dest.as_local() {
             match self.locals[index] {
                 LocalRef::Place(dest) => dest,
@@ -2052,9 +2061,9 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         let tmp = PlaceRef::alloca(bx, fn_ret.layout);
                         tmp.storage_live(bx);
                         llargs.push(tmp.val.llval);
-                        ReturnDest::IndirectOperand(tmp, index)
+                        ReturnDest::IndirectOperand(tmp, index, needs_retag)
                     } else {
-                        ReturnDest::DirectOperand(index)
+                        ReturnDest::DirectOperand(index, needs_retag)
                     };
                 }
                 LocalRef::Operand(_) => {
@@ -2077,7 +2086,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             llargs.push(dest.val.llval);
             ReturnDest::Nothing
         } else {
-            ReturnDest::Store(dest)
+            ReturnDest::Store(dest, needs_retag)
         }
     }
 
@@ -2090,19 +2099,27 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         llval: Bx::Value,
     ) {
         use self::ReturnDest::*;
-
+        let retags_enabled = bx.tcx().sess.opts.unstable_opts.codegen_emit_retag.is_some();
         match dest {
             Nothing => (),
-            Store(dst) => bx.store_arg(ret_abi, llval, dst),
-            IndirectOperand(tmp, index) => {
-                let op = bx.load_operand(tmp);
+            Store(dst, needs_retag) => {
+                bx.store_arg(ret_abi, llval, dst);
+                if retags_enabled && needs_retag {
+                    self.codegen_retag_place(bx, dst, false);
+                }
+            }
+            IndirectOperand(tmp, index, needs_retag) => {
+                let mut op = bx.load_operand(tmp);
+                if retags_enabled && needs_retag {
+                    op = self.codegen_retag_operand(bx, op, false);
+                }
                 tmp.storage_dead(bx);
                 self.overwrite_local(index, LocalRef::Operand(op));
                 self.debug_introduce_local(bx, index);
             }
-            DirectOperand(index) => {
+            DirectOperand(index, needs_retag) => {
                 // If there is a cast, we have to store and reload.
-                let op = if let PassMode::Cast { .. } = ret_abi.mode {
+                let mut op = if let PassMode::Cast { .. } = ret_abi.mode {
                     let tmp = PlaceRef::alloca(bx, ret_abi.layout);
                     tmp.storage_live(bx);
                     bx.store_arg(ret_abi, llval, tmp);
@@ -2112,6 +2129,9 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 } else {
                     OperandRef::from_immediate_or_packed_pair(bx, llval, ret_abi.layout)
                 };
+                if retags_enabled && needs_retag {
+                    op = self.codegen_retag_operand(bx, op, false);
+                }
                 self.overwrite_local(index, LocalRef::Operand(op));
                 self.debug_introduce_local(bx, index);
             }
@@ -2123,11 +2143,11 @@ enum ReturnDest<'tcx, V> {
     /// Do nothing; the return value is indirect or ignored.
     Nothing,
     /// Store the return value to the pointer.
-    Store(PlaceRef<'tcx, V>),
+    Store(PlaceRef<'tcx, V>, bool),
     /// Store an indirect return value to an operand local place.
-    IndirectOperand(PlaceRef<'tcx, V>, mir::Local),
+    IndirectOperand(PlaceRef<'tcx, V>, mir::Local, bool),
     /// Store a direct return value to an operand local place.
-    DirectOperand(mir::Local),
+    DirectOperand(mir::Local, bool),
 }
 
 fn load_cast<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(