Auto merge of #16085 - Austaras:master, r=Veykril

fix: try obligation of `IndexMut` when infer

Closes #15842.

This issue arises because `K` is ambiguous if only inferred from `Index` trait, but is unique if inferred from `IndexMut`, but r-a doesn't use this info.

8 files changed, 74 insertions, 12 deletions

diff --git a/crates/hir-def/src/body/lower.rs b/crates/hir-def/src/body/lower.rs
index a45ec844aba..bc4da360c5a 100644
--- a/crates/hir-def/src/body/lower.rs
+++ b/crates/hir-def/src/body/lower.rs
@@ -622,7 +622,8 @@ impl ExprCollector<'_> {
             ast::Expr::IndexExpr(e) => {
                 let base = self.collect_expr_opt(e.base());
                 let index = self.collect_expr_opt(e.index());
-                self.alloc_expr(Expr::Index { base, index }, syntax_ptr)
+                let is_assignee_expr = self.is_lowering_assignee_expr;
+                self.alloc_expr(Expr::Index { base, index, is_assignee_expr }, syntax_ptr)
             }
             ast::Expr::RangeExpr(e) => {
                 let lhs = e.start().map(|lhs| self.collect_expr(lhs));
diff --git a/crates/hir-def/src/body/pretty.rs b/crates/hir-def/src/body/pretty.rs
index 6ecf1c20d6c..02b19ade44b 100644
--- a/crates/hir-def/src/body/pretty.rs
+++ b/crates/hir-def/src/body/pretty.rs
@@ -376,7 +376,7 @@ impl Printer<'_> {
                     w!(self, ") ");
                 }
             }
