3 files changed, 33 insertions, 14 deletions

diff --git a/crates/ra_hir_ty/src/infer/coerce.rs b/crates/ra_hir_ty/src/infer/coerce.rs
index 0f4dac45e..83c0c2c3f 100644
--- a/crates/ra_hir_ty/src/infer/coerce.rs
+++ b/crates/ra_hir_ty/src/infer/coerce.rs
@@ -134,6 +134,10 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
                 }
             }
 
+            (ty_app!(TypeCtor::Closure { .. }, params), ty_app!(TypeCtor::FnPtr { .. })) => {
+                from_ty = params[0].clone();
+            }
+
             _ => {}
         }
 
diff --git a/crates/ra_hir_ty/src/infer/expr.rs b/crates/ra_hir_ty/src/infer/expr.rs
index 924ad3e81..3af05394c 100644
--- a/crates/ra_hir_ty/src/infer/expr.rs
+++ b/crates/ra_hir_ty/src/infer/expr.rs
@@ -8,7 +8,7 @@ use hir_def::{
     expr::{Array, BinaryOp, Expr, ExprId, Literal, Statement, UnaryOp},
     path::{GenericArg, GenericArgs},
     resolver::resolver_for_expr,
-    AdtId, ContainerId, Lookup, StructFieldId,
+    AdtId, AssocContainerId, Lookup, StructFieldId,
 };
 use hir_expand::name::{name, Name};
 use ra_syntax::ast::RangeOp;
@@ -41,7 +41,7 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
 
     /// Infer type of expression with possibly implicit coerce to the expected type.
     /// Return the type after possible coercion.
-    fn infer_expr_coerce(&mut self, expr: ExprId, expected: &Expectation) -> Ty {
+    pub(super) fn infer_expr_coerce(&mut self, expr: ExprId, expected: &Expectation) -> Ty {
         let ty = self.infer_expr_inner(expr, &expected);
         let ty = if !self.coerce(&ty, &expected.ty) {
             self.result
@@ -102,7 +102,7 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
                 self.infer_expr(*body, &Expectation::has_type(Ty::unit()));
                 Ty::unit()
             }
-            Expr::Lambda { body, args, arg_types } => {
+            Expr::Lambda { body, args, ret_type, arg_types } => {
                 assert_eq!(args.len(), arg_types.len());
 
                 let mut sig_tys = Vec::new();
@@ -118,7 +118,10 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
                 }
 
                 // add return type
-                let ret_ty = self.table.new_type_var();
+                let ret_ty = match ret_type {
+                    Some(type_ref) => self.make_ty(type_ref),
+                    None => self.table.new_type_var(),
+                };
                 sig_tys.push(ret_ty.clone());
                 let sig_ty = Ty::apply(
                     TypeCtor::FnPtr { num_args: sig_tys.len() as u16 - 1 },
@@ -134,7 +137,12 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
                 // infer the body.
                 self.coerce(&closure_ty, &expected.ty);
 
-                self.infer_expr(*body, &Expectation::has_type(ret_ty));
+                let prev_ret_ty = std::mem::replace(&mut self.return_ty, ret_ty.clone());
+
+                self.infer_expr_coerce(*body, &Expectation::has_type(ret_ty));
+
+                self.return_ty = prev_ret_ty;
+
                 closure_ty
             }
             Expr::Call { callee, args } => {
@@ -192,6 +200,9 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
             Expr::Return { expr } => {
                 if let Some(expr) = expr {
                     self.infer_expr_coerce(*expr, &Expectation::has_type(self.return_ty.clone()));
+                } else {
+                    let unit = Ty::unit();
+                    self.coerce(&unit, &self.return_ty.clone());
                 }
                 Ty::simple(TypeCtor::Never)
             }
@@ -422,10 +433,14 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
                 }
             }
             Expr::Index { base, index } => {
-                let _base_ty = self.infer_expr_inner(*base, &Expectation::none());
-                let _index_ty = self.infer_expr(*index, &Expectation::none());
-                // FIXME: use `std::ops::Index::Output` to figure out the real return type
-                Ty::Unknown
+                let base_ty = self.infer_expr_inner(*base, &Expectation::none());
+                let index_ty = self.infer_expr(*index, &Expectation::none());
+
+                self.resolve_associated_type_with_params(
+                    base_ty,
+                    self.resolve_ops_index_output(),
+                    &[index_ty],
+                )
             }
             Expr::Tuple { exprs } => {
                 let mut tys = match &expected.ty {
@@ -672,7 +687,7 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
                 // add obligation for trait implementation, if this is a trait method
                 match def {
                     CallableDef::FunctionId(f) => {
-                        if let ContainerId::TraitId(trait_) = f.lookup(self.db).container {
+                        if let AssocContainerId::TraitId(trait_) = f.lookup(self.db).container {
                             // construct a TraitDef
                             let substs =
                                 a_ty.parameters.prefix(generics(self.db, trait_.into()).len());
diff --git a/crates/ra_hir_ty/src/infer/path.rs b/crates/ra_hir_ty/src/infer/path.rs
index 402a89386..ffd358367 100644
--- a/crates/ra_hir_ty/src/infer/path.rs
+++ b/crates/ra_hir_ty/src/infer/path.rs
@@ -5,7 +5,7 @@ use std::iter;
 use hir_def::{
     path::{Path, PathSegment},
     resolver::{ResolveValueResult, Resolver, TypeNs, ValueNs},
-    AssocItemId, ContainerId, Lookup,
+    AssocContainerId, AssocItemId, Lookup,
 };
 use hir_expand::name::Name;
 
@@ -209,7 +209,7 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
                     AssocItemId::TypeAliasId(_) => unreachable!(),
                 };
                 let substs = match container {
-                    ContainerId::ImplId(impl_id) => {
+                    AssocContainerId::ImplId(impl_id) => {
                         let impl_substs = Substs::build_for_def(self.db, impl_id)
                             .fill(iter::repeat_with(|| self.table.new_type_var()))
                             .build();
@@ -221,7 +221,7 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
                         self.unify(&impl_self_ty, &ty);
                         Some(substs)
                     }
-                    ContainerId::TraitId(trait_) => {
+                    AssocContainerId::TraitId(trait_) => {
                         // we're picking this method
                         let trait_substs = Substs::build_for_def(self.db, trait_)
                             .push(ty.clone())
@@ -237,7 +237,7 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
                         }));
                         Some(substs)
                     }
-                    ContainerId::ModuleId(_) => None,
+                    AssocContainerId::ContainerId(_) => None,
                 };
 
                 self.write_assoc_resolution(id, item.into());