Suggest use .get_mut instead of &mut when source type is HashMap or BTreeMap

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}`",
@@ -1196,7 +1196,11 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                         None
                     }
                     None => {
-                        if name != kw::SelfLower {
+                        let (should_suggest_ref_mut, suggest_get_mut) =
+                            self.suggest_get_mut_or_ref_mut(local, opt_assignment_rhs_span);
+                        if !should_suggest_ref_mut {
+                            suggest_get_mut
+                        } else if name != kw::SelfLower {
                             suggest_ampmut(
                                 self.infcx.tcx,
                                 local_decl.ty,
@@ -1414,6 +1418,79 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
             None => {}
         }
     }
+
+    /// check if the RHS is an overloaded index expression whose source type
+    /// doesn't implement `IndexMut`. This category contains two sub-categories:
+    /// 1. HashMap or BTreeMap
+    /// 2. Other types
+    /// For the first sub-category, we suggest using `.get_mut()` instead of `&mut`,
+    /// For the second sub-category, we still suggest use `&mut`.
+    /// See issue #143732.
+    ///
+    /// the first element of return value is a boolean indicating
+    /// if we should suggest use `&mut`
+    fn suggest_get_mut_or_ref_mut(
+        &self,
+        local: Local,
+        opt_assignment_rhs_span: Option<Span>,
+    ) -> (bool, 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(adt) = ty.ty_adt_def()
+                        && (self.infcx.tcx.is_diagnostic_item(sym::HashMap, adt.did())
+                            || self.infcx.tcx.is_diagnostic_item(sym::BTreeMap, adt.did()))
+                        && 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 Some(bracket_start) = content.find('[')
+                        && let Some(bracket_end) = content.rfind(']')
+                        && bracket_start < bracket_end
+                    {
+                        let map_part = &content[..bracket_start];
+                        let key_part = &content[bracket_start + 1..bracket_end];
+
+                        (
+                            false,
+                            Some(AmpMutSugg {
+                                has_sugg: true,
+                                span: rhs_span,
+                                suggestion: format!("{}.get_mut({}).unwrap()", map_part, key_part),
+                                additional: None,
+                            }),
+                        )
+                    } else {
+                        // Type not implemented `IndexMut`,
+                        // and we did not suggest because of HashMap or BTreeMap
+                        (false, None)
+                    }
+                } else {
+                    // Type Impl IndexMut, we should suggest use `&mut`
+                    (true, None)
+                }
+            })
+            .unwrap_or((true, None))
+    }
 }
 
 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