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}`",
@@ -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-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`.