Auto merge of #17755 - ShoyuVanilla:issue-17738, r=davidbarsky

fix: Apply `IndexMut` obligations for non-assigning mutable index usages

Fixes #17738

Currently, we are pushing `IndexMut` obligations only for assign usages;
https://github.com/rust-lang/rust-analyzer/blob/f982f3fa2c23570c10108e83c1ecc392ea411866/crates/hir-ty/src/infer/expr.rs#L809-L817

Author: bors

src/tools/rust-analyzer/crates/hir-ty/src/infer/mutability.rs

@@ -1,15 +1,18 @@
 //! Finds if an expression is an immutable context or a mutable context, which is used in selecting
 //! between `Deref` and `DerefMut` or `Index` and `IndexMut` or similar.
 
-use chalk_ir::Mutability;
+use chalk_ir::{cast::Cast, Mutability};
 use hir_def::{
     hir::{Array, BinaryOp, BindingAnnotation, Expr, ExprId, PatId, Statement, UnaryOp},
     lang_item::LangItem,
 };
 use hir_expand::name::Name;
 use intern::sym;
 
-use crate::{lower::lower_to_chalk_mutability, Adjust, Adjustment, AutoBorrow, OverloadedDeref};
+use crate::{
+    infer::Expectation, lower::lower_to_chalk_mutability, Adjust, Adjustment, AutoBorrow, Interner,
+    OverloadedDeref, TyBuilder, TyKind,
+};
 
 use super::InferenceContext;
 
@@ -101,7 +104,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, is_assignee_expr: _ } => {
+            &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
@@ -115,18 +118,40 @@ impl InferenceContext<'_> {
                                 .method_by_name(&Name::new_symbol_root(sym::index_mut.clone()))
                             {
                                 *f = index_fn;
+                                let mut base_ty = None;
                                 let base_adjustments = self
                                     .result
                                     .expr_adjustments
                                     .get_mut(&base)
                                     .and_then(|it| it.last_mut());
                                 if let Some(Adjustment {
                                     kind: Adjust::Borrow(AutoBorrow::Ref(mutability)),
-                                    ..
+                                    target,
                                 }) = base_adjustments
                                 {
+                                    // For assignee exprs `IndexMut` obiligations are already applied
+                                    if !is_assignee_expr {
+                                        if let TyKind::Ref(_, _, ty) = target.kind(Interner) {
+                                            base_ty = Some(ty.clone());
+                                        }
+                                    }
                                     *mutability = Mutability::Mut;
                                 }
+
+                                // Apply `IndexMut` obligation for non-assignee expr
+                                if let Some(base_ty) = base_ty {
+                                    let index_ty =
+                                        if let Some(ty) = self.result.type_of_expr.get(index) {
+                                            ty.clone()
+                                        } else {
+                                            self.infer_expr(index, &Expectation::none())
+                                        };
+                                    let trait_ref = TyBuilder::trait_ref(self.db, index_trait)
+                                        .push(base_ty)
+                                        .fill(|_| index_ty.clone().cast(Interner))
+                                        .build();
+                                    self.push_obligation(trait_ref.cast(Interner));
+                                }
                             }
                         }
                     }

src/tools/rust-analyzer/crates/hir-ty/src/tests/regression.rs

@@ -2075,3 +2075,50 @@ impl<'a, T> Trait<'a> for &'a T {
 "#,
     )
 }
+
+#[test]
+fn issue_17738() {
+    check_types(
+        r#"
+//- minicore: index
+use core::ops::{Index, IndexMut};
+
+struct Foo<K, V>(K, V);
+
+struct Bar;
+
+impl Bar {
+    fn bar(&mut self) {}
+}
+
+impl<K, V> Foo<K, V> {
+    fn new(_v: V) -> Self {
+        loop {}
+    }
+}
+
+impl<K, B, V> Index<B> for Foo<K, V> {
+    type Output = V;
+    fn index(&self, _index: B) -> &Self::Output {
+        loop {}
+    }
+}
+
+impl<K, V> IndexMut<K> for Foo<K, V> {
+    fn index_mut(&mut self, _index: K) -> &mut Self::Output {
+        loop {}
+    }
+}
+
+fn test() {
+    let mut t1 = Foo::new(Bar);
+     // ^^^^^^ Foo<&'? (), Bar>
+    t1[&()] = Bar;
+
+    let mut t2 = Foo::new(Bar);
+     // ^^^^^^ Foo<&'? (), Bar>
+    t2[&()].bar();
+}
+"#,
+    )
+}