Suggest use .get_mut instead of &mut when occur overloaded index expression

Signed-off-by: xizheyin <[email protected]>

Author: xizheyin

compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs

@@ -1176,6 +1176,22 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                             self.body.source_info(location).span
                         }
                     });
+
+                // check if the RHS is an overloaded index expression, if so, suggest using .get_mut() instead of &mut, see issue #143732
+                let rhs_content_source = 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)
+                    {
+                        Some(self.borrowed_content_source(place.as_ref()))
+                    }else{
+                        None
+                    }
+                }).flatten();
+
                 match opt_assignment_rhs_span.and_then(|s| s.desugaring_kind()) {
                     // on for loops, RHS points to the iterator part
                     Some(DesugaringKind::ForLoop) => {
@@ -1203,6 +1219,9 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                                 decl_span,
                                 opt_assignment_rhs_span,
                                 opt_ty_info,
+                                rhs_content_source.is_some_and(|s| {
+                                    matches!(s, BorrowedContentSource::OverloadedIndex(_))
+                                }),
                             )
                         } else {
                             match local_decl.local_info() {
@@ -1226,6 +1245,7 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                                     decl_span,
                                     opt_assignment_rhs_span,
                                     opt_ty_info,
+                                    false,
                                 ),
                             }
                         }
@@ -1507,6 +1527,7 @@ fn suggest_ampmut<'tcx>(
     decl_span: Span,
     opt_assignment_rhs_span: Option<Span>,
     opt_ty_info: Option<Span>,
+    is_overloaded_index: bool,
 ) -> Option<AmpMutSugg> {
     // if there is a RHS and it starts with a `&` from it, then check if it is
     // mutable, and if not, put suggest putting `mut ` to make it mutable.
@@ -1553,6 +1574,28 @@ fn suggest_ampmut<'tcx>(
         // if the reference is already mutable then there is nothing we can do
         // here.
         if !is_mut {
+            // If this is an overloaded index expression, suggest using .get_mut() instead of &mut
+            if is_overloaded_index {
+                // Try to extract the expression from &expr[key] to suggest expr.get_mut(key).unwrap()
+                let content = rhs_str.strip_prefix('&')?;
+                if 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,
+                        });
+                    }
+                }
+            }
+
             // shrink the span to just after the `&` in `&variable`
             let span = rhs_span.with_lo(rhs_span.lo() + BytePos(1)).shrink_to_lo();
 

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: aborting due to 1 previous error