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

compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs

@@ -6,7 +6,7 @@ use rustc_abi::FieldIdx;
 use rustc_errors::{Applicability, Diag};
 use rustc_hir::def_id::DefId;
 use rustc_hir::intravisit::Visitor;
-use rustc_hir::{self as hir, BindingMode, ByRef, Node};
+use rustc_hir::{self as hir, BindingMode, ByRef, Expr, Node};
 use rustc_middle::bug;
 use rustc_middle::hir::place::PlaceBase;
 use rustc_middle::mir::visit::PlaceContext;
@@ -669,7 +669,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);
@@ -680,60 +682,160 @@ 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,
+                    //
+                    // which fulfill the requirements of `get_mut`. If Q=K or Q=&{n}K, the requirements
+                    // of `entry` and `insert` are fulfilled too after dereferencing. If K is not
+                    // copy, a subsequent `clone` call may be needed.
+
+                    /// Taken straight from https://doc.rust-lang.org/nightly/nightly-rustc/clippy_utils/fn.peel_hir_ty_refs.html
+                    /// Adapted to mid using https://doc.rust-lang.org/nightly/nightly-rustc/rustc_middle/ty/struct.Ty.html#method.peel_refs
+                    /// Simplified to counting only
+                    /// Peels off all references on the type. Returns the number of references
+                    /// removed.
+                    fn count_ty_refs<'tcx>(mut ty: Ty<'tcx>) -> usize {
+                        let mut count = 0;
+                        while let ty::Ref(_, inner_ty, _) = ty.kind() {
+                            ty = *inner_ty;
+                            count += 1;
+                        }
+                        count
+                    }
+
+                    /// Try to strip `n` `&` reference from an expression.
+                    /// If the expression does not have enough leading `&`, return an Error
+                    /// containing a count of the successfully stripped ones and the stripped
+                    /// expression.
+                    fn strip_n_refs<'a, 'b>(
+                        mut expr: &'a Expr<'b>,
+                        n: usize,
+                    ) -> Result<&'a Expr<'b>, (usize, &'a Expr<'b>)> {
+                        for count in 0..n {
+                            match expr {
+                                Expr {
+                                    kind: ExprKind::AddrOf(hir::BorrowKind::Ref, _, inner),
+                                    ..
+                                } => expr = inner,
+                                _ => return Err((count, expr)),
+                            }
+                        }
+                        Ok(expr)
+                    }
+
+                    // we know ty is a map, with a key type at walk distance 2.
+                    let key_ty = 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);
+
+                        let count_refs_diff: isize =
+                            count_ty_refs(index_ty) as isize - count_ty_refs(key_ty) as isize;
+
+                        let (borrowed_prefix, borrowed_index);
+
+                        // only suggest `insert` and `entry` if index is of type K or &{n}K.
+                        // We use `peel_refs` because borrow lifetimes may differ in both index and
+                        // key. I.e, if they are of the same base type:
+                        if index_ty.peel_refs() == key_ty.peel_refs() {
+                            assert!(
+                                count_refs_diff >= 0,
+                                "compiler bug I think, please report this"
+                            );
+                            // if base type is same, borrowed depth must be exact.
+                            let (deref_prefix, deref_index) =
+                                strip_n_refs(index, count_refs_diff as usize)
+                                    .map(|val| ("".to_string(), val))
+                                    .unwrap_or_else(|(depth, val)| {
+                                        (
+                                            "*".repeat(
+                                                usize::try_from(count_refs_diff)
+                                                    .expect("passed assert")
+                                                    - depth,
+                                            ),
+                                            val,
+                                        )
+                                    });
+
+                            self.err.multipart_suggestion(
+                                format!("use `.insert()` to insert a value into a `{}`", self.ty),
                                 vec![
-                                    // val.insert(index, rv);
+                                    // val.insert({deref_prefix}{deref_index}, rv);
                                     (
-                                        val.span.shrink_to_hi().with_hi(index.span.lo()),
-                                        ".insert(".to_string(),
+                                        val.span.shrink_to_hi().with_hi(deref_index.span.lo()),
+                                        format!(".insert({deref_prefix}"),
                                     ),
                                     (
-                                        index.span.shrink_to_hi().with_hi(rv.span.lo()),
+                                        deref_index.span.shrink_to_hi().with_hi(rv.span.lo()),
                                         ", ".to_string(),
                                     ),
                                     (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({deref_prefix}{deref_index}).insert_entry(rv);
                                     (val.span.shrink_to_lo(), "let val = ".to_string()),
                                     (
-                                        val.span.shrink_to_hi().with_hi(index.span.lo()),
-                                        ".entry(".to_string(),
+                                        val.span.shrink_to_hi().with_hi(deref_index.span.lo()),
+                                        format!(".entry({deref_prefix}"),
                                     ),
                                     (
-                                        index.span.shrink_to_hi().with_hi(rv.span.lo()),
-                                        ").or_insert(".to_string(),
+                                        deref_index.span.shrink_to_hi().with_hi(rv.span.lo()),
+                                        ").insert_entry(".to_string(),
                                     ),
                                     (rv.span.shrink_to_hi(), ")".to_string()),
                                 ],
+                                Applicability::MaybeIncorrect,
+                            );
+
+                            // we can make the next suggestions nicer by stripping as many leading `&` as
+                            // we can, autoderef will do the rest
+                            (borrowed_prefix, borrowed_index) =
+                                match strip_n_refs(index, 0.max(count_refs_diff - 1) as usize) {
+                                    Ok(val) => (String::new(), val),
+                                    // even if we tried to strip more, we can stop there thanks to
+                                    // autoderef
+                                    Err((_depth, val)) => (String::new(), val),
+                                };
+                        } else {
+                            (borrowed_prefix, borrowed_index) = (String::new(), index)
+                        }
+                        // 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({borrowed_prefix}{borrowed_index}) { *v = rv; }
+                                (val.span.shrink_to_lo(), "if let Some(val) = ".to_string()),
+                                (
+                                    val.span.shrink_to_hi().with_hi(borrowed_index.span.lo()),
+                                    format!(".get_mut({borrowed_prefix}"),
+                                ),
+                                (
+                                    borrowed_index.span.shrink_to_hi().with_hi(place.span.hi()),
+                                    ") { *val".to_string(),
+                                ),
+                                (rv.span.shrink_to_hi(), "; }".to_string()),
                             ],
