Suggest use .get_mut instead of &mut when overloaded index type not impl IndexMut

Signed-off-by: xizheyin <xizheyin@smail.nju.edu.cn>

Author: xizheyin

compiler/rustc_borrowck/src/diagnostics/mod.rs

@@ -889,7 +889,7 @@ pub(super) enum BorrowedContentSource<'tcx> {
     DerefMutableRef,
     DerefSharedRef,
     OverloadedDeref(Ty<'tcx>),
-    OverloadedIndex(Ty<'tcx>),
+    OverloadedIndex(Ty<'tcx>, Ty<'tcx>),
 }
 
 impl<'tcx> BorrowedContentSource<'tcx> {
@@ -905,7 +905,7 @@ impl<'tcx> BorrowedContentSource<'tcx> {
                     _ => None,
                 })
                 .unwrap_or_else(|| format!("dereference of `{ty}`")),
-            BorrowedContentSource::OverloadedIndex(ty) => format!("index of `{ty}`"),
+            BorrowedContentSource::OverloadedIndex(ty, _) => format!("index of `{ty}`"),
         }
     }
 
@@ -917,7 +917,7 @@ impl<'tcx> BorrowedContentSource<'tcx> {
             // Overloaded deref and index operators should be evaluated into a
             // temporary. So we don't need a description here.
             BorrowedContentSource::OverloadedDeref(_)
-            | BorrowedContentSource::OverloadedIndex(_) => None,
+            | BorrowedContentSource::OverloadedIndex(_, _) => None,
         }
     }
 
@@ -935,7 +935,7 @@ impl<'tcx> BorrowedContentSource<'tcx> {
                     _ => None,
                 })
                 .unwrap_or_else(|| format!("dereference of `{ty}`")),
-            BorrowedContentSource::OverloadedIndex(ty) => format!("an index of `{ty}`"),
+            BorrowedContentSource::OverloadedIndex(ty, _) => format!("an index of `{ty}`"),
         }
     }
 
@@ -951,7 +951,7 @@ impl<'tcx> BorrowedContentSource<'tcx> {
                 } else if tcx.is_lang_item(trait_id, LangItem::Index)
                     || tcx.is_lang_item(trait_id, LangItem::IndexMut)
                 {
-                    Some(BorrowedContentSource::OverloadedIndex(args.type_at(0)))
+                    Some(BorrowedContentSource::OverloadedIndex(args.type_at(0), args.type_at(1)))
                 } else {
                     None
                 }

compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs

@@ -519,7 +519,7 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                              but it is not implemented for `{ty}`",
                         ));
                     }
-                    Some(BorrowedContentSource::OverloadedIndex(ty)) => {
+                    Some(BorrowedContentSource::OverloadedIndex(ty, _)) => {
                         err.help(format!(
                             "trait `IndexMut` is required to modify indexed content, \
                              but it is not implemented for `{ty}`",
@@ -1176,6 +1176,7 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                             self.body.source_info(location).span
                         }
                     });
+
                 match opt_assignment_rhs_span.and_then(|s| s.desugaring_kind()) {
                     // on for loops, RHS points to the iterator part
                     Some(DesugaringKind::ForLoop) => {
@@ -1196,7 +1197,13 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                         None
                     }
                     None => {
-                        if name != kw::SelfLower {
+                        let suggest_get_mut = self.suggest_get_mut_when_not_impl_index_mut(
+                            local,
+                            opt_assignment_rhs_span,
+                        );
+                        if suggest_get_mut.is_some() {
+                            suggest_get_mut
+                        } else if name != kw::SelfLower {
                             suggest_ampmut(
                                 self.infcx.tcx,
                                 local_decl.ty,
@@ -1414,6 +1421,66 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
             None => {}
         }
     }
+
+    /// check if the RHS is an overloaded index expression from hashmap or btreemap,
+    /// if so, suggest using .get_mut() instead of &mut, see issue #143732
+    /// For example
+    /// ```text
+    /// let mut map = HashMap::new();
+    /// let value = &map["key"];
+    /// ```
+    fn suggest_get_mut_when_not_impl_index_mut(
+        &self,
+        local: Local,
+        opt_assignment_rhs_span: Option<Span>,
+    ) -> Option<AmpMutSugg> {
+        self.find_assignments(local)
+            .first()
+            .map(|&location| {
+                if let Some(mir::Statement {
+                    source_info: _,
+                    kind: mir::StatementKind::Assign(box (_, mir::Rvalue::Ref(_, _, place))),
+                    ..
+                }) = self.body[location.block].statements.get(location.statement_index)
+                    && let BorrowedContentSource::OverloadedIndex(ty, index_ty) =
+                        self.borrowed_content_source(place.as_ref())
+                    && let Some(index_mut_trait) = self.infcx.tcx.lang_items().index_mut_trait()
+                    && !self
+                        .infcx
+                        .type_implements_trait(
+                            index_mut_trait,
+                            [ty, index_ty],
+                            self.infcx.param_env,
+                        )
+                        .must_apply_modulo_regions()
+                {
+                    if let Some(rhs_span) = opt_assignment_rhs_span
+                        && let Ok(rhs_str) =
+                            self.infcx.tcx.sess.source_map().span_to_snippet(rhs_span)
+                        && let Some(content) = rhs_str.strip_prefix('&')
+                        && content.contains('[')
+                        && content.contains(']')
+                    {
+                        let bracket_start = content.find('[')?;
+                        let bracket_end = content.rfind(']')?;
+
+                        if bracket_start < bracket_end {
+                            let map_part = &content[..bracket_start];
+                            let key_part = &content[bracket_start + 1..bracket_end];
+
+                            return Some(AmpMutSugg {
+                                has_sugg: true,
+                                span: rhs_span,
+                                suggestion: format!("{}.get_mut({}).unwrap()", map_part, key_part),
+                                additional: None,
+                            });
+                        }
+                    }
+                }
+                None
+            })
+            .flatten()
+    }
 }
 
 struct BindingFinder {

tests/ui/borrowck/suggest-use-as-mut-for-map.stderr

@@ -6,8 +6,9 @@ LL |     string.push_str("test");
    |
 help: consider changing this to be a mutable reference
    |
-LL |     let string = &mut map[&0];
-   |                   +++
+LL -     let string = &map[&0];
+LL +     let string = map.get_mut(&0).unwrap();
+   |
 
 error[E0596]: cannot borrow `*string` as mutable, as it is behind a `&` reference
   --> $DIR/suggest-use-as-mut-for-map.rs:14:5
@@ -17,8 +18,9 @@ LL |     string.push_str("test");
    |
 help: consider changing this to be a mutable reference
    |
-LL |     let string = &mut map[&0];
-   |                   +++
+LL -     let string = &map[&0];
+LL +     let string = map.get_mut(&0).unwrap();
+   |
 
 error[E0596]: cannot borrow `*string` as mutable, as it is behind a `&` reference
   --> $DIR/suggest-use-as-mut-for-map.rs:18:5