Auto merge of #13209 - lowr:feat/inference-for-generator, r=Veykril

feat: type inference for generators

This PR implements basic type inference for generator and yield expressions.

Things not included in this PR:
- Generator upvars and generator witnesses are not implemented. They are only used to determine auto trait impls, so basic type inference should be fine without them, but method resolutions with auto trait bounds may not be resolved correctly.

Open questions:
- I haven't (yet) implemented `HirDisplay` for `TyKind::Generator`, so generator types are just shown as "{{generator}}" (in tests, inlay hints, hovers, etc.), which is not really nice. How should we show them?
- I added moderate amount of stuffs to minicore. I especially didn't want to add `impl<T> Deref for &T` and `impl<T> Deref for &mut T` exclusively for tests for generators; should I move them into the test fixtures or can they be placed in minicore?

cc #4309

Author: bors

crates/hir-def/src/body/lower.rs

@@ -29,8 +29,9 @@ use crate::{
     builtin_type::{BuiltinFloat, BuiltinInt, BuiltinUint},
     db::DefDatabase,
     expr::{
-        dummy_expr_id, Array, BindingAnnotation, Expr, ExprId, FloatTypeWrapper, Label, LabelId,
-        Literal, MatchArm, Pat, PatId, RecordFieldPat, RecordLitField, Statement,
+        dummy_expr_id, Array, BindingAnnotation, ClosureKind, Expr, ExprId, FloatTypeWrapper,
+        Label, LabelId, Literal, MatchArm, Movability, Pat, PatId, RecordFieldPat, RecordLitField,
+        Statement,
     },
     intern::Interned,
     item_scope::BuiltinShadowMode,
@@ -97,6 +98,7 @@ pub(super) fn lower(
         name_to_pat_grouping: Default::default(),
         is_lowering_inside_or_pat: false,
         is_lowering_assignee_expr: false,
+        is_lowering_generator: false,
     }
     .collect(params, body)
 }
@@ -111,6 +113,7 @@ struct ExprCollector<'a> {
     name_to_pat_grouping: FxHashMap<Name, Vec<PatId>>,
     is_lowering_inside_or_pat: bool,
     is_lowering_assignee_expr: bool,
+    is_lowering_generator: bool,
 }
 
 impl ExprCollector<'_> {
@@ -358,6 +361,7 @@ impl ExprCollector<'_> {
                 self.alloc_expr(Expr::Return { expr }, syntax_ptr)
             }
             ast::Expr::YieldExpr(e) => {
+                self.is_lowering_generator = true;
                 let expr = e.expr().map(|e| self.collect_expr(e));
                 self.alloc_expr(Expr::Yield { expr }, syntax_ptr)
             }
@@ -459,13 +463,31 @@ impl ExprCollector<'_> {
                     .ret_type()
                     .and_then(|r| r.ty())
                     .map(|it| Interned::new(TypeRef::from_ast(&self.ctx(), it)));
+
+                let prev_is_lowering_generator = self.is_lowering_generator;
+                self.is_lowering_generator = false;
+
                 let body = self.collect_expr_opt(e.body());
+
+                let closure_kind = if self.is_lowering_generator {
+                    let movability = if e.static_token().is_some() {
+                        Movability::Static
+                    } else {
+                        Movability::Movable
+                    };
+                    ClosureKind::Generator(movability)
+                } else {
+                    ClosureKind::Closure
+                };
+                self.is_lowering_generator = prev_is_lowering_generator;
+
                 self.alloc_expr(
                     Expr::Closure {
                         args: args.into(),
                         arg_types: arg_types.into(),
                         ret_type,
                         body,
+                        closure_kind,
                     },
                     syntax_ptr,
                 )

crates/hir-def/src/body/pretty.rs

@@ -5,7 +5,7 @@ use std::fmt::{self, Write};
 use syntax::ast::HasName;
 
 use crate::{
-    expr::{Array, BindingAnnotation, Literal, Statement},
+    expr::{Array, BindingAnnotation, ClosureKind, Literal, Movability, Statement},
     pretty::{print_generic_args, print_path, print_type_ref},
     type_ref::TypeRef,
 };
@@ -362,7 +362,10 @@ impl<'a> Printer<'a> {
                 self.print_expr(*index);
                 w!(self, "]");
             }