-            Expr::Index { base, index } => {
+            Expr::Index { base, index, is_assignee_expr: _ } => {
                 self.print_expr(*base);
                 w!(self, "[");
                 self.print_expr(*index);
diff --git a/crates/hir-def/src/hir.rs b/crates/hir-def/src/hir.rs
index 591ee77c70a..5890e818c46 100644
--- a/crates/hir-def/src/hir.rs
+++ b/crates/hir-def/src/hir.rs
@@ -265,6 +265,7 @@ pub enum Expr {
     Index {
         base: ExprId,
         index: ExprId,
+        is_assignee_expr: bool,
     },
     Closure {
         args: Box<[PatId]>,
@@ -432,7 +433,7 @@ impl Expr {
                     f(rhs);
                 }
             }
-            Expr::Index { base, index } => {
+            Expr::Index { base, index, .. } => {
                 f(*base);
                 f(*index);
             }
diff --git a/crates/hir-ty/src/infer/closure.rs b/crates/hir-ty/src/infer/closure.rs
index af74df1032c..58b4f29ec8c 100644
--- a/crates/hir-ty/src/infer/closure.rs
+++ b/crates/hir-ty/src/infer/closure.rs
@@ -598,7 +598,7 @@ impl InferenceContext<'_> {
                     self.consume_expr(expr);
                 }
             }
-            Expr::Index { base, index } => {
+            Expr::Index { base, index, is_assignee_expr: _ } => {
                 self.select_from_expr(*base);
                 self.consume_expr(*index);
             }
diff --git a/crates/hir-ty/src/infer/expr.rs b/crates/hir-ty/src/infer/expr.rs
index 84954ca7e90..b8a7d3ebf79 100644
--- a/crates/hir-ty/src/infer/expr.rs
+++ b/crates/hir-ty/src/infer/expr.rs
@@ -744,7 +744,7 @@ impl InferenceContext<'_> {
                     (RangeOp::Inclusive, _, None) => self.err_ty(),
                 }
             }
-            Expr::Index { base, index } => {
+            Expr::Index { base, index, is_assignee_expr } => {
                 let base_ty = self.infer_expr_inner(*base, &Expectation::none());
                 let index_ty = self.infer_expr(*index, &Expectation::none());
 
@@ -772,11 +772,24 @@ impl InferenceContext<'_> {
                             .build();
                         self.write_method_resolution(tgt_expr, func, substs);
                     }
-                    self.resolve_associated_type_with_params(
-                        self_ty,
-                        self.resolve_ops_index_output(),
-                        &[index_ty.cast(Interner)],
-                    )
+                    let assoc = self.resolve_ops_index_output();
+                    let res = self.resolve_associated_type_with_params(
+                        self_ty.clone(),
+                        assoc,
+                        &[index_ty.clone().cast(Interner)],
+                    );
+
+                    if *is_assignee_expr {
+                        if let Some(index_trait) = self.resolve_lang_trait(LangItem::IndexMut) {
+                            let trait_ref = TyBuilder::trait_ref(self.db, index_trait)
+                                .push(self_ty)
+                                .fill(|_| index_ty.clone().cast(Interner))
+                                .build();
+                            self.push_obligation(trait_ref.cast(Interner));
+                        }
+                    }
+
+                    res
                 } else {
                     self.err_ty()
                 }
diff --git a/crates/hir-ty/src/infer/mutability.rs b/crates/hir-ty/src/infer/mutability.rs
index b8a1af96fba..663ea853231 100644
--- a/crates/hir-ty/src/infer/mutability.rs
+++ b/crates/hir-ty/src/infer/mutability.rs
@@ -96,7 +96,7 @@ impl InferenceContext<'_> {
             Expr::RecordLit { path: _, fields, spread, ellipsis: _, is_assignee_expr: _ } => {
                 self.infer_mut_not_expr_iter(fields.iter().map(|it| it.expr).chain(*spread))
             }
-            &Expr::Index { base, index } => {
+            &Expr::Index { base, index, is_assignee_expr: _ } => {
                 if mutability == Mutability::Mut {
                     if let Some((f, _)) = self.result.method_resolutions.get_mut(&tgt_expr) {
                         if let Some(index_trait) = self
diff --git a/crates/hir-ty/src/mir/lower/as_place.rs b/crates/hir-ty/src/mir/lower/as_place.rs
index 8c078eb4ad7..cb5588a5c13 100644
--- a/crates/hir-ty/src/mir/lower/as_place.rs
+++ b/crates/hir-ty/src/mir/lower/as_place.rs
@@ -218,7 +218,7 @@ impl MirLowerCtx<'_> {
                 self.push_field_projection(&mut r, expr_id)?;
                 Ok(Some((r, current)))
             }
-            Expr::Index { base, index } => {
+            Expr::Index { base, index, is_assignee_expr: _ } => {
                 let base_ty = self.expr_ty_after_adjustments(*base);
                 let index_ty = self.expr_ty_after_adjustments(*index);
                 if index_ty != TyBuilder::usize()
diff --git a/crates/hir-ty/src/tests/traits.rs b/crates/hir-ty/src/tests/traits.rs
index 003ae60e8e5..d270328605a 100644
--- a/crates/hir-ty/src/tests/traits.rs
+++ b/crates/hir-ty/src/tests/traits.rs
@@ -4506,3 +4506,50 @@ fn ttt() {
 "#,
     );
 }
+
+#[test]
+fn infer_borrow() {
+    check_types(
+        r#"
+//- minicore: index
+pub struct SomeMap<K>;
+
+pub trait Borrow<Borrowed: ?Sized> {
+    fn borrow(&self) -> &Borrowed;
+}
+
+impl<T: ?Sized> Borrow<T> for T {
+    fn borrow(&self) -> &T {
+        self
+    }
+}
+
+impl<T: ?Sized> Borrow<T> for &T {
+    fn borrow(&self) -> &T {
+        &**self
+    }
+}
+
+impl<K, KB: Borrow<K>> core::ops::Index<KB> for SomeMap<K> {
+    type Output = ();
+
+    fn index(&self, _: KB) -> &() {
+        &()
+    }
+}
+
+impl<K> core::ops::IndexMut<K> for SomeMap<K> {
+    fn index_mut(&mut self, _: K) -> &mut () {
+        &mut ()
+    }
+}
+
+fn foo() {
+    let mut map = SomeMap;
+    map["a"] = ();
+    map;
+  //^^^ SomeMap<&str>
+}
+"#,
+    );
+}