Auto merge of #723 - nrc:auto-foreign, r=jackh726

Make the Generator trait well-known

r? `@jackh726` There's quite a lot of cargo culting of code from function trait support which I don't fully understand, so I'd appreciate a close review.

As discussed on Zulip, this adds initial support for the Generator trait as a well-known trait. I'll work on `DispatchFromDyn` next, but want to check I'm doing the right things first.

Author: bors

chalk-integration/src/lowering.rs

@@ -1133,6 +1133,7 @@ impl Lower for WellKnownTrait {
             WellKnownTrait::Unpin => rust_ir::WellKnownTrait::Unpin,
             WellKnownTrait::CoerceUnsized => rust_ir::WellKnownTrait::CoerceUnsized,
             WellKnownTrait::DiscriminantKind => rust_ir::WellKnownTrait::DiscriminantKind,
+            WellKnownTrait::Generator => rust_ir::WellKnownTrait::Generator,
         }
     }
 }

chalk-parse/src/ast.rs

@@ -158,6 +158,7 @@ pub enum WellKnownTrait {
     Unpin,
     CoerceUnsized,
     DiscriminantKind,
+    Generator,
 }
 
 #[derive(Clone, PartialEq, Eq, Debug)]

chalk-parse/src/parser.lalrpop

@@ -66,6 +66,7 @@ WellKnownTrait: WellKnownTrait = {
      "#" "[" "lang" "(" "unpin" ")" "]" => WellKnownTrait::Unpin,
      "#" "[" "lang" "(" "coerce_unsized" ")" "]" => WellKnownTrait::CoerceUnsized,
      "#" "[" "lang" "(" "discriminant_kind" ")" "]" => WellKnownTrait::DiscriminantKind,
+     "#" "[" "lang" "(" "generator" ")" "]" => WellKnownTrait::Generator,
 };
 
 AdtReprAttr: AdtReprAttr = {

chalk-solve/src/clauses/builtin_traits.rs

@@ -6,6 +6,7 @@ mod clone;
 mod copy;
 mod discriminant_kind;
 mod fn_family;
+mod generator;
 mod sized;
 mod unsize;
 
@@ -37,13 +38,16 @@ pub fn add_builtin_program_clauses<I: Interner>(
                 clone::add_clone_program_clauses(db, builder, trait_ref, ty, binders)?;
             }
             WellKnownTrait::FnOnce | WellKnownTrait::FnMut | WellKnownTrait::Fn => {
-                fn_family::add_fn_trait_program_clauses(db, builder, well_known, self_ty)?
+                fn_family::add_fn_trait_program_clauses(db, builder, well_known, self_ty)?;
             }
             WellKnownTrait::Unsize => {
                 unsize::add_unsize_program_clauses(db, builder, trait_ref, ty)
             }
             // DiscriminantKind is automatically implemented for all types
             WellKnownTrait::DiscriminantKind => builder.push_fact(trait_ref),
+            WellKnownTrait::Generator => {
+                generator::add_generator_program_clauses(db, builder, self_ty)?;
+            }
             // There are no builtin impls provided for the following traits:
             WellKnownTrait::Unpin | WellKnownTrait::Drop | WellKnownTrait::CoerceUnsized => (),
         }
@@ -73,6 +77,13 @@ pub fn add_builtin_assoc_program_clauses<I: Interner>(
         WellKnownTrait::DiscriminantKind => {
             discriminant_kind::add_discriminant_clauses(db, builder, self_ty)
         }
+        WellKnownTrait::Generator => {
+            let generalized = generalize::Generalize::apply(db.interner(), self_ty);
+
+            builder.push_binders(generalized, |builder, self_ty| {
+                generator::add_generator_program_clauses(db, builder, self_ty)
+            })
+        }
         _ => Ok(()),
     }
 }

chalk-solve/src/clauses/builtin_traits/generator.rs

