InlineFutureIntoFuture is a peephole optimization improving UpvarToLocalProp.

Problem Overview
----------------

When `async fn` is desugared, there's a whole bunch of local-to-local moves that
are easily identified and eliminated. However, there is one exception: the
sugaring of `.await` does `a = IntoFuture::into_future(b);`, and that is no longer obviously a move from the viewpoint of the analysis.

However, for all F: Future, `<F as IntoFuture>::into_future(self)` is
"guaranteed" to be the identity function that returns `self`.

So: this matches `a = <impl Future as IntoFuture>::into_future(b);` and replaces
it with `a = b;`, based on reasoning that libcore's blanket implementation of
IntoFuture for impl Future is an identity function that takes `self` by value.

This transformation, in tandem with UpvarToLocalProp, is enough to address both
case 1 and case 2 of Rust issue 62958.

InlineFutureIntoFuture respects optimization fuel, same as UpvarToLocalProp
(much simpler to implement in this case).

inline-future-into-future: improved comments during code walk for a rubber duck.

MERGEME inline_future_into_future revised internal instrumentation to print out arg0 type

(because that is what is really going to matter and I should be doing more to
let it drive the analysis.)

Updates
-------

respect -Zmir_opt_level=0

Author: pnkfelix

compiler/rustc_mir_transform/src/inline_future_into_future.rs