-            Expr::Closure { args, arg_types, ret_type, body } => {
+            Expr::Closure { args, arg_types, ret_type, body, closure_kind } => {
+                if let ClosureKind::Generator(Movability::Static) = closure_kind {
+                    w!(self, "static ");
+                }
                 w!(self, "|");
                 for (i, (pat, ty)) in args.iter().zip(arg_types.iter()).enumerate() {
                     if i != 0 {

crates/hir-def/src/expr.rs

@@ -198,6 +198,7 @@ pub enum Expr {
         arg_types: Box<[Option<Interned<TypeRef>>]>,
         ret_type: Option<Interned<TypeRef>>,
         body: ExprId,
+        closure_kind: ClosureKind,
     },
     Tuple {
         exprs: Box<[ExprId]>,
@@ -211,6 +212,18 @@ pub enum Expr {
     Underscore,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ClosureKind {
+    Closure,
+    Generator(Movability),
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum Movability {
+    Static,
+    Movable,
+}
+
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub enum Array {
     ElementList { elements: Box<[ExprId]>, is_assignee_expr: bool },

crates/hir-ty/src/builder.rs

@@ -9,8 +9,8 @@ use chalk_ir::{
     AdtId, BoundVar, DebruijnIndex, Scalar,
 };
 use hir_def::{
-    builtin_type::BuiltinType, generics::TypeOrConstParamData, ConstParamId, GenericDefId, TraitId,
-    TypeAliasId,
+    builtin_type::BuiltinType, generics::TypeOrConstParamData, ConstParamId, DefWithBodyId,
+    GenericDefId, TraitId, TypeAliasId,
 };
 use smallvec::SmallVec;
 
@@ -205,6 +205,38 @@ impl TyBuilder<()> {
         )
     }
 
+    /// Creates a `TyBuilder` to build `Substitution` for a generator defined in `parent`.
+    ///
+    /// A generator's substitution consists of:
+    /// - generic parameters in scope on `parent`
+    /// - resume type of generator
+    /// - yield type of generator ([`Generator::Yield`](std::ops::Generator::Yield))
+    /// - return type of generator ([`Generator::Return`](std::ops::Generator::Return))
+    /// in this order.
+    ///
+    /// This method prepopulates the builder with placeholder substitution of `parent`, so you
+    /// should only push exactly 3 `GenericArg`s before building.
+    pub fn subst_for_generator(db: &dyn HirDatabase, parent: DefWithBodyId) -> TyBuilder<()> {
+        let parent_subst = match parent.as_generic_def_id() {
+            Some(parent) => generics(db.upcast(), parent).placeholder_subst(db),
+            // Static initializers *may* contain generators.
+            None => Substitution::empty(Interner),
+        };
+        let builder = TyBuilder::new(
+            (),
+            parent_subst
+                .iter(Interner)
+                .map(|arg| match arg.constant(Interner) {
+                    Some(c) => ParamKind::Const(c.data(Interner).ty.clone()),
+                    None => ParamKind::Type,
+                })
+                // These represent resume type, yield type, and return type of generator.
+                .chain(std::iter::repeat(ParamKind::Type).take(3))
+                .collect(),
+        );
+        builder.use_parent_substs(&parent_subst)
+    }
+
     pub fn build(self) -> Substitution {
         let ((), subst) = self.build_internal();
         subst

crates/hir-ty/src/chalk_db.rs

@@ -11,6 +11,7 @@ use chalk_solve::rust_ir::{self, OpaqueTyDatumBound, WellKnownTrait};
 
 use base_db::CrateId;
 use hir_def::{
+    expr::Movability,
     lang_item::{lang_attr, LangItemTarget},
     AssocItemId, GenericDefId, HasModule, ItemContainerId, Lookup, ModuleId, TypeAliasId,
 };
@@ -26,9 +27,9 @@ use crate::{
     to_assoc_type_id, to_chalk_trait_id,
     traits::ChalkContext,
     utils::generics,
-    AliasEq, AliasTy, BoundVar, CallableDefId, DebruijnIndex, FnDefId, Interner, ProjectionTy,
-    ProjectionTyExt, QuantifiedWhereClause, Substitution, TraitRef, TraitRefExt, Ty, TyBuilder,
-    TyExt, TyKind, WhereClause,
+    wrap_empty_binders, AliasEq, AliasTy, BoundVar, CallableDefId, DebruijnIndex, FnDefId,
+    Interner, ProjectionTy, ProjectionTyExt, QuantifiedWhereClause, Substitution, TraitRef,
+    TraitRefExt, Ty, TyBuilder, TyExt, TyKind, WhereClause,
 };
 
 pub(crate) type AssociatedTyDatum = chalk_solve::rust_ir::AssociatedTyDatum<Interner>;
@@ -372,17 +373,63 @@ impl<'a> chalk_solve::RustIrDatabase<Interner> for ChalkContext<'a> {
     }
     fn generator_datum(
         &self,
-        _: chalk_ir::GeneratorId<Interner>,
+        id: chalk_ir::GeneratorId<Interner>,
     ) -> std::sync::Arc<chalk_solve::rust_ir::GeneratorDatum<Interner>> {
-        // FIXME
-        unimplemented!()
+        let (parent, expr) = self.db.lookup_intern_generator(id.into());
+
+        // We fill substitution with unknown type, because we only need to know whether the generic
+        // params are types or consts to build `Binders` and those being filled up are for
+        // `resume_type`, `yield_type`, and `return_type` of the generator in question.
+        let subst = TyBuilder::subst_for_generator(self.db, parent).fill_with_unknown().build();
+
+        let len = subst.len(Interner);
+        let input_output = rust_ir::GeneratorInputOutputDatum {
+            resume_type: TyKind::BoundVar(BoundVar::new(DebruijnIndex::INNERMOST, len - 3))
+                .intern(Interner),
+            yield_type: TyKind::BoundVar(BoundVar::new(DebruijnIndex::INNERMOST, len - 2))
+                .intern(Interner),
+            return_type: TyKind::BoundVar(BoundVar::new(DebruijnIndex::INNERMOST, len - 1))
+                .intern(Interner),
+            // FIXME: calculate upvars
+            upvars: vec![],
+        };
+
+        let it = subst
+            .iter(Interner)
+            .map(|it| it.constant(Interner).map(|c| c.data(Interner).ty.clone()));
+        let input_output = crate::make_type_and_const_binders(it, input_output);
+
+        let movability = match self.db.body(parent)[expr] {
+            hir_def::expr::Expr::Closure {
+                closure_kind: hir_def::expr::ClosureKind::Generator(movability),
+                ..
+            } => movability,
+            _ => unreachable!("non generator expression interned as generator"),
+        };
+        let movability = match movability {
+            Movability::Static => rust_ir::Movability::Static,
+            Movability::Movable => rust_ir::Movability::Movable,
+        };
+
+        Arc::new(rust_ir::GeneratorDatum { movability, input_output })
     }
     fn generator_witness_datum(
         &self,
-        _: chalk_ir::GeneratorId<Interner>,
+        id: chalk_ir::GeneratorId<Interner>,
     ) -> std::sync::Arc<chalk_solve::rust_ir::GeneratorWitnessDatum<Interner>> {
-        // FIXME
-        unimplemented!()
+        // FIXME: calculate inner types
+        let inner_types =
+            rust_ir::GeneratorWitnessExistential { types: wrap_empty_binders(vec![]) };
+
+        let (parent, _) = self.db.lookup_intern_generator(id.into());
+        // See the comment in `generator_datum()` for unknown types.
+        let subst = TyBuilder::subst_for_generator(self.db, parent).fill_with_unknown().build();
+        let it = subst
+            .iter(Interner)
+            .map(|it| it.constant(Interner).map(|c| c.data(Interner).ty.clone()));
+        let inner_types = crate::make_type_and_const_binders(it, inner_types);
+
+        Arc::new(rust_ir::GeneratorWitnessDatum { inner_types })
     }
 
     fn unification_database(&self) -> &dyn chalk_ir::UnificationDatabase<Interner> {

crates/hir-ty/src/db.rs

@@ -120,6 +120,8 @@ pub trait HirDatabase: DefDatabase + Upcast<dyn DefDatabase> {
     fn intern_impl_trait_id(&self, id: ImplTraitId) -> InternedOpaqueTyId;
     #[salsa::interned]
     fn intern_closure(&self, id: (DefWithBodyId, ExprId)) -> InternedClosureId;
+    #[salsa::interned]
+    fn intern_generator(&self, id: (DefWithBodyId, ExprId)) -> InternedGeneratorId;
 
     #[salsa::invoke(chalk_db::associated_ty_data_query)]
     fn associated_ty_data(&self, id: chalk_db::AssocTypeId) -> Arc<chalk_db::AssociatedTyDatum>;
@@ -233,6 +235,10 @@ impl_intern_key!(InternedOpaqueTyId);
 pub struct InternedClosureId(salsa::InternId);
 impl_intern_key!(InternedClosureId);
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct InternedGeneratorId(salsa::InternId);
+impl_intern_key!(InternedGeneratorId);
+
 /// This exists just for Chalk, because Chalk just has a single `FnDefId` where
 /// we have different IDs for struct and enum variant constructors.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)]

crates/hir-ty/src/display.rs

@@ -20,6 +20,7 @@ use hir_def::{
 };
 use hir_expand::{hygiene::Hygiene, name::Name};
 use itertools::Itertools;
+use smallvec::SmallVec;
 use syntax::SmolStr;
 
 use crate::{
@@ -221,6 +222,7 @@ pub enum DisplaySourceCodeError {
     PathNotFound,
     UnknownType,
     Closure,
+    Generator,
 }
 
 pub enum HirDisplayError {
@@ -783,7 +785,34 @@ impl HirDisplay for Ty {
                 write!(f, "{{unknown}}")?;
             }
             TyKind::InferenceVar(..) => write!(f, "_")?,
-            TyKind::Generator(..) => write!(f, "{{generator}}")?,
+            TyKind::Generator(_, subst) => {
+                if f.display_target.is_source_code() {
+                    return Err(HirDisplayError::DisplaySourceCodeError(
+                        DisplaySourceCodeError::Generator,
+                    ));
+                }
+
+                let subst = subst.as_slice(Interner);
+                let a: Option<SmallVec<[&Ty; 3]>> = subst
+                    .get(subst.len() - 3..)
+                    .map(|args| args.iter().map(|arg| arg.ty(Interner)).collect())
+                    .flatten();
+
+                if let Some([resume_ty, yield_ty, ret_ty]) = a.as_deref() {
+                    write!(f, "|")?;
+                    resume_ty.hir_fmt(f)?;
+                    write!(f, "|")?;
+
+                    write!(f, " yields ")?;
+                    yield_ty.hir_fmt(f)?;
+
+                    write!(f, " -> ")?;
+                    ret_ty.hir_fmt(f)?;
+                } else {
+                    // This *should* be unreachable, but fallback just in case.
+                    write!(f, "{{generator}}")?;
+                }
+            }
             TyKind::GeneratorWitness(..) => write!(f, "{{generator witness}}")?,
         }
         Ok(())

crates/hir-ty/src/infer.rs

@@ -339,7 +339,7 @@ pub struct InferenceResult {
     /// unresolved or missing subpatterns or subpatterns of mismatched types.
     pub type_of_pat: ArenaMap<PatId, Ty>,
     type_mismatches: FxHashMap<ExprOrPatId, TypeMismatch>,
-    /// Interned Unknown to return references to.
+    /// Interned common types to return references to.
     standard_types: InternedStandardTypes,
     /// Stores the types which were implicitly dereferenced in pattern binding modes.
     pub pat_adjustments: FxHashMap<PatId, Vec<Ty>>,
@@ -419,6 +419,8 @@ pub(crate) struct InferenceContext<'a> {
     /// closures, but currently this is the only field that will change there,
     /// so it doesn't make sense.
     return_ty: Ty,
+    /// The resume type and the yield type, respectively, of the generator being inferred.
+    resume_yield_tys: Option<(Ty, Ty)>,
     diverges: Diverges,
     breakables: Vec<BreakableContext>,
 }
@@ -483,6 +485,7 @@ impl<'a> InferenceContext<'a> {
             table: unify::InferenceTable::new(db, trait_env.clone()),
             trait_env,
             return_ty: TyKind::Error.intern(Interner), // set in collect_fn_signature
+            resume_yield_tys: None,
             db,
             owner,
             body,

crates/hir-ty/src/infer/closure.rs

@@ -12,6 +12,7 @@ use crate::{
 use super::{Expectation, InferenceContext};
 
 impl InferenceContext<'_> {
+    // This function handles both closures and generators.
     pub(super) fn deduce_closure_type_from_expectations(
         &mut self,
         closure_expr: ExprId,
@@ -27,6 +28,11 @@ impl InferenceContext<'_> {
         // Deduction from where-clauses in scope, as well as fn-pointer coercion are handled here.
         let _ = self.coerce(Some(closure_expr), closure_ty, &expected_ty);
 
+        // Generators are not Fn* so return early.
+        if matches!(closure_ty.kind(Interner), TyKind::Generator(..)) {
+            return;
+        }
+
         // Deduction based on the expected `dyn Fn` is done separately.
         if let TyKind::Dyn(dyn_ty) = expected_ty.kind(Interner) {
             if let Some(sig) = self.deduce_sig_from_dyn_ty(dyn_ty) {