Do autoderef for indexing

4 files changed, 74 insertions, 7 deletions

diff --git a/crates/ra_hir_ty/src/infer.rs b/crates/ra_hir_ty/src/infer.rs
index 569d46cc3..377f44fa7 100644
--- a/crates/ra_hir_ty/src/infer.rs
+++ b/crates/ra_hir_ty/src/infer.rs
@@ -28,7 +28,7 @@ use hir_def::{
     path::{path, Path},
     resolver::{HasResolver, Resolver, TypeNs},
     type_ref::{Mutability, TypeRef},
-    AdtId, AssocItemId, DefWithBodyId, FunctionId, StructFieldId, TypeAliasId, VariantId,
+    AdtId, AssocItemId, DefWithBodyId, FunctionId, StructFieldId, TraitId, TypeAliasId, VariantId,
 };
 use hir_expand::{diagnostics::DiagnosticSink, name::name};
 use ra_arena::map::ArenaMap;
@@ -540,8 +540,12 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
         Some(struct_.into())
     }
 
+    fn resolve_ops_index(&self) -> Option<TraitId> {
+        self.resolve_lang_item("index")?.as_trait()
+    }
+
     fn resolve_ops_index_output(&self) -> Option<TypeAliasId> {
-        let trait_ = self.resolve_lang_item("index")?.as_trait()?;
+        let trait_ = self.resolve_ops_index()?;
         self.db.trait_data(trait_).associated_type_by_name(&name![Output])
     }
 }
diff --git a/crates/ra_hir_ty/src/infer/expr.rs b/crates/ra_hir_ty/src/infer/expr.rs
index 3db5b2b51..e89cc7298 100644
--- a/crates/ra_hir_ty/src/infer/expr.rs
+++ b/crates/ra_hir_ty/src/infer/expr.rs
@@ -429,11 +429,27 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
                 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],
-                )
+                if let (Some(index_trait), Some(krate)) =
+                    (self.resolve_ops_index(), self.resolver.krate())
+                {
+                    let canonicalized = self.canonicalizer().canonicalize_ty(base_ty);
+                    let self_ty = method_resolution::resolve_indexing_op(
+                        self.db,
+                        &canonicalized.value,
+                        self.trait_env.clone(),
+                        krate,
+                        index_trait,
+                    );
+                    let self_ty =
+                        self_ty.map_or(Ty::Unknown, |t| canonicalized.decanonicalize_ty(t.value));
+                    self.resolve_associated_type_with_params(
+                        self_ty,
+                        self.resolve_ops_index_output(),
+                        &[index_ty],
+                    )
+                } else {
+                    Ty::Unknown
+                }
             }
             Expr::Tuple { exprs } => {
                 let mut tys = match &expected.ty {
diff --git a/crates/ra_hir_ty/src/method_resolution.rs b/crates/ra_hir_ty/src/method_resolution.rs
index 7b0ff8161..74b908c2e 100644
--- a/crates/ra_hir_ty/src/method_resolution.rs
+++ b/crates/ra_hir_ty/src/method_resolution.rs
@@ -447,6 +447,25 @@ fn iterate_inherent_methods<T>(
     None
 }
 
+/// Returns the self type for the index trait call.
+pub fn resolve_indexing_op(
+    db: &impl HirDatabase,
+    ty: &Canonical<Ty>,
+    env: Arc<TraitEnvironment>,
+    krate: CrateId,
+    index_trait: TraitId,
+) -> Option<Canonical<Ty>> {
+    let ty = InEnvironment { value: ty.clone(), environment: env.clone() };
+    let deref_chain = autoderef_method_receiver(db, krate, ty);
+    for ty in deref_chain {
+        let goal = generic_implements_goal(db, env.clone(), index_trait, ty.clone());
+        if db.trait_solve(krate, goal).is_some() {
+            return Some(ty);
+        }
+    }
+    None
+}
+
 fn is_valid_candidate(
     db: &impl HirDatabase,
     name: Option<&Name>,
diff --git a/crates/ra_hir_ty/src/tests/traits.rs b/crates/ra_hir_ty/src/tests/traits.rs
index 7d796d0b9..547010b35 100644
--- a/crates/ra_hir_ty/src/tests/traits.rs
+++ b/crates/ra_hir_ty/src/tests/traits.rs
@@ -568,6 +568,34 @@ mod ops {
 }
 
 #[test]
+fn infer_ops_index_autoderef() {
+    let (db, pos) = TestDB::with_position(
+        r#"
+//- /main.rs crate:main deps:std
+fn test() {
+    let a = &[1u32, 2, 3];
+    let b = a[1];
+    b<|>;
+}
+
+//- /std.rs crate:std
+impl<T> ops::Index<u32> for [T] {
+    type Output = T;
+}
+
+#[prelude_import] use ops::*;
+mod ops {
+    #[lang = "index"]
+    pub trait Index<Idx> {
+        type Output;
+    }
+}
+"#,
+    );
+    assert_eq!("u32", type_at_pos(&db, pos));
+}
+
+#[test]
 fn deref_trait() {
     let t = type_at(
         r#"