Fix 'assign to data in an index of' collection suggestions

compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs

@@ -668,7 +668,9 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                 err: &'a mut Diag<'infcx>,
                 ty: Ty<'tcx>,
                 suggested: bool,
+                infcx: &'a rustc_infer::infer::InferCtxt<'tcx>,
             }
+
             impl<'a, 'infcx, 'tcx> Visitor<'tcx> for SuggestIndexOperatorAlternativeVisitor<'a, 'infcx, 'tcx> {
                 fn visit_stmt(&mut self, stmt: &'tcx hir::Stmt<'tcx>) {
                     hir::intravisit::walk_stmt(self, stmt);
@@ -679,23 +681,76 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                             return;
                         }
                     };
+
+                    // Because of TypeChecking and indexing, we know: index is &Q
+                    // with K: Eq + Hash + Borrow<Q>,
+                    // with Q: Eq + Hash + ?Sized,
+                    //
+                    //
+                    //
+                    // There is no other constraint on the types, therefore we need to look at the
+                    // constraints of the suggestions:
+                    // pub fn get_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
+                    // where
+                    //     K: Borrow<Q>,
+                    //     Q: Hash + Eq + ?Sized,
+                    //
+                    // pub fn insert(&mut self, k: K, v: V) -> Option<V>
+                    // pub fn entry(&mut self, key: K) -> Entry<'_, K, V>
+                    //
+                    // But let's note that there could be also, if imported from hashbrown:
+                    // pub fn entry_ref<'a, 'b, Q>(
+                    //     &'a mut self,
+                    //     key: &'b Q,
+                    // ) -> EntryRef<'a, 'b, K, Q, V, S, A>
+                    // where
+                    //     Q: Hash + Equivalent<K> + ?Sized,
+
+                    /// Testing if index AND key are &Q. We cannot simply test equality because the
+                    /// lifetimes differ.
+                    fn index_and_key_are_same_borrowed_type<'tcx>(
+                        index: Ty<'tcx>,
+                        key: Ty<'tcx>,
+                    ) -> bool {
+                        if let (ty::Ref(_, index_inner_ty, _), ty::Ref(_, key_inner_ty, _)) =
+                            (index.kind(), key.kind())
+                        {
+                            *index_inner_ty == *key_inner_ty
+                        } else {
+                            false
+                        }
+                    }
+                    /// Testing if index is &K:
+                    fn index_is_borrowed_key<'tcx>(index: Ty<'tcx>, key: Ty<'tcx>) -> bool {
+                        if let ty::Ref(_, inner_ty, _) = index.kind() {
+                            *inner_ty == key
+                        } else {
+                            false
+                        }
+                    }
+
+                    // we know ty is a map, with a key type at walk distance 2.
+                    let key_type = self.ty.walk().nth(1).unwrap().expect_ty();
+
                     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!(
-                                "use `.insert()` to insert a value into a `{}`, `.get_mut()` \
-                                to modify it, or the entry API for more flexibility",
-                                self.ty,
-                            ),
-                            vec![
+                        let index_ty =
+                            self.infcx.tcx.typeck(val.hir_id.owner.def_id).expr_ty(index);
+                        // only suggest `insert` and `entry` if index is of type K or &K.
+                        let index_is_borrowed_key = index_is_borrowed_key(index_ty, key_type);
+                        if index_is_borrowed_key
+                            || index_and_key_are_same_borrowed_type(index_ty, key_type)
+                        {
+                            let offset = BytePos(index_is_borrowed_key as u32);
+                            self.err.multipart_suggestion(
+                                format!("use `.insert()` to insert a value into a `{}`", self.ty),
                                 vec![
                                     // val.insert(index, rv);
                                     (
-                                        val.span.shrink_to_hi().with_hi(index.span.lo()),
+                                        val.span.shrink_to_hi().with_hi(index.span.lo() + offset), //remove the
+                                        //leading &
                                         ".insert(".to_string(),
                                     ),
                                     (
@@ -704,35 +759,55 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                                     ),
                                     (rv.span.shrink_to_hi(), ")".to_string()),
                                 ],
+                                Applicability::MaybeIncorrect,
+                            );
+                            self.err.multipart_suggestion(
+                                format!(
+                                    "use the entry API to modify a `{}` for more flexibility",
+                                    self.ty
+                                ),
                                 vec![
-                                    // if let Some(v) = val.get_mut(index) { *v = rv; }
-                                    (val.span.shrink_to_lo(), "if let Some(val) = ".to_string()),
-                                    (
-                                        val.span.shrink_to_hi().with_hi(index.span.lo()),
-                                        ".get_mut(".to_string(),
-                                    ),
-                                    (
-                                        index.span.shrink_to_hi().with_hi(place.span.hi()),
-                                        ") { *val".to_string(),
-                                    ),
-                                    (rv.span.shrink_to_hi(), "; }".to_string()),
-                                ],
-                                vec![
-                                    // let x = val.entry(index).or_insert(rv);
+                                    // let x = val.entry(index).insert_entry(rv);
                                     (val.span.shrink_to_lo(), "let val = ".to_string()),
                                     (
-                                        val.span.shrink_to_hi().with_hi(index.span.lo()),
+                                        val.span.shrink_to_hi().with_hi(index.span.lo() + offset), //remove the
+                                        //leading &
                                         ".entry(".to_string(),
                                     ),
                                     (
                                         index.span.shrink_to_hi().with_hi(rv.span.lo()),
-                                        ").or_insert(".to_string(),
+                                        ").insert_entry(".to_string(),
                                     ),
                                     (rv.span.shrink_to_hi(), ")".to_string()),
                                 ],
+                                Applicability::MaybeIncorrect,
+                            );
+                        }
+                        // in all cases, suggest get_mut because K:Borrow<K> or Q:Borrow<K> as a
+                        // requirement of indexing.
+                        self.err.multipart_suggestion(
+                            format!(
+                                "use `.get_mut()` to modify an existing key in a `{}`",
+                                self.ty,
+                            ),
+                            vec![
+                                // if let Some(v) = val.get_mut(index) { *v = rv; }
+                                (val.span.shrink_to_lo(), "if let Some(val) = ".to_string()),
+                                (
+                                    val.span.shrink_to_hi().with_hi(index.span.lo()),
+                                    ".get_mut(".to_string(),
+                                ),
+                                (
+                                    index.span.shrink_to_hi().with_hi(place.span.hi()),
+                                    ") { *val".to_string(),
+                                ),
+                                (rv.span.shrink_to_hi(), "; }".to_string()),
                             ],
-                            Applicability::MachineApplicable,
+                            Applicability::MaybeIncorrect,
                         );
+                        //FIXME: in the future, include the ref_entry suggestion here if it is added
+                        //to std.
+
                         self.suggested = true;
                     } else if let hir::ExprKind::MethodCall(_path, receiver, _, sp) = expr.kind
                         && let hir::ExprKind::Index(val, index, _) = receiver.kind
