address code review

splitting off futures/generators in trait selection / checking

Author: Swatinem

compiler/rustc_middle/src/traits/mod.rs

@@ -660,6 +660,9 @@ pub enum ImplSource<'tcx, N> {
     /// ImplSource automatically generated for a generator.
     Generator(ImplSourceGeneratorData<'tcx, N>),
 
+    /// ImplSource automatically generated for a generator backing an async future.
+    Future(ImplSourceFutureData<'tcx, N>),
+
     /// ImplSource for a trait alias.
     TraitAlias(ImplSourceTraitAliasData<'tcx, N>),
 
@@ -676,6 +679,7 @@ impl<'tcx, N> ImplSource<'tcx, N> {
             ImplSource::AutoImpl(d) => d.nested,
             ImplSource::Closure(c) => c.nested,
             ImplSource::Generator(c) => c.nested,
+            ImplSource::Future(c) => c.nested,
             ImplSource::Object(d) => d.nested,
             ImplSource::FnPointer(d) => d.nested,
             ImplSource::DiscriminantKind(ImplSourceDiscriminantKindData)
@@ -694,6 +698,7 @@ impl<'tcx, N> ImplSource<'tcx, N> {
             ImplSource::AutoImpl(d) => &d.nested,
             ImplSource::Closure(c) => &c.nested,
             ImplSource::Generator(c) => &c.nested,
+            ImplSource::Future(c) => &c.nested,
             ImplSource::Object(d) => &d.nested,
             ImplSource::FnPointer(d) => &d.nested,
             ImplSource::DiscriminantKind(ImplSourceDiscriminantKindData)
@@ -736,7 +741,11 @@ impl<'tcx, N> ImplSource<'tcx, N> {
                 generator_def_id: c.generator_def_id,
                 substs: c.substs,
                 nested: c.nested.into_iter().map(f).collect(),
-                is_async: c.is_async,
+            }),
+            ImplSource::Future(c) => ImplSource::Future(ImplSourceFutureData {
+                generator_def_id: c.generator_def_id,
+                substs: c.substs,
+                nested: c.nested.into_iter().map(f).collect(),
             }),
             ImplSource::FnPointer(p) => ImplSource::FnPointer(ImplSourceFnPointerData {
                 fn_ty: p.fn_ty,
@@ -795,7 +804,16 @@ pub struct ImplSourceGeneratorData<'tcx, N> {
     /// Nested obligations. This can be non-empty if the generator
     /// signature contains associated types.
     pub nested: Vec<N>,
-    pub is_async: bool,
+}
+
+#[derive(Clone, PartialEq, Eq, TyEncodable, TyDecodable, HashStable, Lift)]
+#[derive(TypeFoldable, TypeVisitable)]
+pub struct ImplSourceFutureData<'tcx, N> {
+    pub generator_def_id: DefId,
+    pub substs: SubstsRef<'tcx>,
+    /// Nested obligations. This can be non-empty if the generator
+    /// signature contains associated types.
+    pub nested: Vec<N>,
 }
 
 #[derive(Clone, PartialEq, Eq, TyEncodable, TyDecodable, HashStable, Lift)]

compiler/rustc_middle/src/traits/select.rs

@@ -127,11 +127,13 @@ pub enum SelectionCandidate<'tcx> {
     /// generated for an `||` expression.
     ClosureCandidate,
 
-    /// Implementation of a `Generator` / `Future` trait by one of the anonymous types
+    /// Implementation of a `Generator` trait by one of the anonymous types
     /// generated for a generator.
-    GeneratorCandidate {
-        is_async: bool,
-    },
+    GeneratorCandidate,
+
+    /// Implementation of a `Future` trait by one of the generator types
+    /// generated for an async construct.
+    FutureCandidate,
 
     /// Implementation of a `Fn`-family trait by one of the anonymous
     /// types generated for a fn pointer type (e.g., `fn(int) -> int`)

compiler/rustc_middle/src/traits/structural_impls.rs

@@ -15,6 +15,8 @@ impl<'tcx, N: fmt::Debug> fmt::Debug for traits::ImplSource<'tcx, N> {
 
             super::ImplSource::Generator(ref d) => write!(f, "{:?}", d),
 
+            super::ImplSource::Future(ref d) => write!(f, "{:?}", d),
+
             super::ImplSource::FnPointer(ref d) => write!(f, "({:?})", d),
 
             super::ImplSource::DiscriminantKind(ref d) => write!(f, "{:?}", d),
@@ -58,6 +60,16 @@ impl<'tcx, N: fmt::Debug> fmt::Debug for traits::ImplSourceGeneratorData<'tcx, N
     }
 }
 
+impl<'tcx, N: fmt::Debug> fmt::Debug for traits::ImplSourceFutureData<'tcx, N> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "ImplSourceFutureData(generator_def_id={:?}, substs={:?}, nested={:?})",
+            self.generator_def_id, self.substs, self.nested
+        )
+    }
+}
+
 impl<'tcx, N: fmt::Debug> fmt::Debug for traits::ImplSourceClosureData<'tcx, N> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(

compiler/rustc_mir_transform/src/generator.rs

@@ -11,10 +11,10 @@
 //! generator in the MIR, since it is used to create the drop glue for the generator. We'd get
 //! infinite recursion otherwise.
 //!
-//! This pass creates the implementation for the Generator::resume function and the drop shim
-//! for the generator based on the MIR input. It converts the generator argument from Self to
-//! &mut Self adding derefs in the MIR as needed. It computes the final layout of the generator
-//! struct which looks like this:
+//! This pass creates the implementation for either the `Generator::resume` or `Future::poll`
+//! function and the drop shim for the generator based on the MIR input.
+//! It converts the generator argument from Self to &mut Self adding derefs in the MIR as needed.
+//! It computes the final layout of the generator struct which looks like this:
 //!     First upvars are stored
 //!     It is followed by the generator state field.
 //!     Then finally the MIR locals which are live across a suspension point are stored.
@@ -32,14 +32,15 @@
 //!     2 - Generator has been poisoned
 //!
 //! It also rewrites `return x` and `yield y` as setting a new generator state and returning
-//! GeneratorState::Complete(x) and GeneratorState::Yielded(y) respectively.
+//! `GeneratorState::Complete(x)` and `GeneratorState::Yielded(y)`,
+//! or `Poll::Ready(x)` and `Poll::Pending` respectively.
 //! MIR locals which are live across a suspension point are moved to the generator struct
 //! with references to them being updated with references to the generator struct.
 //!
 //! The pass creates two functions which have a switch on the generator state giving
 //! the action to take.
 //!
-//! One of them is the implementation of Generator::resume.
+//! One of them is the implementation of `Generator::resume` / `Future::poll`.
 //! For generators with state 0 (unresumed) it starts the execution of the generator.
 //! For generators with state 1 (returned) and state 2 (poisoned) it panics.
 //! Otherwise it continues the execution from the last suspension point.
@@ -53,10 +54,10 @@ use crate::deref_separator::deref_finder;
 use crate::simplify;
 use crate::util::expand_aggregate;
 use crate::MirPass;
-use hir::GeneratorKind;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_hir as hir;
 use rustc_hir::lang_items::LangItem;
+use rustc_hir::GeneratorKind;
 use rustc_index::bit_set::{BitMatrix, BitSet, GrowableBitSet};
 use rustc_index::vec::{Idx, IndexVec};
 use rustc_middle::mir::dump_mir;
@@ -246,14 +247,35 @@ impl<'tcx> TransformVisitor<'tcx> {
     // `core::ops::GeneratorState` only has single element tuple variants,
     // so we can just write to the downcasted first field and then set the
     // discriminant to the appropriate variant.
-    fn make_state(&self, idx: VariantIdx) -> (AggregateKind<'tcx>, Option<Ty<'tcx>>) {
+    fn make_state(
+        &self,
+        val: Operand<'tcx>,
+        source_info: SourceInfo,
+        is_return: bool,
+        statements: &mut Vec<Statement<'tcx>>,
+    ) {
+        let idx = VariantIdx::new(match (is_return, self.is_async_kind) {
+            (true, false) => 1,  // GeneratorState::Complete
+            (false, false) => 0, // GeneratorState::Yielded
+            (true, true) => 0,   // Poll::Ready
+            (false, true) => 1,  // Poll::Pending
+        });
+
         let kind = AggregateKind::Adt(self.state_adt_ref.did(), idx, self.state_substs, None, None);
 
         // `Poll::Pending`
         if self.is_async_kind && idx == VariantIdx::new(1) {
             assert_eq!(self.state_adt_ref.variant(idx).fields.len(), 0);
 
-            return (kind, None);
+            // FIXME(swatinem): assert that `val` is indeed unit?
+            statements.extend(expand_aggregate(
+                Place::return_place(),
+                std::iter::empty(),
+                kind,
+                source_info,
+                self.tcx,
+            ));
+            return;
         }
 
         // else: `Poll::Ready(x)`, `GeneratorState::Yielded(x)` or `GeneratorState::Complete(x)`
@@ -264,7 +286,13 @@ impl<'tcx> TransformVisitor<'tcx> {
             .bound_type_of(self.state_adt_ref.variant(idx).fields[0].did)
             .subst(self.tcx, self.state_substs);
 
-        (kind, Some(ty))
+        statements.extend(expand_aggregate(
+            Place::return_place(),
+            std::iter::once((val, ty)),
+            kind,
+            source_info,
+            self.tcx,
+        ));
     }
 
     // Create a Place referencing a generator struct field
@@ -347,32 +375,7 @@ impl<'tcx> MutVisitor<'tcx> for TransformVisitor<'tcx> {
         if let Some((is_return, resume, v, drop)) = ret_val {
             let source_info = data.terminator().source_info;
             // We must assign the value first in case it gets declared dead below
-            let state_idx = VariantIdx::new(match (is_return, self.is_async_kind) {
-                (true, false) => 1,  // GeneratorState::Complete
-                (false, false) => 0, // GeneratorState::Yielded
-                (true, true) => 0,   // Poll::Ready
-                (false, true) => 1,  // Poll::Pending
-            });
-            let (kind, ty) = self.make_state(state_idx);
-            if let Some(ty) = ty {
-                data.statements.extend(expand_aggregate(
-                    Place::return_place(),
-                    std::iter::once((v, ty)),
-                    kind,
-                    source_info,
-                    self.tcx,
-                ));
-            } else {
-                // FIXME(swatinem): assert that `val` is indeed unit?
-                data.statements.extend(expand_aggregate(
-                    Place::return_place(),
-                    std::iter::empty(),
-                    kind,
-                    source_info,
-                    self.tcx,
-                ));
-            }
-
+            self.make_state(v, source_info, is_return, &mut data.statements);
             let state = if let Some((resume, mut resume_arg)) = resume {
                 // Yield
                 let state = RESERVED_VARIANTS + self.suspension_points.len();

compiler/rustc_trait_selection/src/traits/project.rs

@@ -12,7 +12,8 @@ use super::SelectionContext;
 use super::SelectionError;
 use super::{
     ImplSourceClosureData, ImplSourceDiscriminantKindData, ImplSourceFnPointerData,
-    ImplSourceGeneratorData, ImplSourcePointeeData, ImplSourceUserDefinedData,
+    ImplSourceFutureData, ImplSourceGeneratorData, ImplSourcePointeeData,
+    ImplSourceUserDefinedData,
 };
 use super::{Normalized, NormalizedTy, ProjectionCacheEntry, ProjectionCacheKey};
 
@@ -1556,6 +1557,7 @@ fn assemble_candidates_from_impls<'cx, 'tcx>(
         let eligible = match &impl_source {
             super::ImplSource::Closure(_)
             | super::ImplSource::Generator(_)
+            | super::ImplSource::Future(_)
             | super::ImplSource::FnPointer(_)
             | super::ImplSource::TraitAlias(_) => true,
             super::ImplSource::UserDefined(impl_data) => {
@@ -1844,6 +1846,7 @@ fn confirm_select_candidate<'cx, 'tcx>(
     match impl_source {
         super::ImplSource::UserDefined(data) => confirm_impl_candidate(selcx, obligation, data),
         super::ImplSource::Generator(data) => confirm_generator_candidate(selcx, obligation, data),
+        super::ImplSource::Future(data) => confirm_future_candidate(selcx, obligation, data),
         super::ImplSource::Closure(data) => confirm_closure_candidate(selcx, obligation, data),
         super::ImplSource::FnPointer(data) => confirm_fn_pointer_candidate(selcx, obligation, data),
         super::ImplSource::DiscriminantKind(data) => {
@@ -1884,55 +1887,75 @@ fn confirm_generator_candidate<'cx, 'tcx>(
     debug!(?obligation, ?gen_sig, ?obligations, "confirm_generator_candidate");
 
     let tcx = selcx.tcx();
+    let gen_def_id = tcx.require_lang_item(LangItem::Generator, None);
 
-    let predicate = if impl_source.is_async {
-        let fut_def_id = tcx.require_lang_item(LangItem::Future, None);
-
-        // FIXME(swatinem): this is just copy-pasted from `generator_trait_ref_and_outputs` for now
-        let self_ty = obligation.predicate.self_ty();
-        debug_assert!(!self_ty.has_escaping_bound_vars());
-        let trait_ref =
-            ty::TraitRef { def_id: fut_def_id, substs: tcx.mk_substs_trait(self_ty, &[]) };
-        gen_sig.map_bound(|sig| (trait_ref, sig.return_ty)).map_bound(|(trait_ref, return_ty)| {
-            let name = tcx.associated_item(obligation.predicate.item_def_id).name;
-            let ty = if name == sym::Output { return_ty } else { bug!() };
-
-            ty::ProjectionPredicate {
-                projection_ty: ty::ProjectionTy {
-                    substs: trait_ref.substs,
-                    item_def_id: obligation.predicate.item_def_id,
-                },
-                term: ty.into(),
-            }
-        })
-    } else {
-        let gen_def_id = tcx.require_lang_item(LangItem::Generator, None);
+    let predicate = super::util::generator_trait_ref_and_outputs(
+        tcx,
+        gen_def_id,
+        obligation.predicate.self_ty(),
+        gen_sig,
+    )
+    .map_bound(|(trait_ref, yield_ty, return_ty)| {
+        let name = tcx.associated_item(obligation.predicate.item_def_id).name;
+        let ty = if name == sym::Return {
+            return_ty
+        } else if name == sym::Yield {
+            yield_ty
+        } else {
+            bug!()
+        };
 
-        super::util::generator_trait_ref_and_outputs(
-            tcx,
-            gen_def_id,
-            obligation.predicate.self_ty(),
-            gen_sig,
-        )
-        .map_bound(|(trait_ref, yield_ty, return_ty)| {
-            let name = tcx.associated_item(obligation.predicate.item_def_id).name;
-            let ty = if name == sym::Return {
-                return_ty
-            } else if name == sym::Yield {
-                yield_ty
-            } else {
-                bug!()
-            };
+        ty::ProjectionPredicate {
+            projection_ty: ty::ProjectionTy {
+                substs: trait_ref.substs,
+                item_def_id: obligation.predicate.item_def_id,
+            },
+            term: ty.into(),
+        }
+    });
 
-            ty::ProjectionPredicate {
-                projection_ty: ty::ProjectionTy {
-                    substs: trait_ref.substs,
-                    item_def_id: obligation.predicate.item_def_id,
-                },
-                term: ty.into(),
-            }
-        })
-    };
+    confirm_param_env_candidate(selcx, obligation, predicate, false)
+        .with_addl_obligations(impl_source.nested)
+        .with_addl_obligations(obligations)
+}
+
+fn confirm_future_candidate<'cx, 'tcx>(
+    selcx: &mut SelectionContext<'cx, 'tcx>,
+    obligation: &ProjectionTyObligation<'tcx>,
+    impl_source: ImplSourceFutureData<'tcx, PredicateObligation<'tcx>>,
+) -> Progress<'tcx> {
+    let gen_sig = impl_source.substs.as_generator().poly_sig();
+    let Normalized { value: gen_sig, obligations } = normalize_with_depth(
+        selcx,
+        obligation.param_env,
+        obligation.cause.clone(),
+        obligation.recursion_depth + 1,
+        gen_sig,
+    );
+
+    debug!(?obligation, ?gen_sig, ?obligations, "confirm_future_candidate");
+
+    let tcx = selcx.tcx();
+    let fut_def_id = tcx.require_lang_item(LangItem::Future, None);
+
+    let predicate = super::util::future_trait_ref_and_outputs(
+        tcx,
+        fut_def_id,
+        obligation.predicate.self_ty(),
+        gen_sig,
+    )
+    .map_bound(|(trait_ref, return_ty)| {
+        let name = tcx.associated_item(obligation.predicate.item_def_id).name;
+        let ty = if name == sym::Output { return_ty } else { bug!() };
+
+        ty::ProjectionPredicate {
+            projection_ty: ty::ProjectionTy {
+                substs: trait_ref.substs,
+                item_def_id: obligation.predicate.item_def_id,
+            },
+            term: ty.into(),
+        }
+    });
 
     confirm_param_env_candidate(selcx, obligation, predicate, false)
         .with_addl_obligations(impl_source.nested)