Provide structured suggestion for `hashmap[idx] = val`

6 files changed, 149 insertions, 13 deletions

diff --git a/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs b/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs
index 56721cc3f5c..dd9590016b9 100644
--- a/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs
+++ b/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs
@@ -1,5 +1,8 @@
-use rustc_errors::{Applicability, Diagnostic, DiagnosticBuilder, ErrorGuaranteed};
+use rustc_errors::{
+    Applicability, Diagnostic, DiagnosticBuilder, EmissionGuarantee, ErrorGuaranteed,
+};
 use rustc_hir as hir;
+use rustc_hir::intravisit::Visitor;
 use rustc_hir::Node;
 use rustc_middle::hir::map::Map;
 use rustc_middle::mir::{Mutability, Place, PlaceRef, ProjectionElem};
@@ -614,7 +617,7 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, 'tcx> {
                             "trait `IndexMut` is required to modify indexed content, \
                                 but it is not implemented for `{ty}`",
                         ));
-                        self.suggest_map_index_mut_alternatives(ty, &mut err);
+                        self.suggest_map_index_mut_alternatives(ty, &mut err, span);
                     }
                     _ => (),
                 }
@@ -632,13 +635,120 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, 'tcx> {
         &self,
         ty: Ty<'_>,
         err: &mut DiagnosticBuilder<'_, ErrorGuaranteed>,
+        span: Span,
     ) {
         let Some(adt) = ty.ty_adt_def() else { return };
         let did = adt.did();
         if self.infcx.tcx.is_diagnostic_item(sym::HashMap, did)
             || self.infcx.tcx.is_diagnostic_item(sym::BTreeMap, did)
         {
-            err.help(format!("to modify a `{ty}`, use `.get_mut()`, `.insert()` or the entry API"));
+            struct V<'a, 'b, 'tcx, G: EmissionGuarantee> {
+                assign_span: Span,
+                err: &'a mut DiagnosticBuilder<'b, G>,
+                ty: Ty<'tcx>,
+                suggested: bool,
+            }
+            impl<'a, 'b: 'a, 'hir, 'tcx, G: EmissionGuarantee> Visitor<'hir> for V<'a, 'b, 'tcx, G> {
+                fn visit_stmt(&mut self, stmt: &'hir hir::Stmt<'hir>) {
+                    hir::intravisit::walk_stmt(self, stmt);
+                    let expr = match stmt.kind {
+                        hir::StmtKind::Semi(expr) | hir::StmtKind::Expr(expr) => expr,
+                        hir::StmtKind::Local(hir::Local { init: Some(expr), .. }) => expr,
+                        _ => {
+                            return;
+                        }
+                    };
+                    if let hir::ExprKind::Assign(place, rv, _sp) = expr.kind
+                        && let hir::ExprKind::Index(val, index) = place.kind
+                        && (expr.span == self.assign_span || place.span == self.assign_span)
+                    {
+                        // val[index] = rv;
+                        // ---------- place
+                        self.err.multipart_suggestions(
+                            &format!(
+                                "to modify a `{}`, use `.get_mut()`, `.insert()` or the entry API",
+                                self.ty,
+                            ),
+                            vec![
+                                vec![ // val.insert(index, rv);
+                                    (
+                                        val.span.shrink_to_hi().with_hi(index.span.lo()),
+                                        ".insert(".to_string(),
+                                    ),
+                                    (
+                                        index.span.shrink_to_hi().with_hi(rv.span.lo()),
+                                        ", ".to_string(),
+                                    ),
+                                    (rv.span.shrink_to_hi(), ")".to_string()),
+                                ],
+                                vec![ // val.get_mut(index).map(|v| { *v = rv; });
+                                    (
+                                        val.span.shrink_to_hi().with_hi(index.span.lo()),
+                                        ".get_mut(".to_string(),
+                                    ),
+                                    (
+                                        index.span.shrink_to_hi().with_hi(place.span.hi()),
+                                        ").map(|val| { *val".to_string(),
+                                    ),
+                                    (
+                                        rv.span.shrink_to_hi(),
+                                        "; })".to_string(),
+                                    ),
+                                ],
+                                vec![ // let x = val.entry(index).or_insert(rv);
+                                    (val.span.shrink_to_lo(), "let val = ".to_string()),
+                                    (
+                                        val.span.shrink_to_hi().with_hi(index.span.lo()),
+                                        ".entry(".to_string(),
+                                    ),
+                                    (
+                                        index.span.shrink_to_hi().with_hi(rv.span.lo()),
+                                        ").or_insert(".to_string(),
+                                    ),
+                                    (rv.span.shrink_to_hi(), ")".to_string()),
+                                ],
+                            ].into_iter(),
+                            Applicability::MachineApplicable,
+                        );
+                        self.suggested = true;
+                    } else if let hir::ExprKind::MethodCall(_path, args @ [_, ..], sp) = expr.kind
+                        && let hir::ExprKind::Index(val, index) = args[0].kind
+                        && expr.span == self.assign_span
+                    {
+                        // val[index].path(args..);
+                        self.err.multipart_suggestion(
+                            &format!("to modify a `{}` use `.get_mut()`", self.ty),
+                            vec![
+                                (
+                                    val.span.shrink_to_hi().with_hi(index.span.lo()),
+                                    ".get_mut(".to_string(),
+                                ),
+                                (
+                                    index.span.shrink_to_hi().with_hi(args[0].span.hi()),
+                                    ").map(|val| val".to_string(),
+                                ),
+                                (sp.shrink_to_hi(), ")".to_string()),
+                            ],
+                            Applicability::MachineApplicable,
+                        );
+                        self.suggested = true;
+                    }
+                }
+            }
+            let hir_map = self.infcx.tcx.hir();
+            let def_id = self.body.source.def_id();
+            let hir_id = hir_map.local_def_id_to_hir_id(def_id.as_local().unwrap());
+            let node = hir_map.find(hir_id);
+            let Some(hir::Node::Item(item)) = node else { return; };
+            let hir::ItemKind::Fn(.., body_id) = item.kind else { return; };
+            let body = self.infcx.tcx.hir().body(body_id);
+            let mut v = V { assign_span: span, err, ty, suggested: false };
+            v.visit_body(body);
+            if !v.suggested {
+                err.help(&format!(
+                    "to modify a `{ty}`, use `.get_mut()`, `.insert()` or the entry API",
+                ));
+            }
         }
     }
 
