Add -Zcodegen-emit-retag with nop retags.

Author: icmccorm

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -23,6 +23,7 @@ use super::{CachedLlbb, FunctionCx, LocalRef};
 use crate::base::{self, is_call_from_compiler_builtins_to_upstream_monomorphization};
 use crate::common::{self, IntPredicate};
 use crate::errors::CompilerBuiltinsCannotCall;
+use crate::mir::retag;
 use crate::traits::*;
 use crate::{MemFlags, meth};
 
@@ -263,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 {
@@ -1054,21 +1061,26 @@ 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 = retag::place_needs_retag(&destination);
+
             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 {
+                    dest: self.codegen_place(bx, destination.as_ref()),
+                    needs_retag,
+                }
             };
 
             let args =
@@ -2097,6 +2109,9 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         if fn_ret.is_ignore() {
             return ReturnDest::Nothing;
         }
+
+        let needs_retag = retag::place_needs_retag(&dest);
+
         let dest = if let Some(index) = dest.as_local() {
             match self.locals[index] {
                 LocalRef::Place(dest) => dest,
@@ -2110,9 +2125,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(_) => {
@@ -2135,7 +2150,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 }
         }
     }
 
@@ -2148,19 +2163,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 { dest, needs_retag } => {
+                bx.store_arg(ret_abi, llval, dest);
+                if retags_enabled && needs_retag {
+                    self.codegen_retag_place(bx, dest, 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);
@@ -2170,6 +2193,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);
             }
@@ -2181,11 +2207,28 @@ 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 {
+        /// The destination place
+        dest: PlaceRef<'tcx, V>,
+        /// If this place needs to be retagged
+        needs_retag: bool,
+    },
     /// Store an indirect return value to an operand local place.
-    IndirectOperand(PlaceRef<'tcx, V>, mir::Local),
+    IndirectOperand {
+        /// The temp place where the value is loaded from
+        tmp: PlaceRef<'tcx, V>,
+        /// The local that it is assigned to
+        index: mir::Local,
+        /// If this local needs to be retagged after the assignment
+        needs_retag: bool,
+    },
     /// Store a direct return value to an operand local place.
-    DirectOperand(mir::Local),
+    DirectOperand {
+        /// The destination local
+        index: mir::Local,
+        /// If this local needs to be retagged after the assignment.
+        needs_retag: bool,
+    },
 }
 
 fn load_cast<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(

compiler/rustc_codegen_ssa/src/mir/mod.rs

@@ -23,6 +23,7 @@ mod locals;
 pub mod naked_asm;
 pub mod operand;
 pub mod place;
+mod retag;
 mod rvalue;
 mod statement;
 
@@ -401,7 +402,7 @@ fn arg_local_refs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
         return vec![];
     }
 
-    let args = mir
+    let mut args = mir
         .args_iter()
         .enumerate()
         .map(|(arg_index, local)| {
@@ -535,6 +536,40 @@ fn arg_local_refs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
             }
         })
         .collect::<Vec<_>>();
+    if bx.tcx().sess.opts.unstable_opts.codegen_emit_retag.is_some() {
+        if let ty::InstanceKind::DropGlue(_, _) | ty::InstanceKind::AsyncDropGlue(_, _) =
+            fx.instance.def
+        {
+            // We need to special-case drop glue for now. The first argument is
+            // a raw pointer, but it needs to be treated as if it were `&mut _`.
+            args[0] = dropee_emit_retag(bx, fx, &args[0]);
+        } else {
+            args = args
+                .iter()
+                .map(|arg| match arg {
+                    &LocalRef::Place(place_ref) => {
+                        fx.codegen_retag_place(bx, place_ref, true);
+                        LocalRef::Place(place_ref)
+                    }
+                    &LocalRef::UnsizedPlace(place_ref) => {
+                        let operand = bx.load_operand(place_ref);
+                        let retagged = fx.codegen_retag_operand(bx, operand, true);
+                        assert!(matches!(retagged.val, OperandValue::Pair(_, _)));
+                        retagged.val.store(bx, place_ref);
+                        LocalRef::UnsizedPlace(place_ref)
+                    }
+                    &LocalRef::Operand(operand_ref) => {
+                        let retagged = fx.codegen_retag_operand(bx, operand_ref, true);
+                        LocalRef::Operand(retagged)
+                    }
+                    LocalRef::PendingOperand => LocalRef::PendingOperand,
+                })
+                .collect::<Vec<_>>();
+        }
+        // If we branched during retagging, then we need to update the
+        // start block to the new location.
+        fx.cached_llbbs[mir::START_BLOCK] = CachedLlbb::Some(bx.llbb());
+    }
 
     if fx.instance.def.requires_caller_location(bx.tcx()) {
         let mir_args = if let Some(num_untupled) = num_untupled {
@@ -612,3 +647,31 @@ fn find_cold_blocks<'tcx>(
 
     cold_blocks
 }
+
+fn dropee_emit_retag<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
+    bx: &mut Bx,
+    fx: &mut FunctionCx<'a, 'tcx, Bx>,
+    local: &LocalRef<'tcx, Bx::Value>,
+) -> LocalRef<'tcx, Bx::Value> {
+    // We have `*mut _` as our first argument
+    if let &LocalRef::Operand(OperandRef { val, layout, move_annotation }) = local {
+        if layout.ty.is_raw_ptr()
+            && let Some(deref_ty) = layout.ty.builtin_deref(true)
+        {
+            // Create `&mut _`
+            let lifetime = bx.tcx().lifetimes.re_erased;
+            let subst_ty_kind = ty::Ref(lifetime, deref_ty, ty::Mutability::Mut);
+            let subst_ty = bx.tcx().mk_ty_from_kind(subst_ty_kind);
+            let subst_layout = bx.layout_of(subst_ty);
+
+            // We want the same operand value, but use the reference type for it.
+            let operand_ref = OperandRef { val, layout: subst_layout, move_annotation };
+
+            let retagged = fx.codegen_retag_operand(bx, operand_ref, true);
+
+            // Return the retagged parameter, but use the original layout now.
+            return LocalRef::Operand(OperandRef { val: retagged.val, layout, move_annotation });
+        }
+    }
+    bug!("dropee isn't a raw pointer")
+}

