Auto merge of #139088 - spastorino:ergonomic-ref-counting-2, r=nikomatsakis

Ergonomic ref counting: optimize away clones when possible

This PR build on top of https://github.com/rust-lang/rust/pull/134797. It optimizes codegen of ergonomic ref-counting when the type being `use`d is only known to be copy after monomorphization. We avoid codening a clone and generate bitwise copy instead.

RFC: https://github.com/rust-lang/rfcs/pull/3680
Tracking issue: https://github.com/rust-lang/rust/issues/132290
Project goal: https://github.com/rust-lang/rust-project-goals/issues/107

r? `@nikomatsakis`

This PR could better sit on top of https://github.com/rust-lang/rust/pull/131650 but as it did not land yet I've decided to just do minimal changes. It may be the case that doing what I'm doing regress the performance and we may need to go the full route of https://github.com/rust-lang/rust/pull/131650.
cc `@saethlin` in this regard.

1 files changed, 72 insertions, 7 deletions

diff --git a/compiler/rustc_codegen_ssa/src/mir/mod.rs b/compiler/rustc_codegen_ssa/src/mir/mod.rs
index 0758e5d0456..6a37889217a 100644
--- a/compiler/rustc_codegen_ssa/src/mir/mod.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/mod.rs
@@ -3,7 +3,7 @@ use std::iter;
 use rustc_index::IndexVec;
 use rustc_index::bit_set::DenseBitSet;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
-use rustc_middle::mir::{Local, UnwindTerminateReason, traversal};
+use rustc_middle::mir::{Body, Local, UnwindTerminateReason, traversal};
 use rustc_middle::ty::layout::{FnAbiOf, HasTyCtxt, HasTypingEnv, TyAndLayout};
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt, TypeFoldable, TypeVisitableExt};
 use rustc_middle::{bug, mir, span_bug};
@@ -170,19 +170,29 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
 ) {
     assert!(!instance.args.has_infer());
 
+    let tcx = cx.tcx();
     let llfn = cx.get_fn(instance);
 
-    let mir = cx.tcx().instance_mir(instance.def);
+    let mut mir = tcx.instance_mir(instance.def);
 
     let fn_abi = cx.fn_abi_of_instance(instance, ty::List::empty());
     debug!("fn_abi: {:?}", fn_abi);
 
-    if cx.tcx().codegen_fn_attrs(instance.def_id()).flags.contains(CodegenFnAttrFlags::NAKED) {
+    if tcx.codegen_fn_attrs(instance.def_id()).flags.contains(CodegenFnAttrFlags::NAKED) {
         crate::mir::naked_asm::codegen_naked_asm::<Bx>(cx, &mir, instance);
         return;
     }
 
-    let debug_context = cx.create_function_debug_context(instance, fn_abi, llfn, mir);
+    if tcx.features().ergonomic_clones() {
+        let monomorphized_mir = instance.instantiate_mir_and_normalize_erasing_regions(
+            tcx,
+            ty::TypingEnv::fully_monomorphized(),
+            ty::EarlyBinder::bind(mir.clone()),
+        );
+        mir = tcx.arena.alloc(optimize_use_clone::<Bx>(cx, monomorphized_mir));
+    }
+
+    let debug_context = cx.create_function_debug_context(instance, fn_abi, llfn, &mir);
 
     let start_llbb = Bx::append_block(cx, llfn, "start");
     let mut start_bx = Bx::build(cx, start_llbb);
@@ -194,7 +204,7 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     }
 
     let cleanup_kinds =
-        base::wants_new_eh_instructions(cx.tcx().sess).then(|| analyze::cleanup_kinds(mir));
+        base::wants_new_eh_instructions(tcx.sess).then(|| analyze::cleanup_kinds(&mir));
 
     let cached_llbbs: IndexVec<mir::BasicBlock, CachedLlbb<Bx::BasicBlock>> =
         mir.basic_blocks
@@ -217,7 +227,7 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
         cleanup_kinds,
         landing_pads: IndexVec::from_elem(None, &mir.basic_blocks),
         funclets: IndexVec::from_fn_n(|_| None, mir.basic_blocks.len()),
-        cold_blocks: find_cold_blocks(cx.tcx(), mir),
+        cold_blocks: find_cold_blocks(tcx, mir),
         locals: locals::Locals::empty(),
         debug_context,
         per_local_var_debug_info: None,
@@ -233,7 +243,7 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
         fx.compute_per_local_var_debug_info(&mut start_bx).unzip();
     fx.per_local_var_debug_info = per_local_var_debug_info;
 
-    let traversal_order = traversal::mono_reachable_reverse_postorder(mir, cx.tcx(), instance);
+    let traversal_order = traversal::mono_reachable_reverse_postorder(mir, tcx, instance);
     let memory_locals = analyze::non_ssa_locals(&fx, &traversal_order);
 
     // Allocate variable and temp allocas
@@ -310,6 +320,61 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     }
 }
 
+// FIXME: Move this function to mir::transform when post-mono MIR passes land.
+fn optimize_use_clone<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
+    cx: &'a Bx::CodegenCx,
+    mut mir: Body<'tcx>,
+) -> Body<'tcx> {
+    let tcx = cx.tcx();
+
+    if tcx.features().ergonomic_clones() {
+        for bb in mir.basic_blocks.as_mut() {
+            let mir::TerminatorKind::Call {
+                args,
+                destination,
+                target,
+                call_source: mir::CallSource::Use,
+                ..
+            } = &bb.terminator().kind
+            else {
+                continue;
+            };
+
+            // CallSource::Use calls always use 1 argument.
+            assert_eq!(args.len(), 1);
+            let arg = &args[0];
+
+            // These types are easily available from locals, so check that before
+            // doing DefId lookups to figure out what we're actually calling.
+            let arg_ty = arg.node.ty(&mir.local_decls, tcx);
+
+            let ty::Ref(_region, inner_ty, mir::Mutability::Not) = *arg_ty.kind() else { continue };
+
+            if !tcx.type_is_copy_modulo_regions(cx.typing_env(), inner_ty) {
+                continue;
+            }
+
+            let Some(arg_place) = arg.node.place() else { continue };
+
+            let destination_block = target.unwrap();
+
+            bb.statements.push(mir::Statement {
+                source_info: bb.terminator().source_info,
+                kind: mir::StatementKind::Assign(Box::new((
+                    *destination,
+                    mir::Rvalue::Use(mir::Operand::Copy(
+                        arg_place.project_deeper(&[mir::ProjectionElem::Deref], tcx),
+                    )),
+                ))),
+            });
+
+            bb.terminator_mut().kind = mir::TerminatorKind::Goto { target: destination_block };
+        }
+    }
+
+    mir
+}
+
 /// Produces, for each argument, a `Value` pointing at the
 /// argument's value. As arguments are places, these are always
 /// indirect.