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

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

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}`",
@@ -1158,24 +1158,28 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                 opt_ty_info,
                 ..
             })) => {
-                // check if the RHS is from desugaring
-                let opt_assignment_rhs_span =
-                    self.find_assignments(local).first().map(|&___location| {
+                // check if the RHS is from desugaring, and return rhs of assignment
+                let (opt_assignment_rhs, opt_assignment_rhs_span) = self
+                    .find_assignments(local)
+                    .first()
+                    .map(|&___location| {
+                        let mut rhs = None;
+                        let mut span = self.body.source_info(___location).span;
                         if let Some(mir::Statement {
                             source_info: _,
-                            kind:
-                                mir::StatementKind::Assign(box (
-                                    _,
-                                    mir::Rvalue::Use(mir::Operand::Copy(place)),
-                                )),
+                            kind: mir::StatementKind::Assign(box (_, rvalue)),
                             ..
                         }) = self.body[___location.block].statements.get(___location.statement_index)
                         {
-                            self.body.local_decls[place.local].source_info.span
-                        } else {
-                            self.body.source_info(___location).span
+                            // If there is a assignment, we should return the rhs
+                            rhs = Some(rvalue);
+                            if let mir::Rvalue::Use(mir::Operand::Copy(place)) = rvalue {
+                                span = self.body.local_decls[place.local].source_info.span;
+                            }
                         }
-                    });
+                        (rhs, Some(span))
+                    })
+                    .unwrap_or((None, None));
                 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 +1200,13 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                         None
                     }
                     None => {
-                        if name != kw::SelfLower {
+                        let result = self.suggest_get_mut_or_ref_mut(
+                            opt_assignment_rhs,
+                            opt_assignment_rhs_span,
+                        );
+                        if let ControlFlow::Break(suggest_get_mut) = result {
+                            suggest_get_mut
+                        } else if name != kw::SelfLower {
                             suggest_ampmut(
                                 self.infcx.tcx,
                                 local_decl.ty,
@@ -1414,6 +1424,64 @@ 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.
+    ///
+    /// Break indicates we should not suggest use `&mut`
+    /// Continue indicates we continue to try suggesting use `&mut`
+    fn suggest_get_mut_or_ref_mut(
+        &self,
+        opt_assignment_rhs: Option<&mir::Rvalue<'tcx>>,
+        opt_assignment_rhs_span: Option<Span>,
+    ) -> ControlFlow<Option<AmpMutSugg>, ()> {
+        if let Some(mir::Rvalue::Ref(_, _, place)) = opt_assignment_rhs
+            && 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];
+
+                // We only suggest use `.get_mut()` for HashMap or BTreeMap
+                // We don't try suggesting use `&mut`, so break
+                ControlFlow::Break(Some(AmpMutSugg {
+                    has_sugg: true,
+                    span: rhs_span,
+                    suggestion: format!("{}.get_mut({}).unwrap()", map_part, key_part),
+                    additional: None,
+                }))
+            } else {
+                // Type not implemented `IndexMut` but not HashMap or BTreeMap or error
+                // and we did not suggest either `.get_mut()` or `&mut`, so break
+                ControlFlow::Break(None)
+            }
+        } else {
+            // Type Impl IndexMut, we should continue to try suggesting use `&mut`
+            ControlFlow::Continue(())
+        }
+    }
 }
 
 struct BindingFinder {

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

@@ -0,0 +1,22 @@
+// We don't suggest change `&` to `&mut` with HashMap and BTreeMap
+// instead we suggest using .get_mut() instead of &mut, see issue #143732
+
+//@ run-rustfix
+
+fn main() {
+    let mut map = std::collections::BTreeMap::new();
+    map.insert(0, "string".to_owned());
+
+    let string = map.get_mut(&0).unwrap();
+    string.push_str("test"); //~ ERROR cannot borrow `*string` as mutable, as it is behind a `&` reference [E0596]
+
+    let mut map = std::collections::HashMap::new();
+    map.insert(0, "string".to_owned());
+
+    let string = map.get_mut(&0).unwrap();
+    string.push_str("test"); //~ ERROR cannot borrow `*string` as mutable, as it is behind a `&` reference [E0596]
+
+    let mut vec = vec![String::new(), String::new()];
+    let string = &mut vec[0];
+    string.push_str("test"); //~ ERROR cannot borrow `*string` as mutable, as it is behind a `&` reference [E0596]
+}

tests/ui/borrowck/suggest-use-as-mut-for-map.rs renamed to tests/ui/borrowck/suggest-use-as-mut-for-map-1.rs

@@ -1,24 +1,7 @@
-// We don't suggest change `&` to `&mut`
+// We don't suggest change `&` to `&mut` with HashMap and BTreeMap
 // instead we suggest using .get_mut() instead of &mut, see issue #143732
 
-// Custom type that implements Index<usize> but not IndexMut
-struct ReadOnlyVec<T> {
-    data: Vec<T>,
-}
-
-impl<T> ReadOnlyVec<T> {
-    fn new(data: Vec<T>) -> Self {
-        ReadOnlyVec { data }
-    }
-}
-
-impl<T> std::ops::Index<usize> for ReadOnlyVec<T> {
-    type Output = T;
-    
-    fn index(&self, index: usize) -> &Self::Output {
-        &self.data[index]
-    }
-}
+//@ run-rustfix
 
 fn main() {
     let mut map = std::collections::BTreeMap::new();
@@ -36,9 +19,4 @@ fn main() {
     let mut vec = vec![String::new(), String::new()];
     let string = &vec[0];
     string.push_str("test"); //~ ERROR cannot borrow `*string` as mutable, as it is behind a `&` reference [E0596]
-
-    // Example with our custom ReadOnlyVec type
-    let read_only_vec = ReadOnlyVec::new(vec![String::new(), String::new()]);
-    let string_ref = &read_only_vec[0];
-    string_ref.push_str("test"); //~ ERROR cannot borrow `*string_ref` as mutable, as it is behind a `&` reference [E0596]
 }

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

@@ -1,27 +1,29 @@
 error[E0596]: cannot borrow `*string` as mutable, as it is behind a `&` reference
-  --> $DIR/suggest-use-as-mut-for-map.rs:28:5
+  --> $DIR/suggest-use-as-mut-for-map-1.rs:11:5
    |
 LL |     string.push_str("test");
    |     ^^^^^^ `string` is a `&` reference, so the data it refers to cannot be borrowed as mutable
    |
 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:34:5
+  --> $DIR/suggest-use-as-mut-for-map-1.rs:17:5
    |
 LL |     string.push_str("test");
    |     ^^^^^^ `string` is a `&` reference, so the data it refers to cannot be borrowed as mutable
    |
 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:38:5
+  --> $DIR/suggest-use-as-mut-for-map-1.rs:21:5
    |
 LL |     string.push_str("test");
    |     ^^^^^^ `string` is a `&` reference, so the data it refers to cannot be borrowed as mutable
@@ -31,17 +33,6 @@ help: consider changing this to be a mutable reference
 LL |     let string = &mut vec[0];
    |                   +++
 
-error[E0596]: cannot borrow `*string_ref` as mutable, as it is behind a `&` reference
-  --> $DIR/suggest-use-as-mut-for-map.rs:43:5
-   |
-LL |     string_ref.push_str("test");
-   |     ^^^^^^^^^^ `string_ref` is a `&` reference, so the data it refers to cannot be borrowed as mutable
-   |
-help: consider changing this to be a mutable reference
-   |
-LL |     let string_ref = &mut read_only_vec[0];
-   |                       +++
-
-error: aborting due to 4 previous errors
+error: aborting due to 3 previous errors
 
 For more information about this error, try `rustc --explain E0596`.

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

@@ -0,0 +1,27 @@
+// We don't suggest anything, see issue #143732
+
+// Custom type that implements Index<usize> but not IndexMut
+struct ReadOnlyVec<T> {
+    data: Vec<T>,
+}
+
+impl<T> ReadOnlyVec<T> {
+    fn new(data: Vec<T>) -> Self {
+        ReadOnlyVec { data }
+    }
+}
+
+impl<T> std::ops::Index<usize> for ReadOnlyVec<T> {
+    type Output = T;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        &self.data[index]
+    }
+}
+
+fn main() {
+    // Example with our custom ReadOnlyVec type
+    let read_only_vec = ReadOnlyVec::new(vec![String::new(), String::new()]);
+    let string_ref = &read_only_vec[0];
+    string_ref.push_str("test"); //~ ERROR cannot borrow `*string_ref` as mutable, as it is behind a `&` reference [E0596]
+}

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

@@ -0,0 +1,9 @@
+error[E0596]: cannot borrow `*string_ref` as mutable, as it is behind a `&` reference
+  --> $DIR/suggest-use-as-mut-for-map-2.rs:26:5
+   |
+LL |     string_ref.push_str("test");
+   |     ^^^^^^^^^^ `string_ref` is a `&` reference, so the data it refers to cannot be borrowed as mutable
+
+error: aborting due to 1 previous error
+
+For more information about this error, try `rustc --explain E0596`.