Implement generators

`GeneratorDatum` and `GeneratorWitnessDatum` are introduced, which
mirror `TyKind::Generator` and `TyKind::GeneratorWitness` in rustc.

We handle auto traits for `GeneratorDatum` using the existing
`constituent_types` method. For `GeneratorWitnessDatum`, we generate a
`forall<'a, 'b, ... 'z>` goal to reflect the fact that our lifetimes are
existentially bound.

Unresolved questions:
* The implementation of `GeneratorWitnessDatum` needs careful review -
  I'm not certain that I fully understand how `Binders` works.
* The syntax implemented in `parser.lalrpop` can be bikesheeded. We
  might want to use something other than `for<>` for the witness types,
  to reflect that fact that we have existential rather than universal
  quantification.
* The generator grammar only allows quantifying over liftimes for the
  witness types - however, this gets lowered to a `Binders`, which
  supports types and constants as well. Should we do anything else
  to ensure we never end up with types/consts where they're not
  expected?
* I added a test to ensure that we treat the witness lifetimes as
  existentially quantified - for example, if we have
  `for<'a, 'b> Foo<'a, 'b>`, we should only be able to use an auto
  trait impl that is fully generic over lifetimes - e.g.
  `impl<'a> Send for Foo<'a, 'b>` will not be used. I *believe*
  that the test is showing this behavior - it ends up returning
  region constraints that require `'a` and `'b` to outlive each other.
  Since these are 'existential' lifetimes (which cannot be substituted),
  this is an unsatisfiable constraint. However, I'm not sure where we
  should be detecting that this is unsatisfiable.

Author: Aaron1011

chalk-integration/src/db.rs

@@ -8,12 +8,13 @@ use crate::{
 };
 use chalk_ir::{
     AdtId, ApplicationTy, AssocTypeId, Binders, Canonical, CanonicalVarKinds, ClosureId,
-    ConstrainedSubst, Environment, FnDefId, GenericArg, Goal, ImplId, InEnvironment, OpaqueTyId,
-    ProgramClause, ProgramClauses, Substitution, TraitId, Ty, UCanonical,
+    ConstrainedSubst, Environment, FnDefId, GeneratorId, GenericArg, Goal, ImplId, InEnvironment,
+    OpaqueTyId, ProgramClause, ProgramClauses, Substitution, TraitId, Ty, UCanonical,
 };
 use chalk_solve::rust_ir::{
     AdtDatum, AdtRepr, AssociatedTyDatum, AssociatedTyValue, AssociatedTyValueId, ClosureKind,
-    FnDefDatum, FnDefInputsAndOutputDatum, ImplDatum, OpaqueTyDatum, TraitDatum, WellKnownTrait,
+    FnDefDatum, FnDefInputsAndOutputDatum, GeneratorDatum, GeneratorWitnessDatum, ImplDatum,
+    OpaqueTyDatum, TraitDatum, WellKnownTrait,
 };
 use chalk_solve::{RustIrDatabase, Solution, SubstitutionResult};
 use salsa::Database;
@@ -106,6 +107,17 @@ impl RustIrDatabase<ChalkIr> for ChalkDatabase {
         self.program_ir().unwrap().adt_datum(id)
     }
 
+    fn generator_datum(&self, id: GeneratorId<ChalkIr>) -> Arc<GeneratorDatum<ChalkIr>> {
+        self.program_ir().unwrap().generator_datum(id)
+    }
+
+    fn generator_witness_datum(
+        &self,
+        id: GeneratorId<ChalkIr>,
+    ) -> Arc<GeneratorWitnessDatum<ChalkIr>> {
+        self.program_ir().unwrap().generator_witness_datum(id)
+    }
+
     fn adt_repr(&self, id: AdtId<ChalkIr>) -> AdtRepr {
         self.program_ir().unwrap().adt_repr(id)
     }

chalk-integration/src/lib.rs

@@ -27,6 +27,7 @@ pub enum TypeSort {
     Closure,
     Trait,
     Opaque,
+    Generator,
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]

chalk-integration/src/lowering.rs

@@ -2,6 +2,7 @@ mod env;
 mod program_lowerer;
 
 use chalk_ir::cast::{Cast, Caster};