compiler/rustc_codegen_ssa/src/mir/retag.rs

@@ -0,0 +1,39 @@
+//! Experimental support for emitting retags as function calls in generated code.
+
+use rustc_middle::mir::{Place, Rvalue, WithRetag};
+
+use crate::mir::FunctionCx;
+use crate::mir::operand::OperandRef;
+use crate::mir::place::PlaceRef;
+use crate::traits::BuilderMethods;
+
+pub(crate) fn place_needs_retag(place: &Place<'_>) -> bool {
+    // We're not interested in tracking stores to "outside" locations
+    !place.is_indirect_first_projection()
+}
+
+pub(crate) fn rvalue_needs_retag(rvalue: &Rvalue<'_>) -> bool {
+    // `Ref` has its own internal retagging
+    !matches!(rvalue, Rvalue::Ref(..)) && !matches!(rvalue, Rvalue::Use(.., WithRetag::No))
+}
+
+impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
+    /// Retags the pointers within an [`OperandRef`].
+    pub(crate) fn codegen_retag_operand(
+        &mut self,
+        _bx: &mut Bx,
+        operand: OperandRef<'tcx, Bx::Value>,
+        _is_fn_entry: bool,
+    ) -> OperandRef<'tcx, Bx::Value> {
+        operand
+    }
+
+    /// Retags the pointers within a [`PlaceRef`].
+    pub(crate) fn codegen_retag_place(
+        &mut self,
+        _bx: &mut Bx,
+        _place_ref: PlaceRef<'tcx, Bx::Value>,
+        _is_fn_entry: bool,
+    ) {
+    }
+}

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -515,7 +515,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 let mk_ref = move |tcx: TyCtxt<'tcx>, ty: Ty<'tcx>| {
                     Ty::new_ref(tcx, tcx.lifetimes.re_erased, ty, bk.to_mutbl_lossy())
                 };
-                self.codegen_place_to_pointer(bx, place, mk_ref)
+                let op = self.codegen_place_to_pointer(bx, place, mk_ref);
+                if self.cx.tcx().sess.opts.unstable_opts.codegen_emit_retag.is_some() {
+                    self.codegen_retag_operand(bx, op, false)
+                } else {
+                    op
+                }
             }
 
             mir::Rvalue::RawPtr(kind, place) => {

compiler/rustc_codegen_ssa/src/mir/statement.rs

@@ -3,6 +3,7 @@ use rustc_middle::{bug, span_bug, ty};
 use tracing::instrument;
 
 use super::{FunctionCx, LocalRef};
+use crate::mir::retag;
 use crate::traits::*;
 
 impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
@@ -12,9 +13,18 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         self.set_debug_loc(bx, statement.source_info);
         match statement.kind {
             mir::StatementKind::Assign(box (ref place, ref rvalue)) => {
+                let needs_retag = bx.tcx().sess.opts.unstable_opts.codegen_emit_retag.is_some()
+                    && retag::rvalue_needs_retag(rvalue)
+                    && retag::place_needs_retag(place);
+
                 if let Some(index) = place.as_local() {
                     match self.locals[index] {
-                        LocalRef::Place(cg_dest) => self.codegen_rvalue(bx, cg_dest, rvalue),
+                        LocalRef::Place(cg_dest) => {
+                            self.codegen_rvalue(bx, cg_dest, rvalue);
+                            if needs_retag {
+                                self.codegen_retag_place(bx, cg_dest, false);
+                            }
+                        }
                         LocalRef::UnsizedPlace(cg_indirect_dest) => {
                             let ty = cg_indirect_dest.layout.ty;
                             span_bug!(
@@ -24,7 +34,10 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                             );
                         }
                         LocalRef::PendingOperand => {
-                            let operand = self.codegen_rvalue_operand(bx, rvalue);
+                            let mut operand = self.codegen_rvalue_operand(bx, rvalue);
+                            if needs_retag {
+                                operand = self.codegen_retag_operand(bx, operand, false);
+                            }
                             self.overwrite_local(index, LocalRef::Operand(operand));
                             self.debug_introduce_local(bx, index);
                         }
@@ -39,12 +52,19 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
 
                             // If the type is zero-sized, it's already been set here,
                             // but we still need to make sure we codegen the operand
-                            self.codegen_rvalue_operand(bx, rvalue);
+                            // and emit a retag.
+                            let operand = self.codegen_rvalue_operand(bx, rvalue);
+                            if needs_retag {
+                                self.codegen_retag_operand(bx, operand, false);
+                            }
                         }
                     }
                 } else {
                     let cg_dest = self.codegen_place(bx, place.as_ref());
                     self.codegen_rvalue(bx, cg_dest, rvalue);
+                    if needs_retag {
+                        self.codegen_retag_place(bx, cg_dest, false);
+                    }
                 }
             }
             mir::StatementKind::SetDiscriminant { box ref place, variant_index } => {

compiler/rustc_interface/src/tests.rs

@@ -775,6 +775,7 @@ fn test_unstable_options_tracking_hash() {
         })
     );
     tracked!(codegen_backend, Some("abc".to_string()));
+    tracked!(codegen_emit_retag, Some(CodegenRetagOptions::default()));
     tracked!(
         coverage_options,
         CoverageOptions {