@@ -768,6 +843,7 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                 err,
                 ty,
                 suggested: false,
+                infcx: self.infcx,
             };
             v.visit_body(&body);
             if !v.suggested {

tests/ui/borrowck/index-mut-help.stderr

@@ -18,16 +18,20 @@ 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: use `.insert()` to insert a value into a `HashMap<&str, String>`, `.get_mut()` to modify it, or the entry API for more flexibility
+help: use `.insert()` to insert a value into a `HashMap<&str, String>`
    |
 LL -     map["peter"] = "0".to_string();
 LL +     map.insert("peter", "0".to_string());
    |
+help: use the entry API to modify a `HashMap<&str, String>` for more flexibility
+   |
 LL -     map["peter"] = "0".to_string();
-LL +     if let Some(val) = map.get_mut("peter") { *val = "0".to_string(); };
+LL +     let val = map.entry("peter").insert_entry("0".to_string());
+   |
+help: use `.get_mut()` to modify an existing key in a `HashMap<&str, String>`
    |
 LL -     map["peter"] = "0".to_string();
-LL +     let val = map.entry("peter").or_insert("0".to_string());
+LL +     if let Some(val) = map.get_mut("peter") { *val = "0".to_string(); };
    |
 
 error[E0596]: cannot borrow data in an index of `HashMap<&str, String>` as mutable

tests/ui/borrowck/index-mut-suggest-borrow-issue.rs

@@ -0,0 +1,84 @@
+// When mutably indexing a type that implements `Index` but not `IndexMut`, a
+// special 'help' message is added to the output.
+//
+// The suggestions have different requirements on the index type.
+// map[&idx] = val;
+// If idx is later used, it results in a borrow  issue if both:
+// - the the suggestion that was chosen requires the idx, not a ref to it.
+// - idx's  type is not Copy.
+//
+// For now, we suggest all options, regardless of the Copy-ness of the idx.
+use std::collections::HashMap;
+
+/// With a copy type, subsequent reuse of idx is not an issue.
+fn copy_type() {
+    // ===&str===, a copy type.
+
+    let mut map = HashMap::<&str, u32>::new();
+    // earlier, peter is initialised with 22.
+    // map["peter"] = 22; //ERROR
+    map.insert("peter", 22);
+
+    // at some point, if we get a &str variable peter again
+    let peter = "peter";
+    // and we want to use it to update the map but still want to use it later?
+    // map[&peter] = 23; // ERROR
+    // we could insert again, and because &T are copy, we can use peter even if we use peter later.
+    map.insert(peter, 23); // WORKS
+    println!("my name is {peter}"); // WORKS because &str is Copy
+    // and we could use a &&str too in this case, because &str:Borrow<&str> (because T:Borrow<T>)
+    if let Some(val) = map.get_mut(&peter) {
+        *val = 23;
+    }; // WORKS
+    println!("my name is {peter}"); // WORKS because &str is Copy
+    // even a &str directly, (because rust auto borrows peter -> &peter ?)
+    if let Some(val) = map.get_mut(peter) {
+        *val = 24;
+    }; // WORKS
+}
+
+/// With a non-copy type, subsequent reuse of idx is an issue for `insert` and `entry`.
+fn non_copy_type_insert() {
+    // ===STRING===, a non-copy type
+
+    let mut map = HashMap::<String, u32>::new();
+    // earlier, peter is initialised with 22.
+    // map[&"peter".to_string()] = 22; // ERROR cannot assign
+    map.insert("peter".to_string(), 22);
+
+    // at some point, if we get a String variable peter again
+    let peter = "peter".to_string();
+    // and we want to use it to update the map but still want to use it later?
+    // map[&peter] = 23; // ERROR cannot assign
+    // we could insert again, but we cannot use peter after.
+    map.insert(peter, 23); // WORKS
+    println!("my name is {peter}"); //~ ERROR: borrow of moved value: `peter` [E0382]
+}
+
+/// With a non-copy type, subsequent reuse of idx is not an issue for `get_mut`.
+fn non_copy_type_get_mut() {
+    // ===STRING===, a non-copy type
+
+    let mut map = HashMap::<String, u32>::new();
+    // earlier, peter is initialised with 22.
+    // map["peter".to_string()] = 22; // ERROR cannot assign
+    map.insert("peter".to_string(), 22);
+
+    // at some point, if we get a String variable peter again
+    let peter = "peter".to_string();
+    // and we want to use it to update the map but still want to use it later?
+    // map[&peter] = 23; // ERROR cannot assign
+    // we can use a &String in this case, so get_mut is always fine.
+    if let Some(val) = map.get_mut(&peter) {
+        *val = 23;
+    }; // WORKS
+    println!("my name is {peter}"); // WORKS
+    // or a &str because String:Borrow<str>) and "peter" is &str.
+
+    if let Some(val) = map.get_mut("peter") {
+        *val = 24;
+    }; // WORKS
+    println!("my name is {peter}"); // WORKS
+}
+
+fn main() {}

tests/ui/borrowck/index-mut-suggest-borrow-issue.stderr

@@ -0,0 +1,19 @@
+error[E0382]: borrow of moved value: `peter`
+  --> $DIR/index-mut-suggest-borrow-issue.rs:55:27
+   |
+LL |     let peter = "peter".to_string();
+   |         ----- move occurs because `peter` has type `String`, which does not implement the `Copy` trait
+...
+LL |     map.insert(peter, 23); // WORKS
+   |                ----- value moved here
+LL |     println!("my name is {peter}");
+   |                           ^^^^^ value borrowed here after move
+   |
+help: consider cloning the value if the performance cost is acceptable
+   |
+LL |     map.insert(peter.clone(), 23); // WORKS
+   |                     ++++++++
+
+error: aborting due to 1 previous error
+
+For more information about this error, try `rustc --explain E0382`.

tests/ui/collections/btreemap/btreemap-index-mut-2.stderr

@@ -5,16 +5,20 @@ 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: use `.insert()` to insert a value into a `BTreeMap<u32, u32>`, `.get_mut()` to modify it, or the entry API for more flexibility
+help: use `.insert()` to insert a value into a `BTreeMap<u32, u32>`
    |
 LL -         map[&0] = 1;
-LL +         map.insert(&0, 1);
+LL +         map.insert(0, 1);
+   |
+help: use the entry API to modify a `BTreeMap<u32, u32>` for more flexibility
    |
 LL -         map[&0] = 1;
-LL +         if let Some(val) = map.get_mut(&0) { *val = 1; };
+LL +         let val = map.entry(0).insert_entry(1);
+   |
+help: use `.get_mut()` to modify an existing key in a `BTreeMap<u32, u32>`
    |
 LL -         map[&0] = 1;
-LL +         let val = map.entry(&0).or_insert(1);
+LL +         if let Some(val) = map.get_mut(&0) { *val = 1; };
    |
 
 error: aborting due to 1 previous error

tests/ui/collections/btreemap/btreemap-index-mut.stderr

@@ -5,16 +5,20 @@ 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: use `.insert()` to insert a value into a `BTreeMap<u32, u32>`, `.get_mut()` to modify it, or the entry API for more flexibility
+help: use `.insert()` to insert a value into a `BTreeMap<u32, u32>`
    |
 LL -     map[&0] = 1;
-LL +     map.insert(&0, 1);
+LL +     map.insert(0, 1);
+   |
+help: use the entry API to modify a `BTreeMap<u32, u32>` for more flexibility
    |
 LL -     map[&0] = 1;
-LL +     if let Some(val) = map.get_mut(&0) { *val = 1; };
+LL +     let val = map.entry(0).insert_entry(1);
+   |
+help: use `.get_mut()` to modify an existing key in a `BTreeMap<u32, u32>`
    |
 LL -     map[&0] = 1;
-LL +     let val = map.entry(&0).or_insert(1);
+LL +     if let Some(val) = map.get_mut(&0) { *val = 1; };
    |
 
 error: aborting due to 1 previous error

tests/ui/collections/hashmap/hashmap-index-mut.stderr

@@ -5,16 +5,20 @@ 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: use `.insert()` to insert a value into a `HashMap<u32, u32>`, `.get_mut()` to modify it, or the entry API for more flexibility
+help: use `.insert()` to insert a value into a `HashMap<u32, u32>`
    |
 LL -     map[&0] = 1;
-LL +     map.insert(&0, 1);
+LL +     map.insert(0, 1);
+   |
+help: use the entry API to modify a `HashMap<u32, u32>` for more flexibility
    |
 LL -     map[&0] = 1;
-LL +     if let Some(val) = map.get_mut(&0) { *val = 1; };
+LL +     let val = map.entry(0).insert_entry(1);
+   |
+help: use `.get_mut()` to modify an existing key in a `HashMap<u32, u32>`
    |
 LL -     map[&0] = 1;
-LL +     let val = map.entry(&0).or_insert(1);
+LL +     if let Some(val) = map.get_mut(&0) { *val = 1; };
    |
 
 error: aborting due to 1 previous error