+
 use chalk_ir::{
     self, BoundVar, ClausePriority, DebruijnIndex, ImplId, QuantifiedWhereClauses, Substitution,
     TyKind,
@@ -78,6 +79,21 @@ trait LowerParameterMap {
         // as an object. Actually the handling of object types is
         // probably just kind of messed up right now. That's ok.
     }
+
+    fn parameter_map(&self) -> ParameterMap {
+        self.all_parameters()
+            .into_iter()
+            .zip((0..).map(|i| BoundVar::new(DebruijnIndex::INNERMOST, i)))
+            .map(|(k, v)| {
+                let (kind, name) = k.into();
+                (name, chalk_ir::WithKind::new(kind, v))
+            })
+            .collect()
+    }
+
+    fn interner(&self) -> &ChalkIr {
+        &ChalkIr
+    }
 }
 
 macro_rules! lower_param_map {
@@ -709,6 +725,9 @@ impl LowerWithEnv for Ty {
                     TypeLookup::Opaque(id) => {
                         (chalk_ir::TypeName::OpaqueType(id), env.opaque_kind(id))
                     }
+                    TypeLookup::Generator(id) => {
+                        (chalk_ir::TypeName::Generator(id), env.generator_kind(id))
+                    }
 
                     TypeLookup::Foreign(_) | TypeLookup::Trait(_) => {
                         panic!("Unexpected apply type")
@@ -1070,6 +1089,8 @@ pub fn lower_goal(goal: &Goal, program: &LoweredProgram) -> LowerResult<chalk_ir
         closure_ids: &program.closure_ids,
         trait_ids: &program.trait_ids,
         opaque_ty_ids: &program.opaque_ty_ids,
+        generator_ids: &program.generator_ids,
+        generator_kinds: &program.generator_kinds,
         adt_kinds: &program.adt_kinds,
         fn_def_kinds: &program.fn_def_kinds,
         closure_kinds: &program.closure_kinds,

chalk-integration/src/lowering/env.rs

@@ -1,6 +1,7 @@
 use chalk_ir::interner::HasInterner;
 use chalk_ir::{
-    self, AdtId, BoundVar, ClosureId, DebruijnIndex, FnDefId, OpaqueTyId, TraitId, VariableKinds,
+    self, AdtId, BoundVar, ClosureId, DebruijnIndex, FnDefId, GeneratorId, OpaqueTyId, TraitId,
+    VariableKinds,
 };
 use chalk_ir::{cast::Cast, ForeignDefId, WithKind};
 use chalk_parse::ast::*;
@@ -15,13 +16,15 @@ pub type AdtIds = BTreeMap<Ident, chalk_ir::AdtId<ChalkIr>>;
 pub type FnDefIds = BTreeMap<Ident, chalk_ir::FnDefId<ChalkIr>>;
 pub type ClosureIds = BTreeMap<Ident, chalk_ir::ClosureId<ChalkIr>>;
 pub type TraitIds = BTreeMap<Ident, chalk_ir::TraitId<ChalkIr>>;
+pub type GeneratorIds = BTreeMap<Ident, chalk_ir::GeneratorId<ChalkIr>>;
 pub type OpaqueTyIds = BTreeMap<Ident, chalk_ir::OpaqueTyId<ChalkIr>>;
 pub type AdtKinds = BTreeMap<chalk_ir::AdtId<ChalkIr>, TypeKind>;
 pub type FnDefKinds = BTreeMap<chalk_ir::FnDefId<ChalkIr>, TypeKind>;
 pub type ClosureKinds = BTreeMap<chalk_ir::ClosureId<ChalkIr>, TypeKind>;
 pub type TraitKinds = BTreeMap<chalk_ir::TraitId<ChalkIr>, TypeKind>;
 pub type AutoTraits = BTreeMap<chalk_ir::TraitId<ChalkIr>, bool>;
 pub type OpaqueTyKinds = BTreeMap<chalk_ir::OpaqueTyId<ChalkIr>, TypeKind>;
+pub type GeneratorKinds = BTreeMap<chalk_ir::GeneratorId<ChalkIr>, TypeKind>;
 pub type AssociatedTyLookups = BTreeMap<(chalk_ir::TraitId<ChalkIr>, Ident), AssociatedTyLookup>;
 pub type AssociatedTyValueIds =
     BTreeMap<(chalk_ir::ImplId<ChalkIr>, Ident), AssociatedTyValueId<ChalkIr>>;
@@ -46,6 +49,8 @@ pub struct Env<'k> {
     pub associated_ty_lookups: &'k AssociatedTyLookups,
     pub auto_traits: &'k AutoTraits,
     pub foreign_ty_ids: &'k ForeignIds,
+    pub generator_ids: &'k GeneratorIds,
+    pub generator_kinds: &'k GeneratorKinds,
     /// GenericArg identifiers are used as keys, therefore
     /// all identifiers in an environment must be unique (no shadowing).
     pub parameter_map: ParameterMap,
@@ -78,6 +83,7 @@ pub enum TypeLookup<'k> {
     Opaque(OpaqueTyId<ChalkIr>),
     Foreign(ForeignDefId<ChalkIr>),
     Trait(TraitId<ChalkIr>),
+    Generator(GeneratorId<ChalkIr>),
 }
 
 impl Env<'_> {
@@ -128,6 +134,9 @@ impl Env<'_> {
             Ok(TypeLookup::Closure(id)) => {
                 apply(self.closure_kind(id), chalk_ir::TypeName::Closure(id))
             }
+            Ok(TypeLookup::Generator(id)) => {
+                apply(self.generator_kind(id), chalk_ir::TypeName::Generator(id))
+            }
             Ok(TypeLookup::Opaque(id)) => Ok(chalk_ir::TyData::Alias(chalk_ir::AliasTy::Opaque(
                 chalk_ir::OpaqueTy {
                     opaque_ty_id: id,
@@ -162,6 +171,8 @@ impl Env<'_> {
             Ok(TypeLookup::Foreign(*id))
         } else if let Some(id) = self.trait_ids.get(&name.str) {
             Ok(TypeLookup::Trait(*id))
+        } else if let Some(id) = self.generator_ids.get(&name.str) {
+            Ok(TypeLookup::Generator(*id))
         } else {
             Err(RustIrError::NotStruct(name.clone()))
         }
@@ -203,6 +214,10 @@ impl Env<'_> {
         &self.opaque_ty_kinds[&id]
     }
 
+    pub fn generator_kind(&self, id: chalk_ir::GeneratorId<ChalkIr>) -> &TypeKind {
+        &self.generator_kinds[&id]
+    }
+
     pub fn lookup_associated_ty(
         &self,
         trait_id: TraitId<ChalkIr>,

chalk-integration/src/lowering/program_lowerer.rs

@@ -1,11 +1,12 @@
 use chalk_ir::cast::Cast;
 use chalk_ir::{
-    self, AdtId, AssocTypeId, BoundVar, ClosureId, DebruijnIndex, FnDefId, ForeignDefId, ImplId,
-    OpaqueTyId, TraitId, TyKind, VariableKinds,
+    self, AdtId, AssocTypeId, BoundVar, ClosureId, DebruijnIndex, FnDefId, ForeignDefId,
+    GeneratorId, ImplId, OpaqueTyId, TraitId, TyKind, VariableKinds,
 };
 use chalk_parse::ast::*;
 use chalk_solve::rust_ir::{
-    self, Anonymize, AssociatedTyValueId, OpaqueTyDatum, OpaqueTyDatumBound,
+    self, Anonymize, AssociatedTyValueId, GeneratorDatum, GeneratorInputOutputDatum,
+    GeneratorWitnessDatum, GeneratorWitnessExistential, OpaqueTyDatum, OpaqueTyDatumBound,
 };
 use rust_ir::IntoWhereClauses;
 use std::collections::{BTreeMap, HashSet};
@@ -32,6 +33,8 @@ pub(super) struct ProgramLowerer {
     opaque_ty_ids: OpaqueTyIds,
     adt_kinds: AdtKinds,
     fn_def_kinds: FnDefKinds,
+    generator_ids: GeneratorIds,
+    generator_kinds: GeneratorKinds,
     closure_kinds: ClosureKinds,
     trait_kinds: TraitKinds,
     opaque_ty_kinds: OpaqueTyKinds,
@@ -125,6 +128,11 @@ impl ProgramLowerer {
                     self.foreign_ty_ids
                         .insert(ident.str.clone(), ForeignDefId(raw_id));
                 }
+                Item::GeneratorDefn(defn) => {
+                    let id = GeneratorId(raw_id);
+                    self.generator_ids.insert(defn.name.str.clone(), id);
+                    self.generator_kinds.insert(id, defn.lower_type_kind()?);
+                }
                 Item::Impl(_) => continue,
                 Item::Clause(_) => continue,
             };
@@ -145,6 +153,8 @@ impl ProgramLowerer {
         let mut associated_ty_data = BTreeMap::new();
         let mut associated_ty_values = BTreeMap::new();
         let mut opaque_ty_data = BTreeMap::new();
+        let mut generator_data = BTreeMap::new();
+        let mut generator_witness_data = BTreeMap::new();
         let mut hidden_opaque_types = BTreeMap::new();
         let mut custom_clauses = Vec::new();
 
@@ -160,6 +170,8 @@ impl ProgramLowerer {
                 trait_kinds: &self.trait_kinds,
                 opaque_ty_ids: &self.opaque_ty_ids,
                 opaque_ty_kinds: &self.opaque_ty_kinds,
+                generator_ids: &self.generator_ids,
+                generator_kinds: &self.generator_kinds,
                 associated_ty_lookups: &self.associated_ty_lookups,
                 parameter_map: BTreeMap::new(),
                 auto_traits: &self.auto_traits,
@@ -356,6 +368,51 @@ impl ProgramLowerer {
                         );
                     }
                 }
+                Item::GeneratorDefn(ref defn) => {
+                    let variable_kinds = defn
+                        .variable_kinds
+                        .iter()
+                        .map(|k| k.lower())
+                        .collect::<Vec<_>>();
+
+                    let witness_lifetimes = defn
+                        .witness_lifetimes
+                        .iter()
+                        .map(|i| VariableKind::Lifetime(i.clone()).lower())
+                        .collect::<Vec<_>>();
+
+                    let input_output = empty_env.in_binders(variable_kinds.clone(), |env| {
+                        let yield_type = defn.yield_ty.lower(&env)?;
+                        let resume_type = defn.resume_ty.lower(&env)?;
+                        let return_type = defn.return_ty.lower(&env)?;
+                        let upvars: Result<Vec<_>, _> =
+                            defn.upvars.iter().map(|ty| ty.lower(&env)).collect();
+
+                        Ok(GeneratorInputOutputDatum {
+                            resume_type,
+                            yield_type,
+                            return_type,
+                            upvars: upvars?,
+                        })
+                    })?;
+
+                    let inner_types = empty_env.in_binders(variable_kinds, |env| {
+                        let witnesses = env.in_binders(witness_lifetimes, |env| {
+                            let witnesses: Result<Vec<_>, _> =
+                                defn.witness_types.iter().map(|ty| ty.lower(&env)).collect();
+                            witnesses
+                        })?;
+
+                        Ok(GeneratorWitnessExistential { types: witnesses })
+                    })?;
+
+                    let generator_datum = GeneratorDatum { input_output };
+                    let generator_witness = GeneratorWitnessDatum { inner_types };
+
+                    let id = self.generator_ids[&defn.name.str];
+                    generator_data.insert(id, Arc::new(generator_datum));
+                    generator_witness_data.insert(id, Arc::new(generator_witness));
+                }
                 Item::Foreign(_) => {}
             }
         }
@@ -375,6 +432,10 @@ impl ProgramLowerer {
             fn_def_data,
             closure_inputs_and_output,
             closure_closure_kind,
+            generator_ids: self.generator_ids,
+            generator_kinds: self.generator_kinds,
+            generator_data,
+            generator_witness_data,
             trait_data,
             well_known_traits,
             impl_data,
@@ -416,6 +477,11 @@ lower_type_kind!(AdtDefn, Adt, |defn: &AdtDefn| defn.all_parameters());
 lower_type_kind!(FnDefn, FnDef, |defn: &FnDefn| defn.all_parameters());
 lower_type_kind!(ClosureDefn, Closure, |defn: &ClosureDefn| defn
     .all_parameters());
+lower_type_kind!(GeneratorDefn, Generator, |defn: &GeneratorDefn| defn
+    .variable_kinds
+    .iter()
+    .map(|k| k.lower())
+    .collect::<Vec<_>>());
 lower_type_kind!(TraitDefn, Trait, |defn: &TraitDefn| defn
     .variable_kinds
     .iter()

chalk-integration/src/program.rs

@@ -4,14 +4,14 @@ use chalk_ir::could_match::CouldMatch;
 use chalk_ir::debug::Angle;
 use chalk_ir::{
     debug::SeparatorTraitRef, AdtId, AliasTy, ApplicationTy, AssocTypeId, Binders,
-    CanonicalVarKinds, ClosureId, FnDefId, ForeignDefId, GenericArg, Goal, Goals, ImplId, Lifetime,
-    OpaqueTy, OpaqueTyId, ProgramClause, ProgramClauseImplication, ProgramClauses, ProjectionTy,
-    Substitution, TraitId, Ty, TyData,
+    CanonicalVarKinds, ClosureId, FnDefId, ForeignDefId, GeneratorId, GenericArg, Goal, Goals,
+    ImplId, Lifetime, OpaqueTy, OpaqueTyId, ProgramClause, ProgramClauseImplication,
+    ProgramClauses, ProjectionTy, Substitution, TraitId, Ty, TyData,
 };
 use chalk_solve::rust_ir::{
     AdtDatum, AdtRepr, AssociatedTyDatum, AssociatedTyValue, AssociatedTyValueId, ClosureKind,
-    FnDefDatum, FnDefInputsAndOutputDatum, ImplDatum, ImplType, OpaqueTyDatum, TraitDatum,
-    WellKnownTrait,
+    FnDefDatum, FnDefInputsAndOutputDatum, GeneratorDatum, GeneratorWitnessDatum, ImplDatum,
+    ImplType, OpaqueTyDatum, TraitDatum, WellKnownTrait,
 };
 use chalk_solve::split::Split;
 use chalk_solve::RustIrDatabase;
@@ -37,6 +37,15 @@ pub struct Program {
 
     pub closure_kinds: BTreeMap<ClosureId<ChalkIr>, TypeKind>,
 
+    /// For each generator
+    pub generator_ids: BTreeMap<Identifier, GeneratorId<ChalkIr>>,
+
+    pub generator_kinds: BTreeMap<GeneratorId<ChalkIr>, TypeKind>,
+
+    pub generator_data: BTreeMap<GeneratorId<ChalkIr>, Arc<GeneratorDatum<ChalkIr>>>,
+
+    pub generator_witness_data: BTreeMap<GeneratorId<ChalkIr>, Arc<GeneratorWitnessDatum<ChalkIr>>>,
+
     /// From trait name to item-id. Used during lowering only.
     pub trait_ids: BTreeMap<Identifier, TraitId<ChalkIr>>,
 
@@ -380,6 +389,17 @@ impl RustIrDatabase<ChalkIr> for Program {
         self.adt_data[&id].clone()
     }
 
+    fn generator_datum(&self, id: GeneratorId<ChalkIr>) -> Arc<GeneratorDatum<ChalkIr>> {
+        self.generator_data[&id].clone()
+    }
+
+    fn generator_witness_datum(
+        &self,
+        id: GeneratorId<ChalkIr>,
+    ) -> Arc<GeneratorWitnessDatum<ChalkIr>> {
+        self.generator_witness_data[&id].clone()
+    }
+
     fn adt_repr(&self, id: AdtId<ChalkIr>) -> AdtRepr {
         self.adt_reprs[&id]
     }

chalk-ir/src/cast.rs

@@ -315,6 +315,15 @@ where
     }
 }
 
+impl<I> CastTo<TypeName<I>> for GeneratorId<I>
+where
+    I: Interner,
+{
+    fn cast_to(self, _interner: &I) -> TypeName<I> {
+        TypeName::Generator(self)
+    }
+}
+
 impl<I> CastTo<TypeName<I>> for FnDefId<I>
 where
     I: Interner,

chalk-ir/src/debug.rs

@@ -35,6 +35,13 @@ impl<I: Interner> Debug for ClosureId<I> {
     }
 }
 
+impl<I: Interner> Debug for GeneratorId<I> {
+    fn fmt(&self, fmt: &mut Formatter<'_>) -> std::fmt::Result {
+        I::debug_generator_id(*self, fmt)
+            .unwrap_or_else(|| write!(fmt, "GeneratorId({:?})", self.0))
+    }
+}
+
 impl<I: Interner> Debug for ForeignDefId<I> {
     fn fmt(&self, fmt: &mut Formatter<'_>) -> std::fmt::Result {
         I::debug_foreign_def_id(*self, fmt)
@@ -205,6 +212,8 @@ impl<I: Interner> Debug for TypeName<I> {
             TypeName::Never => write!(fmt, "Never"),
             TypeName::Array => write!(fmt, "{{array}}"),
             TypeName::Closure(id) => write!(fmt, "{{closure:{:?}}}", id),
+            TypeName::Generator(generator) => write!(fmt, "{:?}", generator),
+            TypeName::GeneratorWitness(witness) => write!(fmt, "{:?}", witness),
             TypeName::Foreign(foreign_ty) => write!(fmt, "{:?}", foreign_ty),
             TypeName::Error => write!(fmt, "{{error}}"),
         }

chalk-ir/src/fold/boring_impls.rs

@@ -303,6 +303,7 @@ id_fold!(AssocTypeId);
 id_fold!(OpaqueTyId);
 id_fold!(FnDefId);
 id_fold!(ClosureId);
+id_fold!(GeneratorId);
 id_fold!(ForeignDefId);
 
 impl<I: Interner, TI: TargetInterner<I>> SuperFold<I, TI> for ProgramClauseData<I> {