@@ -0,0 +1,162 @@
+//! Converts `y = <Future as IntoFuture>::into_future(x);` into just `y = x;`,
+//! since we "know" that matches the behavior of the blanket implementation of
+//! IntoFuture for F where F: Future.
+//!
+//! FIXME: determine such coalescing is sound. In particular, check whether
+//! specialization could foil our plans here!
+//!
+//! This is meant to enhance the effectiveness of the upvar-to-local-prop
+//! transformation in reducing the size of the generators constructed by the
+//! compiler.
+
+use crate::MirPass;
+use rustc_index::IndexVec;
+use rustc_middle::mir::interpret::ConstValue;
+use rustc_middle::mir::visit::MutVisitor;
+use rustc_middle::mir::*;
+use rustc_middle::ty::{self, Ty, TyCtxt};
+use rustc_span::def_id::DefId;
+
+pub struct InlineFutureIntoFuture;
+impl<'tcx> MirPass<'tcx> for InlineFutureIntoFuture {
+    fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
+        sess.mir_opt_level() > 0 // on by default w/o -Zmir-opt-level=0
+    }
+
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        let Some(into_future_fn_def_id) = tcx.lang_items().into_future_fn() else { return; };
+        let Some(future_trait_def_id) = tcx.lang_items().future_trait() else { return; };
+        let mir_source_def_id = body.source.def_id();
+        trace!("Running InlineFutureIntoFuture on {:?}", body.source);
+        let local_decls = body.local_decls().to_owned();
+        let mut v = Inliner {
+            tcx,
+            into_future_fn_def_id,
+            future_trait_def_id,
+            mir_source_def_id,
+            local_decls,
+        };
+        v.visit_body(body);
+    }
+}
+
+struct Inliner<'tcx> {
+    tcx: TyCtxt<'tcx>,
+    mir_source_def_id: DefId,
+    into_future_fn_def_id: DefId,
+    future_trait_def_id: DefId,
+    local_decls: IndexVec<Local, LocalDecl<'tcx>>,
+}
+
+#[derive(Copy, Clone, PartialEq, Eq)]
+enum FoundImplFuture {
+    Yes,
+    No,
+}
+
+#[derive(Copy, Clone, PartialEq, Eq)]
+enum FoundIntoFutureCall {
+    Yes,
+    No,
+}
+
+struct ImplFutureCallingIntoFuture<'tcx> {
+    args: Vec<Operand<'tcx>>,
+    destination: Place<'tcx>,
+    target: Option<BasicBlock>,
+}
+
+impl<'tcx> Inliner<'tcx> {
+    // This verifies that `ty` implements `Future`, according to the where
+    // clauses (i.e. predicates) attached to the source code identified by
+    // `mir_source_def_id`).
+    fn does_ty_impl_future(&self, ty: Ty<'tcx>) -> FoundImplFuture {
+        let mir_source_predicates = self.tcx.predicates_of(self.mir_source_def_id);
+        let predicates = mir_source_predicates.instantiate_identity(self.tcx);
+        for pred in &predicates.predicates {
+            let Some(kind) = pred.kind().no_bound_vars() else { continue; };
+            let ty::ClauseKind::Trait(trait_pred) = kind else { continue; };
+            let ty::TraitPredicate { trait_ref, polarity: ty::ImplPolarity::Positive } = trait_pred else { continue };
+
+            // FIXME: justify ignoring `substs` below. My current argument is
+            // that `trait Future` has no generic parameters, and the blanket
+            // impl of `IntoFuture` for all futures does not put any constrants
+            // on the associated items of those futures. But it is worth running
+            // this by a trait system expert to validate.
+            let ty::TraitRef { def_id: trait_def_id, .. } = trait_ref;
+            let self_ty = trait_ref.self_ty();
+            if trait_def_id == self.future_trait_def_id {
+                if self_ty == ty {
+                    return FoundImplFuture::Yes;
+                }
+            }
+        }
+        FoundImplFuture::No
+    }
+}
+
+impl<'tcx> MutVisitor<'tcx> for Inliner<'tcx> {
+    fn tcx<'a>(&'a self) -> TyCtxt<'tcx> {
+        self.tcx
+    }
+    fn visit_basic_block_data(&mut self, _bb: BasicBlock, bb_data: &mut BasicBlockData<'tcx>) {
+        let Some(term) = &mut bb_data.terminator else { return; };
+        let Some(result) = self.analyze_terminator(term) else { return; };
+        let ImplFutureCallingIntoFuture {
+            args, destination: dest, target: Some(target)
+        } = result else { return; };
+
+        // At this point, we have identified this terminator as a call to the
+        // associated function `<impl Future as IntoFuture>::into_future`
+        // Due to our knowledge of how libcore implements Future and IntoFuture,
+        // we know we can replace such a call with a trivial move.
+
+        let Some(arg0) = args.get(0) else { return; };
+
+        trace!("InlineFutureIntoFuture bb_data args:{args:?} dest:{dest:?} target:{target:?}");
+
+        bb_data.statements.push(Statement {
+            source_info: term.source_info,
+            kind: StatementKind::Assign(Box::new((dest, Rvalue::Use(arg0.clone())))),
+        });
+        term.kind = TerminatorKind::Goto { target }
+    }
+}
+
+impl<'tcx> Inliner<'tcx> {
+    fn analyze_terminator(
+        &mut self,
+        term: &mut Terminator<'tcx>,
+    ) -> Option<ImplFutureCallingIntoFuture<'tcx>> {
+        let mut found = (FoundImplFuture::No, FoundIntoFutureCall::No);
+        let &TerminatorKind::Call {
+            ref func, ref args, destination, target, fn_span: _, unwind: _, call_source: _
+        } = &term.kind else { return None; };
+        let Operand::Constant(c) = func else { return None; };
+        let ConstantKind::Val(val_const, const_ty) = c.literal else { return None; };
+        let ConstValue::ZeroSized = val_const else { return None; };
+        let ty::FnDef(fn_def_id, substs) =  const_ty.kind() else { return None; };
+        if *fn_def_id == self.into_future_fn_def_id {
+            found.1 = FoundIntoFutureCall::Yes;
+        } else {
+            trace!("InlineFutureIntoFuture bail as this is not `into_future` invocation.");
+            return None;
+        }
+        let arg0_ty = args.get(0).map(|arg0| arg0.ty(&self.local_decls, self.tcx()));
+        trace!("InlineFutureIntoFuture substs:{substs:?} args:{args:?} arg0 ty:{arg0_ty:?}");
+        let Some(arg0_ty) = arg0_ty else { return None; };
+        found.0 = self.does_ty_impl_future(arg0_ty);
+        if let (FoundImplFuture::Yes, FoundIntoFutureCall::Yes) = found {
+            trace!("InlineFutureIntoFuture can replace {term:?}, a {func:?} call, with move");
+            if !self.tcx.consider_optimizing(|| {
+                format!("InlineFutureIntoFuture {:?}", self.mir_source_def_id)
+            }) {
+                return None;
+            }
+            let args = args.clone();
+            Some(ImplFutureCallingIntoFuture { args, destination, target })
+        } else {
+            None
+        }
+    }
+}

compiler/rustc_mir_transform/src/lib.rs

@@ -76,6 +76,7 @@ mod ffi_unwind_calls;
 mod function_item_references;
 mod generator;
 pub mod inline;
+mod inline_future_into_future;
 mod instsimplify;
 mod large_enums;
 mod lower_intrinsics;
@@ -492,6 +493,9 @@ fn run_runtime_lowering_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         // `AddRetag` needs to run after `ElaborateDrops`. Otherwise it should run fairly late,
         // but before optimizations begin.
         &elaborate_box_derefs::ElaborateBoxDerefs,
+        // `InlineFutureIntoFuture` needs to run before `UpvarToLocalProp`, because its
+        // purpose is to enhance the effectiveness of the latter transformation.
+        &inline_future_into_future::InlineFutureIntoFuture,
         // `UpvarToLocalProp` needs to run before `generator::StateTransform`, because its
         // purpose is to coalesce locals into their original upvars before fresh space is
         // allocated for them in the generator.

tests/ui/async-await/future-sizes/async-awaiting-fut.stdout

