Rollup merge of #152216 - GTimothy:map-diagnostics-fix, r=JonathanBrouwer

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

fixes #150001
fixes rust-lang/rust-analyzer#16076
fixes #134917

The issues are threefold and linked:
 1. Assigning data to a ~~collection~~BTreeMap/HashMap suggests 3 solutions all marked as `MachineApplicable`
 2. The suggestions are slightly wrong with regards to their borrowing needs.
 3. The suggestions are not guaranteed to produce code that is valid, and suggestion number two is equivalent to an update, not an insertion.

This PR:
 - splits the large triple suggestion into three
 - sets them to `MaybeIncorrect`
 - automatically determines the required borrowing to use.

 I think this solution may not be very elegant, expecially the key typechecking / borrowing part, but it works. I am however very open to any improvement/change :)

 edit: edited to replace 'collection' with  BTreeMap/HashMap'

Author: JonathanBrouwer

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,166 @@ 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 (borrowed_prefix, borrowed_index);
+
+                        // only suggest `insert` and `entry` if index is of type K or &{n}K or *{n}K (when there is a Borrow impl for this case).
+                        // 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() {
+                            let (index_refs, key_refs) =
+                                (count_ty_refs(index_ty), count_ty_refs(key_ty));
+
+                            let (deref_prefix, deref_index) = if index_refs >= key_refs {
+                                // index is &{n}K
+                                strip_n_refs(index, index_refs - key_refs)
+                                    .map(|val| ("".to_string(), val))
+                                    .unwrap_or_else(|(depth, val)| {
+                                        (
+                                            if key_refs == 0 {
+                                                "*".repeat(
+                                                    (index_refs-key_refs).checked_sub(depth).expect("return depth from strip_n_refs should be smaller than the input")
+                                                )
+                                            } else {
+                                                String::new() //if key K is a ref, autoderef finish this for us.
+                                            },
+                                            val,
+                                        )
+                                    })
+                            } else {
+                                // in this case the minimal ref addition works for all subcases
+                                ("&".repeat(key_refs - index_refs), index)
+                            };
+
+                            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) = (
+                                String::new(),
+                                if index_refs > key_refs {
+                                    strip_n_refs(index, index_refs - key_refs - 1)
+                                        .unwrap_or_else(|(_depth, val)| val)
+                                    // even if we tried to strip more, we can stop there thanks to autoderef
+                                } else {
+                                    // when the diff is negative or zero, we already are in the index=&Q case.
+                                    index
+                                },
+                            );
+                        } 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 +877,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