@@ -0,0 +1,76 @@
+use crate::clauses::ClauseBuilder;
+use crate::rust_ir::WellKnownTrait;
+use crate::{Interner, RustIrDatabase, TraitRef};
+use chalk_ir::cast::Cast;
+use chalk_ir::{AliasTy, Floundered, Normalize, ProjectionTy, Substitution, Ty, TyKind};
+
+/// Add implicit impls of the generator trait, i.e., add a clause that all generators implement
+/// `Generator` and clauses for `Generator`'s associated types.
+pub fn add_generator_program_clauses<I: Interner>(
+    db: &dyn RustIrDatabase<I>,
+    builder: &mut ClauseBuilder<'_, I>,
+    self_ty: Ty<I>,
+) -> Result<(), Floundered> {
+    let interner = db.interner();
+
+    match self_ty.kind(interner) {
+        TyKind::Generator(id, substitution) => {
+            let generator_datum = db.generator_datum(*id);
+            let generator_io_datum = generator_datum
+                .input_output
+                .clone()
+                .substitute(interner, &substitution);
+
+            let trait_id = db.well_known_trait_id(WellKnownTrait::Generator).unwrap();
+            let trait_datum = db.trait_datum(trait_id);
+            assert_eq!(
+                trait_datum.associated_ty_ids.len(),
+                2,
+                "Generator trait should have exactly two associated types, found {:?}",
+                trait_datum.associated_ty_ids
+            );
+
+            let substitution = Substitution::from_iter(
+                interner,
+                &[
+                    self_ty.cast(interner),
+                    generator_io_datum.resume_type.cast(interner),
+                ],
+            );
+
+            // generator: Generator<resume_type>
+            builder.push_fact(TraitRef {
+                trait_id,
+                substitution: substitution.clone(),
+            });
+
+            // `Generator::Yield`
+            let yield_id = trait_datum.associated_ty_ids[0];
+            let yield_alias = AliasTy::Projection(ProjectionTy {
+                associated_ty_id: yield_id,
+                substitution: substitution.clone(),
+            });
+            builder.push_fact(Normalize {
+                alias: yield_alias,
+                ty: generator_io_datum.yield_type,
+            });
+
+            // `Generator::Return`
+            let return_id = trait_datum.associated_ty_ids[1];
+            let return_alias = AliasTy::Projection(ProjectionTy {
+                associated_ty_id: return_id,
+                substitution,
+            });
+            builder.push_fact(Normalize {
+                alias: return_alias,
+                ty: generator_io_datum.return_type,
+            });
+
+            Ok(())
+        }
+
+        // Generator trait is non-enumerable
+        TyKind::InferenceVar(..) | TyKind::BoundVar(_) | TyKind::Alias(..) => Err(Floundered),
+        _ => Ok(()),
+    }
+}

chalk-solve/src/display/items.rs

@@ -203,6 +203,7 @@ impl<I: Interner> RenderAsRust<I> for TraitDatum<I> {
                 WellKnownTrait::Unpin => "unpin",
                 WellKnownTrait::CoerceUnsized => "coerce_unsized",
                 WellKnownTrait::DiscriminantKind => "discriminant_kind",
+                WellKnownTrait::Generator => "generator",
             };
             writeln!(f, "#[lang({})]", name)?;
         }

chalk-solve/src/rust_ir.rs

@@ -262,6 +262,7 @@ pub enum WellKnownTrait {
     Unpin,
     CoerceUnsized,
     DiscriminantKind,
+    Generator,
 }
 
 chalk_ir::const_visit!(WellKnownTrait);

chalk-solve/src/wf.rs

@@ -426,7 +426,8 @@ where
             | WellKnownTrait::FnMut
             | WellKnownTrait::Unsize
             | WellKnownTrait::Sized
-            | WellKnownTrait::DiscriminantKind => false,
+            | WellKnownTrait::DiscriminantKind
+            | WellKnownTrait::Generator => false,
         };
 
         if is_legal {

tests/test/generators.rs

@@ -6,6 +6,11 @@ fn generator_test() {
         program {
             #[auto] trait Send { }
 
+            #[lang(generator)]
+            trait Generator<R> {
+                type Yield;
+                type Return;
+            }
 
             struct StructOne {}
             struct NotSend {}
@@ -37,6 +42,10 @@ fn generator_test() {
                 witnesses []
             }
 
+            generator gen_with_types<U>[resume = U, yield = StructOne] -> NotSend {
+                upvars []
+                witnesses []
+            }
         }
 
         goal {
@@ -51,6 +60,36 @@ fn generator_test() {
             "Unique"
         }
 
+        goal {
+            empty_gen: Generator<()>
+        } yields {
+            "Unique"
+        }
+
+        goal {
+            forall<T> {
+                gen_with_types<T>: Generator<T>
+            }
+        } yields {
+            "Unique"
+        }
+
+        goal {
+            forall<T> {
+                Normalize(<gen_with_types<T> as Generator<T>>::Yield -> StructOne)
+            }
+        } yields {
+            "Unique"
+        }
+
+        goal {
+            forall<T> {
+                Normalize(<gen_with_types<T> as Generator<T>>::Return -> NotSend)
+            }
+        } yields {
+            "Unique"
+        }
+
         goal {
             forall<T> {
                 upvar_lifetime_restrict<T>: Send