-                            Applicability::MachineApplicable,
+                            Applicability::MaybeIncorrect,
                         );
+
                         self.suggested = true;
                     } else if let hir::ExprKind::MethodCall(_path, receiver, _, sp) = expr.kind
                         && let hir::ExprKind::Index(val, index, _) = receiver.kind
@@ -769,6 +871,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-help2.rs

@@ -0,0 +1,153 @@
+// When mutably indexing a type that implements `Index` but not `IndexMut`, a
+// special 'help' message is added to the output.
+// ... Except when it is not!
+use std::borrow::Borrow;
+use std::collections::HashMap;
+
+#[derive(Hash, Eq, PartialEq)]
+struct A {}
+#[derive(Hash, Eq, PartialEq, Copy, Clone)]
+struct ACopy {}
+#[derive(Hash, Eq, PartialEq)]
+struct B {}
+#[derive(Hash, Eq, PartialEq)]
+struct C {}
+#[derive(Hash, Eq, PartialEq)]
+struct D {}
+
+impl Borrow<C> for D {
+    fn borrow(&self) -> &C {
+        &C {}
+    }
+}
+
+/// In this test the key is a A type.
+fn test_a_index() {
+    let mut map = HashMap::<A, u32>::new();
+
+    let index = &&&&&A {};
+    // index gets autodereferenced
+    map[index] = 23; //~ ERROR E0594
+    map[&&&&&&&&index] = 23; //~ ERROR E0594
+}
+/// In this test the key is a A type. Auto-dereferencing to &A works but dereferencing to A need
+/// the exact amount of `*`
+fn test_a_2() {
+    let mut map = HashMap::<A, u32>::new();
+
+    let index = &&&&&A {};
+    // complete dereferencing needs to be exact
+    map.insert(*index, 23); //~ ERROR E0308
+    // index gets autodereferenced
+    if let Some(val) = map.get_mut(index) {
+        *val = 23;
+    } // passes
+}
+/// In this test the key is a A type. Could not be merged with 2 because compiler only shows error
+/// 0308
+fn test_a_3() {
+    let mut map = HashMap::<A, u32>::new();
+
+    let index = &&&&&A {};
+    // A does not implement Copy so a Clone might be required
+    map.insert(*****index, 23); //~ ERROR E0507
+}
+
+/// In this test the key is a ACopy type.
+fn test_acopy_index() {
+    let mut map = HashMap::<ACopy, u32>::new();
+
+    let index = &&&&&ACopy {};
+    // index gets autodereferenced
+    map[index] = 23; //~ ERROR E0594
+}
+/// In this test the key is a ACopy type. Auto-dereferencing to &A works but dereferencing to A
+/// need the exact amount of `*`.
+fn test_acopy_2() {
+    let mut map = HashMap::<ACopy, u32>::new();
+
+    let index = &&&&&ACopy {};
+    // complete dereferencing needs to be exact
+    map.insert(*index, 23); //~ ERROR E0308
+
+    // index gets autodereferenced
+    if let Some(val) = map.get_mut(index) {
+        *val = 23;
+    } // passes
+}
+/// In this test the key is a ACopy type. Could not be merged with 2 because compiler only shows
+/// error 0308
+fn test_acopy_3() {
+    let mut map = HashMap::<ACopy, u32>::new();
+
+    let index = &&&&&ACopy {};
+    map.insert(*****index, 23); // no E057 error in this case because ACopy is Copy
+}
+
+/// In this test the key type is B-reference. The autodereferencing does not work in this case for
+/// both for the map[index] part and `get_mut` call.
+/// This leads to E0277 errors.
+fn test_b() {
+    let mut map = HashMap::<&B, u32>::new();
+
+    let index = &&&&&B {};
+    // index does NOT get autorederefenced
+    map[index] = 23; //~ ERROR E0277
+
+    // index does NOT get autorederefenced
+    if let Some(val) = map.get_mut(index) { //~ ERROR E0277
+        *val = 23;
+    }
+    if let Some(val) = map.get_mut(***index) {
+        *val = 23;
+    } // passes
+}
+
+
+/// In this test the key type is C.
+/// The `Borrow<C> for D` implementation changes nothing here, same error as for test_a.
+fn test_c() {
+    let mut map = HashMap::<C, u32>::new();
+
+    let index = &&&&&C {};
+    // index gets autodereferenced
+    map[index] = 23; //~ ERROR E0594
+
+    // index gets autodereferenced
+    if let Some(val) = map.get_mut(index) {
+        *val = 23;
+    } // passes
+}
+
+/// In this test the key type is D. The `Borrow<C> for D` implementation seems to prevent
+/// autodereferencing.
+/// The autodereferencing does not work in this case for both for the map[index] part and `get_mut`
+/// call.
+/// This leads to E0277 errors.
+fn test_d() {
+    let mut map = HashMap::<D, u32>::new();
+
+    let index = &&&&&D {};
+    // index does NOT get autorederefenced
+    map[index] = 23; //~ ERROR E0277
+
+    // index does NOT get autorederefenced
+    if let Some(val) = map.get_mut(index) {//~ ERROR E0277
+        *val = 23;
+    }
+    if let Some(val) = map.get_mut(****index) {
+        *val = 23;
+    } // passes
+}
+
+fn main() {
+    test_a_index();
+    test_a_2();
+    test_a_3();
+    test_acopy_index();
+    test_acopy_2();
+    test_acopy_3();
+    test_b();
+    test_c();
+    test_d();
+}