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.

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.

compiler/rustc_middle/src/mir/syntax.rs

@@ -652,6 +652,8 @@ pub enum CallSource {
     /// Other types of desugaring that did not come from the HIR, but we don't care about
     /// for diagnostics (yet).
     Misc,
+    /// Use of value, generating a clone function call
+    Use,
     /// Normal function call, no special source
     Normal,
 }

compiler/rustc_mir_build/src/builder/expr/into.rs

@@ -328,7 +328,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                             destination,
                             target: Some(success),
                             unwind: UnwindAction::Unreachable,
-                            call_source: CallSource::Misc,
+                            call_source: CallSource::Use,
                             fn_span: expr_span,
                         },
                     );

rustfmt.toml

@@ -53,4 +53,8 @@ ignore = [
     # Code automatically generated and included.
     "compiler/rustc_codegen_gcc/src/intrinsic/archs.rs",
     "compiler/rustc_codegen_gcc/example",
+
+    # Rustfmt doesn't support use closures yet
+    "tests/mir-opt/ergonomic-clones/closure.rs",
+    "tests/codegen/ergonomic-clones/closure.rs",
 ]

tests/codegen/ergonomic-clones/closure.rs

@@ -0,0 +1,55 @@
+//@ compile-flags: -C no-prepopulate-passes -Copt-level=0 -Zmir-opt-level=0
+
+#![crate_type = "lib"]
+
+#![feature(ergonomic_clones)]
+#![allow(incomplete_features)]
+
+use std::clone::UseCloned;
+
+pub fn ergonomic_clone_closure_move() -> String {
+    let s = String::from("hi");
+
+    // CHECK-NOT: ; call core::clone::impls::<impl core::clone::Clone for String>::clone
+    let cl = use || s;
+    cl()
+}
+
+#[derive(Clone)]
+struct Foo;
+
+impl UseCloned for Foo {}
+
+pub fn ergonomic_clone_closure_use_cloned() -> Foo {
+    let f = Foo;
+
+    // CHECK: ; call <closure::Foo as core::clone::Clone>::clone
+    let f1 = use || f;
+
+    // CHECK: ; call <closure::Foo as core::clone::Clone>::clone
+    let f2 = use || f;
+
+    f
+}
+
+pub fn ergonomic_clone_closure_copy() -> i32 {
+    let i = 1;
+
+    // CHECK-NOT: ; call core::clone::impls::<impl core::clone::Clone for i32>::clone
+    let i1 = use || i;
+
+    // CHECK-NOT: ; call core::clone::impls::<impl core::clone::Clone for i32>::clone
+    let i2 = use || i;
+
+    i
+}
+
+pub fn ergonomic_clone_closure_use_cloned_generics<T: UseCloned>(f: T) -> T {
+    // CHECK-NOT: ; call core::clone::impls::<impl core::clone::Clone for i32>::clone
+    let f1 = use || f;
+
+    // CHECK-NOT: ; call core::clone::impls::<impl core::clone::Clone for i32>::clone
+    let f2 = use || f;
+
+    f
+}

tests/crashes/129372.rs

@@ -1,52 +0,0 @@
-//@ known-bug: #129372
-//@ compile-flags: -Cdebuginfo=2 -Copt-level=0
-
-pub struct Wrapper<T>(T);
-struct Struct;
-
-pub trait TraitA {
-    type AssocA<'t>;
-}
-pub trait TraitB {
-    type AssocB;
-}
-
-pub fn helper(v: impl MethodTrait) {
-    let _local_that_causes_ice = v.method();
-}
-
-pub fn main() {
-    helper(Wrapper(Struct));
-}
-
-pub trait MethodTrait {
-    type Assoc<'a>;
-
-    fn method(self) -> impl for<'a> FnMut(&'a ()) -> Self::Assoc<'a>;
-}
-
-impl<T: TraitB> MethodTrait for T
-where
-    <T as TraitB>::AssocB: TraitA,
-{
-    type Assoc<'a> = <T::AssocB as TraitA>::AssocA<'a>;
-
-    fn method(self) -> impl for<'a> FnMut(&'a ()) -> Self::Assoc<'a> {
-        move |_| loop {}
-    }
-}
-
-impl<T, B> TraitB for Wrapper<B>
-where
-    B: TraitB<AssocB = T>,
-{
-    type AssocB = T;
-}
-
-impl TraitB for Struct {
-    type AssocB = Struct;
-}
-
-impl TraitA for Struct {
-    type AssocA<'t> = Self;
-}

tests/mir-opt/ergonomic-clones/closure.rs

@@ -0,0 +1,55 @@
+#![crate_type = "lib"]
+#![feature(ergonomic_clones)]
+#![allow(incomplete_features)]
+
+use std::clone::UseCloned;
+
+pub fn ergonomic_clone_closure_move() -> String {
+    // CHECK-LABEL: fn ergonomic_clone_closure_move(
+    // CHECK: _0 = move (_1.0: std::string::String);
+    // CHECK-NOT: <String as Clone>::clone
+    let s = String::from("hi");
+
+    let cl = use || s;
+    cl()
+}
+
+#[derive(Clone)]
+struct Foo;
+
+impl UseCloned for Foo {}
+
+pub fn ergonomic_clone_closure_use_cloned() -> Foo {
+    // CHECK-LABEL: fn ergonomic_clone_closure_use_cloned(
+    // CHECK: <Foo as Clone>::clone
+    let f = Foo;
+
+    let f1 = use || f;
+
+    let f2 = use || f;
+
+    f
+}
+
+pub fn ergonomic_clone_closure_copy() -> i32 {
+    // CHECK-LABEL: fn ergonomic_clone_closure_copy(
+    // CHECK: _0 = copy ((*_1).0: i32);
+    // CHECK-NOT: <i32 as Clone>::clone
+    let i = 1;
+
+    let i1 = use || i;
+
+    let i2 = use || i;
+
+    i
+}
+
+pub fn ergonomic_clone_closure_use_cloned_generics<T: UseCloned>(f: T) -> T {
+    // CHECK-LABEL: fn ergonomic_clone_closure_use_cloned_generics(
+    // CHECK: <T as Clone>::clone
+    let f1 = use || f;
+
+    let f2 = use || f;
+
+    f
+}