|
| 1 | +//! Converts `y = <Future as IntoFuture>::into_future(x);` into just `y = x;`, |
| 2 | +//! since we "know" that matches the behavior of the blanket implementation of |
| 3 | +//! IntoFuture for F where F: Future. |
| 4 | +//! |
| 5 | +//! FIXME: determine such coalescing is sound. In particular, check whether |
| 6 | +//! specialization could foil our plans here! |
| 7 | +//! |
| 8 | +//! This is meant to enhance the effectiveness of the upvar-to-local-prop |
| 9 | +//! transformation in reducing the size of the generators constructed by the |
| 10 | +//! compiler. |
| 11 | +
|
| 12 | +use crate::MirPass; |
| 13 | +use rustc_index::IndexVec; |
| 14 | +use rustc_middle::mir::interpret::ConstValue; |
| 15 | +use rustc_middle::mir::visit::MutVisitor; |
| 16 | +use rustc_middle::mir::*; |
| 17 | +use rustc_middle::ty::{self, Ty, TyCtxt}; |
| 18 | +use rustc_span::def_id::DefId; |
| 19 | + |
| 20 | +pub struct InlineFutureIntoFuture; |
| 21 | +impl<'tcx> MirPass<'tcx> for InlineFutureIntoFuture { |
| 22 | + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { |
| 23 | + sess.mir_opt_level() > 0 // on by default w/o -Zmir-opt-level=0 |
| 24 | + } |
| 25 | + |
| 26 | + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
| 27 | + let Some(into_future_fn_def_id) = tcx.lang_items().into_future_fn() else { return; }; |
| 28 | + let Some(future_trait_def_id) = tcx.lang_items().future_trait() else { return; }; |
| 29 | + let mir_source_def_id = body.source.def_id(); |
| 30 | + trace!("Running InlineFutureIntoFuture on {:?}", body.source); |
| 31 | + let local_decls = body.local_decls().to_owned(); |
| 32 | + let mut v = Inliner { |
| 33 | + tcx, |
| 34 | + into_future_fn_def_id, |
| 35 | + future_trait_def_id, |
| 36 | + mir_source_def_id, |
| 37 | + local_decls, |
| 38 | + }; |
| 39 | + v.visit_body(body); |
| 40 | + } |
| 41 | +} |
| 42 | + |
| 43 | +struct Inliner<'tcx> { |
| 44 | + tcx: TyCtxt<'tcx>, |
| 45 | + mir_source_def_id: DefId, |
| 46 | + into_future_fn_def_id: DefId, |
| 47 | + future_trait_def_id: DefId, |
| 48 | + local_decls: IndexVec<Local, LocalDecl<'tcx>>, |
| 49 | +} |
| 50 | + |
| 51 | +#[derive(Copy, Clone, PartialEq, Eq)] |
| 52 | +enum FoundImplFuture { |
| 53 | + Yes, |
| 54 | + No, |
| 55 | +} |
| 56 | + |
| 57 | +#[derive(Copy, Clone, PartialEq, Eq)] |
| 58 | +enum FoundIntoFutureCall { |
| 59 | + Yes, |
| 60 | + No, |
| 61 | +} |
| 62 | + |
| 63 | +struct ImplFutureCallingIntoFuture<'tcx> { |
| 64 | + args: Vec<Operand<'tcx>>, |
| 65 | + destination: Place<'tcx>, |
| 66 | + target: Option<BasicBlock>, |
| 67 | +} |
| 68 | + |
| 69 | +impl<'tcx> Inliner<'tcx> { |
| 70 | + // This verifies that `ty` implements `Future`, according to the where |
| 71 | + // clauses (i.e. predicates) attached to the source code identified by |
| 72 | + // `mir_source_def_id`). |
| 73 | + fn does_ty_impl_future(&self, ty: Ty<'tcx>) -> FoundImplFuture { |
| 74 | + let mir_source_predicates = self.tcx.predicates_of(self.mir_source_def_id); |
| 75 | + let predicates = mir_source_predicates.instantiate_identity(self.tcx); |
| 76 | + for pred in &predicates.predicates { |
| 77 | + let Some(kind) = pred.kind().no_bound_vars() else { continue; }; |
| 78 | + let ty::ClauseKind::Trait(trait_pred) = kind else { continue; }; |
| 79 | + let ty::TraitPredicate { trait_ref, polarity: ty::ImplPolarity::Positive } = trait_pred else { continue }; |
| 80 | + |
| 81 | + // FIXME: justify ignoring `substs` below. My current argument is |
| 82 | + // that `trait Future` has no generic parameters, and the blanket |
| 83 | + // impl of `IntoFuture` for all futures does not put any constrants |
| 84 | + // on the associated items of those futures. But it is worth running |
| 85 | + // this by a trait system expert to validate. |
| 86 | + let ty::TraitRef { def_id: trait_def_id, .. } = trait_ref; |
| 87 | + let self_ty = trait_ref.self_ty(); |
| 88 | + if trait_def_id == self.future_trait_def_id { |
| 89 | + if self_ty == ty { |
| 90 | + return FoundImplFuture::Yes; |
| 91 | + } |
| 92 | + } |
| 93 | + } |
| 94 | + FoundImplFuture::No |
| 95 | + } |
| 96 | +} |
| 97 | + |
| 98 | +impl<'tcx> MutVisitor<'tcx> for Inliner<'tcx> { |
| 99 | + fn tcx<'a>(&'a self) -> TyCtxt<'tcx> { |
| 100 | + self.tcx |
| 101 | + } |
| 102 | + fn visit_basic_block_data(&mut self, _bb: BasicBlock, bb_data: &mut BasicBlockData<'tcx>) { |
| 103 | + let Some(term) = &mut bb_data.terminator else { return; }; |
| 104 | + let Some(result) = self.analyze_terminator(term) else { return; }; |
| 105 | + let ImplFutureCallingIntoFuture { |
| 106 | + args, destination: dest, target: Some(target) |
| 107 | + } = result else { return; }; |
| 108 | + |
| 109 | + // At this point, we have identified this terminator as a call to the |
| 110 | + // associated function `<impl Future as IntoFuture>::into_future` |
| 111 | + // Due to our knowledge of how libcore implements Future and IntoFuture, |
| 112 | + // we know we can replace such a call with a trivial move. |
| 113 | + |
| 114 | + let Some(arg0) = args.get(0) else { return; }; |
| 115 | + |
| 116 | + trace!("InlineFutureIntoFuture bb_data args:{args:?} dest:{dest:?} target:{target:?}"); |
| 117 | + |
| 118 | + bb_data.statements.push(Statement { |
| 119 | + source_info: term.source_info, |
| 120 | + kind: StatementKind::Assign(Box::new((dest, Rvalue::Use(arg0.clone())))), |
| 121 | + }); |
| 122 | + term.kind = TerminatorKind::Goto { target } |
| 123 | + } |
| 124 | +} |
| 125 | + |
| 126 | +impl<'tcx> Inliner<'tcx> { |
| 127 | + fn analyze_terminator( |
| 128 | + &mut self, |
| 129 | + term: &mut Terminator<'tcx>, |
| 130 | + ) -> Option<ImplFutureCallingIntoFuture<'tcx>> { |
| 131 | + let mut found = (FoundImplFuture::No, FoundIntoFutureCall::No); |
| 132 | + let &TerminatorKind::Call { |
| 133 | + ref func, ref args, destination, target, fn_span: _, unwind: _, call_source: _ |
| 134 | + } = &term.kind else { return None; }; |
| 135 | + let Operand::Constant(c) = func else { return None; }; |
| 136 | + let ConstantKind::Val(val_const, const_ty) = c.literal else { return None; }; |
| 137 | + let ConstValue::ZeroSized = val_const else { return None; }; |
| 138 | + let ty::FnDef(fn_def_id, substs) = const_ty.kind() else { return None; }; |
| 139 | + if *fn_def_id == self.into_future_fn_def_id { |
| 140 | + found.1 = FoundIntoFutureCall::Yes; |
| 141 | + } else { |
| 142 | + trace!("InlineFutureIntoFuture bail as this is not `into_future` invocation."); |
| 143 | + return None; |
| 144 | + } |
| 145 | + let arg0_ty = args.get(0).map(|arg0| arg0.ty(&self.local_decls, self.tcx())); |
| 146 | + trace!("InlineFutureIntoFuture substs:{substs:?} args:{args:?} arg0 ty:{arg0_ty:?}"); |
| 147 | + let Some(arg0_ty) = arg0_ty else { return None; }; |
| 148 | + found.0 = self.does_ty_impl_future(arg0_ty); |
| 149 | + if let (FoundImplFuture::Yes, FoundIntoFutureCall::Yes) = found { |
| 150 | + trace!("InlineFutureIntoFuture can replace {term:?}, a {func:?} call, with move"); |
| 151 | + if !self.tcx.consider_optimizing(|| { |
| 152 | + format!("InlineFutureIntoFuture {:?}", self.mir_source_def_id) |
| 153 | + }) { |
| 154 | + return None; |
| 155 | + } |
| 156 | + let args = args.clone(); |
| 157 | + Some(ImplFutureCallingIntoFuture { args, destination, target }) |
| 158 | + } else { |
| 159 | + None |
| 160 | + } |
| 161 | + } |
| 162 | +} |
0 commit comments