Add std::ops::Index support for infering

5 files changed, 63 insertions, 5 deletions

diff --git a/crates/ra_hir_def/src/path.rs b/crates/ra_hir_def/src/path.rs
index 3b26e8337..9e37ac416 100644
--- a/crates/ra_hir_def/src/path.rs
+++ b/crates/ra_hir_def/src/path.rs
@@ -254,6 +254,7 @@ macro_rules! __known_path {
     (std::ops::Try) => {};
     (std::ops::Neg) => {};
     (std::ops::Not) => {};
+    (std::ops::Index) => {};
     ($path:path) => {
         compile_error!("Please register your known path in the path module")
     };
diff --git a/crates/ra_hir_expand/src/name.rs b/crates/ra_hir_expand/src/name.rs
index 59d8214fd..fd02ffa4e 100644
--- a/crates/ra_hir_expand/src/name.rs
+++ b/crates/ra_hir_expand/src/name.rs
@@ -161,6 +161,7 @@ pub mod known {
         Range,
         Neg,
         Not,
+        Index,
         // Builtin macros
         file,
         column,
diff --git a/crates/ra_hir_ty/src/infer.rs b/crates/ra_hir_ty/src/infer.rs
index 98ba05fc2..14bfdde3d 100644
--- a/crates/ra_hir_ty/src/infer.rs
+++ b/crates/ra_hir_ty/src/infer.rs
@@ -363,14 +363,28 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
     }
 
     fn resolve_associated_type(&mut self, inner_ty: Ty, assoc_ty: Option<TypeAliasId>) -> Ty {
+        self.resolve_associated_type_with_params(inner_ty, assoc_ty, &[])
+    }
+
+    fn resolve_associated_type_with_params(
+        &mut self,
+        inner_ty: Ty,
+        assoc_ty: Option<TypeAliasId>,
+        params: &[Ty],
+    ) -> Ty {
         match assoc_ty {
             Some(res_assoc_ty) => {
                 let ty = self.table.new_type_var();
+                let mut builder = Substs::builder(1 + params.len()).push(inner_ty);
+                for ty in params {
+                    builder = builder.push(ty.clone());
+                }
+
                 let projection = ProjectionPredicate {
                     ty: ty.clone(),
                     projection_ty: ProjectionTy {
                         associated_ty: res_assoc_ty,
-                        parameters: Substs::single(inner_ty),
+                        parameters: builder.build(),
                     },
                 };
                 self.obligations.push(Obligation::Projection(projection));
@@ -517,6 +531,12 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
         let struct_ = self.resolver.resolve_known_struct(self.db, &path)?;
         Some(struct_.into())
     }
+
+    fn resolve_ops_index_output(&self) -> Option<TypeAliasId> {
+        let path = path![std::ops::Index];
+        let trait_ = self.resolver.resolve_known_trait(self.db, &path)?;
+        self.db.trait_data(trait_).associated_type_by_name(&name![Output])
+    }
 }
 
 /// The kinds of placeholders we need during type inference. There's separate
diff --git a/crates/ra_hir_ty/src/infer/expr.rs b/crates/ra_hir_ty/src/infer/expr.rs
index 924ad3e81..011c6c5c6 100644
--- a/crates/ra_hir_ty/src/infer/expr.rs
+++ b/crates/ra_hir_ty/src/infer/expr.rs
@@ -422,10 +422,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 {
diff --git a/crates/ra_hir_ty/src/tests/traits.rs b/crates/ra_hir_ty/src/tests/traits.rs
index 802937cb0..2d92a5eec 100644
--- a/crates/ra_hir_ty/src/tests/traits.rs
+++ b/crates/ra_hir_ty/src/tests/traits.rs
@@ -427,6 +427,38 @@ fn indexing_arrays() {
 }
 
 #[test]
+fn infer_ops_index() {
+    let (db, pos) = TestDB::with_position(
+        r#"
+//- /main.rs crate:main deps:std
+
+struct Bar;
+struct Foo;
+
+impl std::ops::Index<u32> for Bar {
+    type Output = Foo;
+}
+
+fn test() {
+    let a = Bar;
+    let b = a[1];
+    b<|>;
+}
+
+//- /std.rs crate:std
+
+#[prelude_import] use ops::*;
+mod ops {
+    pub trait Index<Idx> {
+        type Output;
+    }
+}
+"#,
+    );
+    assert_eq!("Foo", type_at_pos(&db, pos));
+}
+
+#[test]
 fn deref_trait() {
     let t = type_at(
         r#"