diff --git a/src/test/ui/borrowck/index-mut-help.rs b/src/test/ui/borrowck/index-mut-help.rs
index d57ef975d96..35266e113a6 100644
--- a/src/test/ui/borrowck/index-mut-help.rs
+++ b/src/test/ui/borrowck/index-mut-help.rs
@@ -1,10 +1,9 @@
 // When mutably indexing a type that implements `Index` but not `IndexMut`, a
 // special 'help' message is added to the output.
+use std::collections::HashMap;
 
 
 fn main() {
-    use std::collections::HashMap;
-
     let mut map = HashMap::new();
     map.insert("peter", "23".to_string());
 
diff --git a/src/test/ui/borrowck/index-mut-help.stderr b/src/test/ui/borrowck/index-mut-help.stderr
index 0ce60e3eb1d..f42d7e01554 100644
--- a/src/test/ui/borrowck/index-mut-help.stderr
+++ b/src/test/ui/borrowck/index-mut-help.stderr
@@ -1,23 +1,33 @@
 error[E0596]: cannot borrow data in an index of `HashMap<&str, String>` as mutable
-  --> $DIR/index-mut-help.rs:11:5
+  --> $DIR/index-mut-help.rs:10:5
    |
 LL |     map["peter"].clear();
    |     ^^^^^^^^^^^^^^^^^^^^ cannot borrow as mutable
    |
    = help: trait `IndexMut` is required to modify indexed content, but it is not implemented for `HashMap<&str, String>`
-   = help: to modify a `HashMap<&str, String>`, use `.get_mut()`, `.insert()` or the entry API
+help: to modify a `HashMap<&str, String>` use `.get_mut()`
+   |
+LL |     map.get_mut("peter").map(|val| val.clear());
+   |        ~~~~~~~~~       ~~~~~~~~~~~~~~~        +
 
 error[E0594]: cannot assign to data in an index of `HashMap<&str, String>`
-  --> $DIR/index-mut-help.rs:12:5
+  --> $DIR/index-mut-help.rs:11:5
    |
 LL |     map["peter"] = "0".to_string();
    |     ^^^^^^^^^^^^ cannot assign
    |
    = help: trait `IndexMut` is required to modify indexed content, but it is not implemented for `HashMap<&str, String>`
-   = help: to modify a `HashMap<&str, String>`, use `.get_mut()`, `.insert()` or the entry API
+help: to modify a `HashMap<&str, String>`, use `.get_mut()`, `.insert()` or the entry API
+   |
+LL |     map.insert("peter", "0".to_string());
+   |        ~~~~~~~~       ~                +
+LL |     map.get_mut("peter").map(|val| { *val = "0".to_string(); });
+   |        ~~~~~~~~~       ~~~~~~~~~~~~~~~~~~                  ++++
+LL |     let val = map.entry("peter").or_insert("0".to_string());
+   |     +++++++++    ~~~~~~~       ~~~~~~~~~~~~               +
 
 error[E0596]: cannot borrow data in an index of `HashMap<&str, String>` as mutable
-  --> $DIR/index-mut-help.rs:13:13
+  --> $DIR/index-mut-help.rs:12:13
    |
 LL |     let _ = &mut map["peter"];
    |             ^^^^^^^^^^^^^^^^^ cannot borrow as mutable
diff --git a/src/test/ui/btreemap/btreemap-index-mut.stderr b/src/test/ui/btreemap/btreemap-index-mut.stderr
index 260f7100074..26f2a4c4b29 100644
--- a/src/test/ui/btreemap/btreemap-index-mut.stderr
+++ b/src/test/ui/btreemap/btreemap-index-mut.stderr
@@ -5,7 +5,14 @@ LL |     map[&0] = 1;
    |     ^^^^^^^^^^^ cannot assign
    |
    = help: trait `IndexMut` is required to modify indexed content, but it is not implemented for `BTreeMap<u32, u32>`
-   = help: to modify a `BTreeMap<u32, u32>`, use `.get_mut()`, `.insert()` or the entry API
+help: to modify a `BTreeMap<u32, u32>`, use `.get_mut()`, `.insert()` or the entry API
+   |
+LL |     map.insert(&0, 1);
+   |        ~~~~~~~~  ~  +
+LL |     map.get_mut(&0).map(|val| { *val = 1; });
+   |        ~~~~~~~~~  ~~~~~~~~~~~~~~~~~~    ++++
+LL |     let val = map.entry(&0).or_insert(1);
+   |     +++++++++    ~~~~~~~  ~~~~~~~~~~~~ +
 
 error: aborting due to previous error
 
diff --git a/src/test/ui/hashmap/hashmap-index-mut.stderr b/src/test/ui/hashmap/hashmap-index-mut.stderr
index c72b380f466..c1948ab6271 100644
--- a/src/test/ui/hashmap/hashmap-index-mut.stderr
+++ b/src/test/ui/hashmap/hashmap-index-mut.stderr
@@ -5,7 +5,14 @@ LL |     map[&0] = 1;
    |     ^^^^^^^^^^^ cannot assign
    |
    = help: trait `IndexMut` is required to modify indexed content, but it is not implemented for `HashMap<u32, u32>`
-   = help: to modify a `HashMap<u32, u32>`, use `.get_mut()`, `.insert()` or the entry API
+help: to modify a `HashMap<u32, u32>`, use `.get_mut()`, `.insert()` or the entry API
+   |
+LL |     map.insert(&0, 1);
+   |        ~~~~~~~~  ~  +
+LL |     map.get_mut(&0).map(|val| { *val = 1; });
+   |        ~~~~~~~~~  ~~~~~~~~~~~~~~~~~~    ++++
+LL |     let val = map.entry(&0).or_insert(1);
+   |     +++++++++    ~~~~~~~  ~~~~~~~~~~~~ +
 
 error: aborting due to previous error
 
diff --git a/src/test/ui/issues/issue-41726.stderr b/src/test/ui/issues/issue-41726.stderr
index 9c70ab7d971..b05c1fb14ef 100644
--- a/src/test/ui/issues/issue-41726.stderr
+++ b/src/test/ui/issues/issue-41726.stderr
@@ -5,7 +5,10 @@ LL |         things[src.as_str()].sort();
    |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^ cannot borrow as mutable
    |
    = help: trait `IndexMut` is required to modify indexed content, but it is not implemented for `HashMap<String, Vec<String>>`
-   = help: to modify a `HashMap<String, Vec<String>>`, use `.get_mut()`, `.insert()` or the entry API
+help: to modify a `HashMap<String, Vec<String>>` use `.get_mut()`
+   |
+LL |         things.get_mut(src.as_str()).map(|val| val.sort());
+   |               ~~~~~~~~~            ~~~~~~~~~~~~~~~       +
 
 error: aborting due to previous error