Make the Generator trait well-known

nrc · nrc · commit ce31db804dbf · 2021-08-17T09:00:00.000+12:00
Signed-off-by: Nick Cameron &lt;nrc@ncameron.org&gt;
diff --git a/chalk-integration/src/lowering.rs b/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,
         }
     }
 }
diff --git a/chalk-parse/src/ast.rs b/chalk-parse/src/ast.rs
@@ -158,6 +158,7 @@ pub enum WellKnownTrait {
     Unpin,
     CoerceUnsized,
     DiscriminantKind,
+    Generator,
 }
 
 #[derive(Clone, PartialEq, Eq, Debug)]
diff --git a/chalk-parse/src/parser.lalrpop b/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 = {
diff --git a/chalk-solve/src/clauses/builtin_traits.rs b/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(()),
     }
 }
diff --git a/chalk-solve/src/clauses/builtin_traits/generator.rs b/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(()),
+    }
+}
diff --git a/chalk-solve/src/display/items.rs b/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)?;
         }
diff --git a/chalk-solve/src/rust_ir.rs b/chalk-solve/src/rust_ir.rs
@@ -262,6 +262,7 @@ pub enum WellKnownTrait {
     Unpin,
     CoerceUnsized,
     DiscriminantKind,
+    Generator,
 }
 
 chalk_ir::const_visit!(WellKnownTrait);
diff --git a/chalk-solve/src/wf.rs b/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 {
diff --git a/tests/test/generators.rs b/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

Original file line number	Diff line number	Diff line change
`@@ -1133,6 +1133,7 @@ impl Lower for WellKnownTrait {`
`1133`	`1133`	`WellKnownTrait::Unpin => rust_ir::WellKnownTrait::Unpin,`
`1134`	`1134`	`WellKnownTrait::CoerceUnsized => rust_ir::WellKnownTrait::CoerceUnsized,`
`1135`	`1135`	`WellKnownTrait::DiscriminantKind => rust_ir::WellKnownTrait::DiscriminantKind,`
	`1136`	`+ WellKnownTrait::Generator => rust_ir::WellKnownTrait::Generator,`
`1136`	`1137`	`}`
`1137`	`1138`	`}`
`1138`	`1139`	`}`
Original file line number	Diff line number	Diff line change
`@@ -158,6 +158,7 @@ pub enum WellKnownTrait {`
`158`	`158`	`Unpin,`
`159`	`159`	`CoerceUnsized,`
`160`	`160`	`DiscriminantKind,`
	`161`	`+ Generator,`
`161`	`162`	`}`
`162`	`163`
`163`	`164`	`#[derive(Clone, PartialEq, Eq, Debug)]`
Original file line number	Diff line number	Diff line change
`@@ -6,6 +6,7 @@ mod clone;`
`6`	`6`	`mod copy;`
`7`	`7`	`mod discriminant_kind;`
`8`	`8`	`mod fn_family;`
	`9`	`+mod generator;`
`9`	`10`	`mod sized;`
`10`	`11`	`mod unsize;`
`11`	`12`
`@@ -37,13 +38,16 @@ pub fn add_builtin_program_clauses<I: Interner>(`
`37`	`38`	`clone::add_clone_program_clauses(db, builder, trait_ref, ty, binders)?;`
`38`	`39`	`}`
`39`	`40`	`WellKnownTrait::FnOnce \| WellKnownTrait::FnMut \| WellKnownTrait::Fn => {`
`40`		`- fn_family::add_fn_trait_program_clauses(db, builder, well_known, self_ty)?`
	`41`	`+ fn_family::add_fn_trait_program_clauses(db, builder, well_known, self_ty)?;`
`41`	`42`	`}`
`42`	`43`	`WellKnownTrait::Unsize => {`
`43`	`44`	`unsize::add_unsize_program_clauses(db, builder, trait_ref, ty)`
`44`	`45`	`}`
`45`	`46`	`// DiscriminantKind is automatically implemented for all types`
`46`	`47`	`WellKnownTrait::DiscriminantKind => builder.push_fact(trait_ref),`
	`48`	`+ WellKnownTrait::Generator => {`
	`49`	`+ generator::add_generator_program_clauses(db, builder, self_ty)?;`
	`50`	`+ }`
`47`	`51`	`// There are no builtin impls provided for the following traits:`
`48`	`52`	`WellKnownTrait::Unpin \| WellKnownTrait::Drop \| WellKnownTrait::CoerceUnsized => (),`
`49`	`53`	`}`
`@@ -73,6 +77,13 @@ pub fn add_builtin_assoc_program_clauses<I: Interner>(`
`73`	`77`	`WellKnownTrait::DiscriminantKind => {`
`74`	`78`	`discriminant_kind::add_discriminant_clauses(db, builder, self_ty)`
`75`	`79`	`}`
	`80`	`+ WellKnownTrait::Generator => {`
	`81`	`+ let generalized = generalize::Generalize::apply(db.interner(), self_ty);`
	`82`	`+`
	`83`	`+ builder.push_binders(generalized, \|builder, self_ty\| {`
	`84`	`+ generator::add_generator_program_clauses(db, builder, self_ty)`
	`85`	`+ })`
	`86`	`+ }`
`76`	`87`	`_ => Ok(()),`
`77`	`88`	`}`
`78`	`89`	`}`
Original file line number	Diff line number	Diff line change
`@@ -203,6 +203,7 @@ impl<I: Interner> RenderAsRust<I> for TraitDatum<I> {`
`203`	`203`	`WellKnownTrait::Unpin => "unpin",`
`204`	`204`	`WellKnownTrait::CoerceUnsized => "coerce_unsized",`
`205`	`205`	`WellKnownTrait::DiscriminantKind => "discriminant_kind",`
	`206`	`+ WellKnownTrait::Generator => "generator",`
`206`	`207`	`};`
`207`	`208`	`writeln!(f, "#[lang({})]", name)?;`
`208`	`209`	`}`
Original file line number	Diff line number	Diff line change
`@@ -262,6 +262,7 @@ pub enum WellKnownTrait {`
`262`	`262`	`Unpin,`
`263`	`263`	`CoerceUnsized,`
`264`	`264`	`DiscriminantKind,`
	`265`	`+ Generator,`
`265`	`266`	`}`
`266`	`267`
`267`	`268`	`chalk_ir::const_visit!(WellKnownTrait);`