From bacf926a77a3e4040ccaf64cce1b1517e41a8d47 Mon Sep 17 00:00:00 2001
From: Unreal Hoang <unrealhoang@gmail.com>
Date: Sat, 20 Jul 2019 20:11:18 +0900
Subject: [PATCH] infer type for await

by projecting inner_ty to Future::Output alias
---
 crates/ra_hir/src/name.rs     |  3 +++
 crates/ra_hir/src/ty/infer.rs | 42 +++++++++++++++++++++++++++++++++--
 crates/ra_hir/src/ty/tests.rs | 35 +++++++++++++++++++++++++++++
 3 files changed, 78 insertions(+), 2 deletions(-)

diff --git a/crates/ra_hir/src/name.rs b/crates/ra_hir/src/name.rs
index c589f8aba6e..6d14eea8ecf 100644
--- a/crates/ra_hir/src/name.rs
+++ b/crates/ra_hir/src/name.rs
@@ -118,6 +118,9 @@ fn as_name(&self) -> Name {
 pub(crate) const OPS: Name = Name::new(SmolStr::new_inline_from_ascii(3, b"ops"));
 pub(crate) const TRY: Name = Name::new(SmolStr::new_inline_from_ascii(3, b"Try"));
 pub(crate) const OK: Name = Name::new(SmolStr::new_inline_from_ascii(2, b"Ok"));
+pub(crate) const FUTURE_MOD: Name = Name::new(SmolStr::new_inline_from_ascii(6, b"future"));
+pub(crate) const FUTURE_TYPE: Name = Name::new(SmolStr::new_inline_from_ascii(6, b"Future"));
+pub(crate) const OUTPUT: Name = Name::new(SmolStr::new_inline_from_ascii(6, b"Output"));
 
 fn resolve_name(text: &SmolStr) -> SmolStr {
     let raw_start = "r#";
diff --git a/crates/ra_hir/src/ty/infer.rs b/crates/ra_hir/src/ty/infer.rs
index 6eae595a91e..a82dff711db 100644
--- a/crates/ra_hir/src/ty/infer.rs
+++ b/crates/ra_hir/src/ty/infer.rs
@@ -1114,8 +1114,24 @@ fn infer_expr(&mut self, tgt_expr: ExprId, expected: &Expectation) -> Ty {
                 .unwrap_or(Ty::Unknown);
                 self.insert_type_vars(ty)
             }
-            Expr::Await { .. } => {
-                Ty::Unknown
+            Expr::Await { expr } => {
+                let inner_ty = self.infer_expr(*expr, &Expectation::none());
+                let ty = match self.resolve_future_future_output() {
+                    Some(future_future_output_alias) => {
+                        let ty = self.new_type_var();
+                        let projection = ProjectionPredicate {
+                            ty: ty.clone(),
+                            projection_ty: ProjectionTy {
+                                associated_ty: future_future_output_alias,
+                                parameters: vec![inner_ty].into(),
+                            },
+                        };
+                        self.obligations.push(Obligation::Projection(projection));
+                        self.resolve_ty_as_possible(&mut vec![], ty)
+                    }
+                    None => Ty::Unknown,
+                };
+                ty
             }
             Expr::Try { expr } => {
                 let inner_ty = self.infer_expr(*expr, &Expectation::none());
@@ -1371,6 +1387,28 @@ fn resolve_ops_try_ok(&self) -> Option<TypeAlias> {
             _ => None,
         }
     }
+
+    fn resolve_future_future_output(&self) -> Option<TypeAlias> {
+        let future_future_path = Path {
+            kind: PathKind::Abs,
+            segments: vec![
+                PathSegment { name: name::STD, args_and_bindings: None },
+                PathSegment { name: name::FUTURE_MOD, args_and_bindings: None },
+                PathSegment { name: name::FUTURE_TYPE, args_and_bindings: None },
+            ],
+        };
+
+        match self
+            .resolver
+            .resolve_path_segments(self.db, &future_future_path)
+            .into_fully_resolved()
+        {
+            PerNs { types: Some(Def(Trait(trait_))), .. } => {
+                Some(trait_.associated_type_by_name(self.db, name::OUTPUT)?)
+            }
+            _ => None,
+        }
+    }
 }
 
 /// The ID of a type variable.
diff --git a/crates/ra_hir/src/ty/tests.rs b/crates/ra_hir/src/ty/tests.rs
index 265740e5420..d4bfcb888fa 100644
--- a/crates/ra_hir/src/ty/tests.rs
+++ b/crates/ra_hir/src/ty/tests.rs
@@ -20,6 +20,41 @@
 // against snapshots of the expected results using insta. Use cargo-insta to
 // update the snapshots.
 
+#[test]
+fn infer_await() {
+    let (mut db, pos) = MockDatabase::with_position(
+        r#"
+//- /main.rs
+
+struct IntFuture;
+
+impl Future for IntFuture {
+    type Output = u64;
+}
+
+fn test() {
+    let r = IntFuture;
+    let v = r.await;
+    v<|>;
+}
+
+//- /std.rs
+#[prelude_import] use future::*;
+mod future {
+    trait Future {
+        type Output;
+    }
+}
+
+"#,
+    );
+    db.set_crate_graph_from_fixture(crate_graph! {
+        "main": ("/main.rs", ["std"]),
+        "std": ("/std.rs", []),
+    });
+    assert_eq!("u64", type_at_pos(&db, pos));
+}
+
 #[test]
 fn infer_try() {
     let (mut db, pos) = MockDatabase::with_position(