@@ -1,39 +1,38 @@
-print-type-size type: `[async fn body@$DIR/async-awaiting-fut.rs:21:21: 24:2]`: 2053 bytes, alignment: 1 bytes
+print-type-size type: `[async fn body@$DIR/async-awaiting-fut.rs:21:21: 24:2]`: 1029 bytes, alignment: 1 bytes
 print-type-size     discriminant: 1 bytes
 print-type-size     variant `Unresumed`: 0 bytes
-print-type-size     variant `Suspend0`: 2052 bytes
-print-type-size         local `.__awaitee`: 2052 bytes
+print-type-size     variant `Suspend0`: 1028 bytes
+print-type-size         local `.__awaitee`: 1028 bytes
 print-type-size     variant `Returned`: 0 bytes
 print-type-size     variant `Panicked`: 0 bytes
-print-type-size type: `[async fn body@$DIR/async-awaiting-fut.rs:10:64: 19:2]`: 2052 bytes, alignment: 1 bytes
+print-type-size type: `[async fn body@$DIR/async-awaiting-fut.rs:10:64: 19:2]`: 1028 bytes, alignment: 1 bytes
 print-type-size     discriminant: 1 bytes
 print-type-size     variant `Unresumed`: 1025 bytes
 print-type-size         upvar `.fut`: 1025 bytes, offset: 0 bytes, alignment: 1 bytes
 print-type-size     variant `Suspend0`: 1027 bytes
 print-type-size         upvar `.fut`: 1025 bytes, offset: 0 bytes, alignment: 1 bytes
 print-type-size         padding: 1 bytes
-print-type-size         local `..generator_field3`: 1 bytes, alignment: 1 bytes
+print-type-size         local `..generator_field2`: 1 bytes, alignment: 1 bytes
 print-type-size         local `.__awaitee`: 1 bytes
-print-type-size     variant `Suspend1`: 2051 bytes
+print-type-size     variant `Suspend1`: 1026 bytes
 print-type-size         upvar `.fut`: 1025 bytes, offset: 0 bytes, alignment: 1 bytes
 print-type-size         padding: 1 bytes
-print-type-size         local `..generator_field3`: 1 bytes, alignment: 1 bytes
-print-type-size         local `.__awaitee`: 1025 bytes
+print-type-size         local `..generator_field2`: 1 bytes, alignment: 1 bytes
 print-type-size     variant `Suspend2`: 1027 bytes
 print-type-size         upvar `.fut`: 1025 bytes, offset: 0 bytes, alignment: 1 bytes
 print-type-size         padding: 1 bytes
-print-type-size         local `..generator_field3`: 1 bytes, alignment: 1 bytes
+print-type-size         local `..generator_field2`: 1 bytes, alignment: 1 bytes
 print-type-size         local `.__awaitee`: 1 bytes
 print-type-size     variant `Returned`: 1025 bytes
 print-type-size         upvar `.fut`: 1025 bytes, offset: 0 bytes, alignment: 1 bytes
 print-type-size     variant `Panicked`: 1025 bytes
 print-type-size         upvar `.fut`: 1025 bytes, offset: 0 bytes, alignment: 1 bytes
-print-type-size type: `std::mem::ManuallyDrop<[async fn body@$DIR/async-awaiting-fut.rs:10:64: 19:2]>`: 2052 bytes, alignment: 1 bytes
-print-type-size     field `.value`: 2052 bytes
-print-type-size type: `std::mem::MaybeUninit<[async fn body@$DIR/async-awaiting-fut.rs:10:64: 19:2]>`: 2052 bytes, alignment: 1 bytes
-print-type-size     variant `MaybeUninit`: 2052 bytes
+print-type-size type: `std::mem::ManuallyDrop<[async fn body@$DIR/async-awaiting-fut.rs:10:64: 19:2]>`: 1028 bytes, alignment: 1 bytes
+print-type-size     field `.value`: 1028 bytes
+print-type-size type: `std::mem::MaybeUninit<[async fn body@$DIR/async-awaiting-fut.rs:10:64: 19:2]>`: 1028 bytes, alignment: 1 bytes
+print-type-size     variant `MaybeUninit`: 1028 bytes
 print-type-size         field `.uninit`: 0 bytes
-print-type-size         field `.value`: 2052 bytes
+print-type-size         field `.value`: 1028 bytes
 print-type-size type: `[async fn body@$DIR/async-awaiting-fut.rs:8:35: 8:37]`: 1025 bytes, alignment: 1 bytes
 print-type-size     discriminant: 1 bytes
 print-type-size     variant `Unresumed`: 1024 bytes

tests/ui/async-await/future-sizes/future-as-arg.rs

@@ -11,6 +11,6 @@ fn main() {
     let actual = std::mem::size_of_val(
         &use_future(use_future(use_future(use_future(use_future(test([0; 16])))))));
     // Not using an exact number in case it slightly changes over different commits
-    let expected = 550;
+    let expected = 20;
     assert!(actual > expected, "expected: >{expected}, actual: {actual}");
 }