Merge #1562

1562: Continue support for .await r=matklad a=unrealhoang

- add await expr to ast and HIR Expr
- infer type for `.await`

Co-authored-by: Unreal Hoang <[email protected]>

4 files changed, 87 insertions, 0 deletions

diff --git a/crates/ra_hir/src/expr.rs b/crates/ra_hir/src/expr.rs
index 6eed67f50..4dcea19a9 100644
--- a/crates/ra_hir/src/expr.rs
+++ b/crates/ra_hir/src/expr.rs
@@ -220,6 +220,9 @@ pub enum Expr {
         expr: ExprId,
         name: Name,
     },
+    Await {
+        expr: ExprId,
+    },
     Try {
         expr: ExprId,
     },
@@ -359,6 +362,7 @@ impl Expr {
                 f(*rhs);
             }
             Expr::Field { expr, .. }
+            | Expr::Await { expr }
             | Expr::Try { expr }
             | Expr::Cast { expr, .. }
             | Expr::Ref { expr, .. }
@@ -729,6 +733,10 @@ where
                 };
                 self.alloc_expr(Expr::Field { expr, name }, syntax_ptr)
             }
+            ast::ExprKind::AwaitExpr(e) => {
+                let expr = self.collect_expr_opt(e.expr());
+                self.alloc_expr(Expr::Await { expr }, syntax_ptr)
+            }
             ast::ExprKind::TryExpr(e) => {
                 let expr = self.collect_expr_opt(e.expr());
                 self.alloc_expr(Expr::Try { expr }, syntax_ptr)
diff --git a/crates/ra_hir/src/name.rs b/crates/ra_hir/src/name.rs
index c589f8aba..6d14eea8e 100644
--- a/crates/ra_hir/src/name.rs
+++ b/crates/ra_hir/src/name.rs
@@ -118,6 +118,9 @@ pub(crate) const ITEM: Name = Name::new(SmolStr::new_inline_from_ascii(4, b"Item
 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 0e030576d..a82dff711 100644
--- a/crates/ra_hir/src/ty/infer.rs
+++ b/crates/ra_hir/src/ty/infer.rs
@@ -1114,6 +1114,25 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
                 .unwrap_or(Ty::Unknown);
                 self.insert_type_vars(ty)
             }
+            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());
                 let ty = match self.resolve_ops_try_ok() {
@@ -1368,6 +1387,28 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
             _ => 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 7b2bdeb3f..706500484 100644
--- a/crates/ra_hir/src/ty/tests.rs
+++ b/crates/ra_hir/src/ty/tests.rs
@@ -21,6 +21,41 @@ use crate::{
 // 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(
         r#"