Add tests for type inference for generators

3 files changed, 148 insertions, 0 deletions

diff --git a/crates/hir-ty/src/tests/coercion.rs b/crates/hir-ty/src/tests/coercion.rs
index bf59fadc2c3..d301595bcd9 100644
--- a/crates/hir-ty/src/tests/coercion.rs
+++ b/crates/hir-ty/src/tests/coercion.rs
@@ -295,6 +295,24 @@ fn foo() {
 }
 
 #[test]
+fn generator_yield_return_coerce() {
+    check_no_mismatches(
+        r#"
+fn test() {
+    let g = || {
+        yield &1u32;
+        yield &&1u32;
+        if true {
+            return &1u32;
+        }
+        &&1u32
+    };
+}
+        "#,
+    );
+}
+
+#[test]
 fn assign_coerce() {
     check_no_mismatches(
         r"
diff --git a/crates/hir-ty/src/tests/simple.rs b/crates/hir-ty/src/tests/simple.rs
index 4ea103e5d9e..e6ff0762caa 100644
--- a/crates/hir-ty/src/tests/simple.rs
+++ b/crates/hir-ty/src/tests/simple.rs
@@ -1918,6 +1918,88 @@ fn closure_return_inferred() {
 }
 
 #[test]
+fn generator_types_inferred() {
+    check_infer(
+        r#"
+//- minicore: generator, deref
+use core::ops::{Generator, GeneratorState};
+use core::pin::Pin;
+
+fn f(v: i64) {}
+fn test() {
+    let mut g = |r| {
+        let a = yield 0;
+        let a = yield 1;
+        let a = yield 2;
+        "return value"
+    };
+
+    match Pin::new(&mut g).resume(0usize) {
+        GeneratorState::Yielded(y) => { f(y); }
+        GeneratorState::Complete(r) => {}
+    }
+}
+        "#,
+        expect![[r#"
+            70..71 'v': i64
+            78..80 '{}': ()
+            91..362 '{     ...   } }': ()
+            101..106 'mut g': {generator}
+            109..218 '|r| { ...     }': {generator}
+            110..111 'r': usize
+            113..218 '{     ...     }': &str
+            127..128 'a': usize
+            131..138 'yield 0': usize
+            137..138 '0': i64
+            152..153 'a': usize
+            156..163 'yield 1': usize
+            162..163 '1': i64
+            177..178 'a': usize
+            181..188 'yield 2': usize
+            187..188 '2': i64
+            198..212 '"return value"': &str
+            225..360 'match ...     }': ()
+            231..239 'Pin::new': fn new<&mut {generator}>(&mut {generator}) -> Pin<&mut {generator}>
+            231..247 'Pin::n...mut g)': Pin<&mut {generator}>
+            231..262 'Pin::n...usize)': GeneratorState<i64, &str>
+            240..246 '&mut g': &mut {generator}
+            245..246 'g': {generator}
+            255..261 '0usize': usize
+            273..299 'Genera...ded(y)': GeneratorState<i64, &str>
+            297..298 'y': i64
+            303..312 '{ f(y); }': ()
+            305..306 'f': fn f(i64)
+            305..309 'f(y)': ()
+            307..308 'y': i64
+            321..348 'Genera...ete(r)': GeneratorState<i64, &str>
+            346..347 'r': &str
+            352..354 '{}': ()
+        "#]],
+    );
+}
+
+#[test]
+fn generator_resume_yield_return_unit() {
+    check_no_mismatches(
+        r#"
+//- minicore: generator, deref
+use core::ops::{Generator, GeneratorState};
+use core::pin::Pin;
+fn test() {
+    let mut g = || {
+        let () = yield;
+    };
+
+    match Pin::new(&mut g).resume(()) {
+        GeneratorState::Yielded(()) => {}
+        GeneratorState::Complete(()) => {}
+    }
+}
+        "#,
+    );
+}
+
+#[test]
 fn fn_pointer_return() {
     check_infer(
         r#"
diff --git a/crates/test-utils/src/minicore.rs b/crates/test-utils/src/minicore.rs
index 6df29db4745..10386b5b7bc 100644
--- a/crates/test-utils/src/minicore.rs
+++ b/crates/test-utils/src/minicore.rs
@@ -37,6 +37,7 @@
 //!     add:
 //!     as_ref: sized
 //!     drop:
+//!     generator: pin
 
 pub mod marker {
     // region:sized
@@ -182,6 +183,19 @@ pub mod ops {
             type Target: ?Sized;
             fn deref(&self) -> &Self::Target;
         }
+
+        impl<T: ?Sized> Deref for &T {
+            type Target = T;
+            fn deref(&self) -> &T {
+                loop {}
+            }
+        }
+        impl<T: ?Sized> Deref for &mut T {
+            type Target = T;
+            fn deref(&self) -> &T {
+                loop {}
+            }
+        }
         // region:deref_mut
         #[lang = "deref_mut"]
         pub trait DerefMut: Deref {
@@ -347,6 +361,27 @@ pub mod ops {
         fn add(self, rhs: Rhs) -> Self::Output;
     }
     // endregion:add
+
+    // region:generator
+    mod generator {
+        use crate::pin::Pin;
+
+        #[lang = "generator"]
+        pub trait Generator<R = ()> {
+            type Yield;
+            #[lang = "generator_return"]
+            type Return;
+            fn resume(self: Pin<&mut Self>, arg: R) -> GeneratorState<Self::Yield, Self::Return>;
+        }
+
+        #[lang = "generator_state"]
+        pub enum GeneratorState<Y, R> {
+            Yielded(Y),
+            Complete(R),
+        }
+    }
+    pub use self::generator::{Generator, GeneratorState};
+    // endregion:generator
 }
 
 // region:eq
@@ -455,6 +490,19 @@ pub mod pin {
     pub struct Pin<P> {
         pointer: P,
     }
+    impl<P> Pin<P> {
+        pub fn new(pointer: P) -> Pin<P> {
+            loop {}
+        }
+    }
+    // region:deref
+    impl<P: crate::ops::Deref> crate::ops::Deref for Pin<P> {
+        type Target = P::Target;
+        fn deref(&self) -> &P::Target {
+            loop {}
+        }
+    }
+    // endregion:deref
 }
 // endregion:pin