Fix misattributed type inference error span for index expressions

compiler/rustc_hir_typeck/src/cast.rs

@@ -725,9 +725,21 @@ impl<'a, 'tcx> CastCheck<'tcx> {
         );
     }
 
+    fn expr_span_for_type_resolution(&self, fcx: &FnCtxt<'a, 'tcx>) -> Span {
+        if let hir::ExprKind::Index(_, idx, _) = self.expr.kind
+            && fcx.resolve_vars_if_possible(self.expr_ty).is_ty_var()
+            && fcx.resolve_vars_if_possible(fcx.node_ty(idx.hir_id)).is_ty_var()
+        {
+            index_operand_ambiguity_span(idx)
+        } else {
+            self.expr_span
+        }
+    }
+
     #[instrument(skip(fcx), level = "debug")]
     pub(crate) fn check(mut self, fcx: &FnCtxt<'a, 'tcx>) {
-        self.expr_ty = fcx.structurally_resolve_type(self.expr_span, self.expr_ty);
+        let expr_span = self.expr_span_for_type_resolution(fcx);
+        self.expr_ty = fcx.structurally_resolve_type(expr_span, self.expr_ty);
         self.cast_ty = fcx.structurally_resolve_type(self.cast_span, self.cast_ty);
 
         debug!("check_cast({}, {:?} as {:?})", self.expr.hir_id, self.expr_ty, self.cast_ty);
@@ -1208,3 +1220,10 @@ impl<'a, 'tcx> CastCheck<'tcx> {
         }
     }
 }
+
+fn index_operand_ambiguity_span(expr: &hir::Expr<'_>) -> Span {
+    match expr.kind {
+        hir::ExprKind::MethodCall(segment, ..) => segment.ident.span,
+        _ => expr.span,
+    }
+}

compiler/rustc_hir_typeck/src/fn_ctxt/adjust_fulfillment_errors.rs

@@ -4,8 +4,11 @@ use rustc_hir as hir;
 use rustc_hir::def::{DefKind, Res};
 use rustc_hir::def_id::DefId;
 use rustc_infer::traits::ObligationCauseCode;
-use rustc_middle::ty::{self, Ty, TyCtxt, TypeSuperVisitable, TypeVisitable, TypeVisitor};
+use rustc_middle::ty::{
+    self, Ty, TyCtxt, TypeSuperVisitable, TypeVisitable, TypeVisitableExt, TypeVisitor,
+};
 use rustc_span::{Span, kw};
+use rustc_trait_selection::infer::InferCtxtExt;
 use rustc_trait_selection::traits;
 
 use crate::FnCtxt;
@@ -37,6 +40,10 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         &self,
         error: &mut traits::FulfillmentError<'tcx>,
     ) -> bool {
+        if self.adjust_binop_index_operand(error) {
+            return true;
+        }
+
         let (def_id, hir_id, idx, flavor) = match *error.obligation.cause.code().peel_derives() {
             ObligationCauseCode::WhereClauseInExpr(def_id, _, hir_id, idx) => {
                 (def_id, hir_id, idx, ClauseFlavor::Where)
@@ -165,6 +172,119 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         }
     }
 
+    fn adjust_binop_index_operand(&self, error: &mut traits::FulfillmentError<'tcx>) -> bool {
+        let ObligationCauseCode::BinOp { lhs_hir_id, rhs_hir_id, .. } =
+            *error.obligation.cause.code().peel_derives()
+        else {
+            return false;
+        };
+        if !matches!(error.code, traits::FulfillmentErrorCode::Ambiguity { .. })
+            || !error.obligation.predicate.has_infer()
+        {
+            return false;
+        }
+
+        let hir::Node::Expr(lhs_expr) = self.tcx.hir_node(lhs_hir_id) else {
+            return false;
+        };
+        let hir::Node::Expr(rhs_expr) = self.tcx.hir_node(rhs_hir_id) else {
+            return false;
+        };
+        let Some(binop) = self.binop_for_operands(lhs_hir_id, rhs_hir_id) else {
+            return false;
+        };
+        let hir::ExprKind::Index(indexed_expr, idx, _) = rhs_expr.kind else {
+            return false;
+        };
+        if !self.resolve_vars_if_possible(self.node_ty(idx.hir_id)).is_ty_var() {
+            return false;
+        }
+        let lhs_ty = self.resolve_vars_if_possible(self.node_ty(lhs_expr.hir_id));
+        let indexed_ty = self.resolve_vars_if_possible(self.node_ty(indexed_expr.hir_id));
+        let rhs_ty = match *indexed_ty.kind() {
+            ty::Array(element_ty, _) | ty::Slice(element_ty) => element_ty,
+            ty::Ref(_, pointee_ty, _) => match *pointee_ty.kind() {
+                ty::Array(element_ty, _) | ty::Slice(element_ty) => element_ty,
+                _ => self.resolve_vars_if_possible(self.node_ty(rhs_expr.hir_id)),
+            },
+            _ => self.resolve_vars_if_possible(self.node_ty(rhs_expr.hir_id)),
+        };
+        if !self.binop_accepts_types(binop.node, lhs_ty, rhs_ty) {
+            return false;
+        }
+
+        error.obligation.cause.span = match idx.kind {
+            hir::ExprKind::MethodCall(segment, ..) => segment.ident.span,
+            _ => idx.span,
+        };
+        true
+    }
+
+    fn binop_for_operands(
+        &self,
+        lhs_hir_id: hir::HirId,
+        rhs_hir_id: hir::HirId,
+    ) -> Option<hir::BinOp> {
+        let hir::Node::Expr(parent_expr) = self.tcx.parent_hir_node(rhs_hir_id) else {
+            return None;
+        };
+        let hir::ExprKind::Binary(binop, lhs_expr, rhs_expr) = parent_expr.kind else {
+            return None;
+        };
+        (lhs_expr.hir_id == lhs_hir_id && rhs_expr.hir_id == rhs_hir_id).then_some(binop)
+    }
+
+    fn binop_accepts_types(
+        &self,
+        binop: hir::BinOpKind,
+        lhs_ty: Ty<'tcx>,
+        rhs_ty: Ty<'tcx>,
+    ) -> bool {
+        let lhs_ty = self.deref_ty_if_possible(lhs_ty);
+        let rhs_ty = self.deref_ty_if_possible(rhs_ty);
+        if lhs_ty.references_error() || rhs_ty.references_error() {
+            return true;
+        }
+
+        match binop {
+            hir::BinOpKind::Shl | hir::BinOpKind::Shr => {
+                lhs_ty.is_integral() && rhs_ty.is_integral()
+            }
+            hir::BinOpKind::Add
+            | hir::BinOpKind::Sub
+            | hir::BinOpKind::Mul
+            | hir::BinOpKind::Div
+            | hir::BinOpKind::Rem => {
+                self.can_eq(self.param_env, lhs_ty, rhs_ty)
+                    && (lhs_ty.is_integral() || lhs_ty.is_floating_point())
+                    && (rhs_ty.is_integral() || rhs_ty.is_floating_point())
+            }
+            hir::BinOpKind::BitXor | hir::BinOpKind::BitAnd | hir::BinOpKind::BitOr => {
+                self.can_eq(self.param_env, lhs_ty, rhs_ty)
+                    && ((lhs_ty.is_integral() && rhs_ty.is_integral())
+                        || (lhs_ty.is_bool() && rhs_ty.is_bool()))
+            }
+            hir::BinOpKind::Eq
+            | hir::BinOpKind::Ne
+            | hir::BinOpKind::Lt
+            | hir::BinOpKind::Le
+            | hir::BinOpKind::Ge
+            | hir::BinOpKind::Gt => {
+                self.can_eq(self.param_env, lhs_ty, rhs_ty)
+                    && lhs_ty.is_scalar()
+                    && rhs_ty.is_scalar()
+            }
+            hir::BinOpKind::And | hir::BinOpKind::Or => lhs_ty.is_bool() && rhs_ty.is_bool(),
+        }
+    }
+
+    fn deref_ty_if_possible(&self, ty: Ty<'tcx>) -> Ty<'tcx> {
+        match ty.kind() {
+            ty::Ref(_, ty, hir::Mutability::Not) => *ty,
+            _ => ty,
+        }
+    }
+
     fn point_at_expr_if_possible(
         &self,
         error: &mut traits::FulfillmentError<'tcx>,

compiler/rustc_hir_typeck/src/op.rs

@@ -313,6 +313,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 }
 
                 self.write_method_call_and_enforce_effects(expr.hir_id, expr.span, method);
+
                 method.sig.output()
             }
             // error types are considered "builtin"

tests/ui/type-inference/index-expr-ambiguous-type.rs

@@ -0,0 +1,54 @@
+// Regression test for #156738
+//
+// When the index type in `arr[idx]` is ambiguous, the error should point
+// at the index sub-expression, not the whole indexing expression or
+// surrounding operators.
+
+fn with_cast() {
+    let bad_idx = 0u8;
+    let _foo = [1, 2, 3][bad_idx.into()] as i32;
+    //~^ ERROR type annotations needed
+}
+
+fn with_binop() {
+    let bad_idx = 0u8;
+    let _foo = 0 + [1, 2, 3][bad_idx.into()];
+    //~^ ERROR type annotations needed
+}
+
+fn standalone() {
+    let bad_idx = 0u8;
+    let _foo = [1, 2, 3][bad_idx.into()];
+    //~^ ERROR type annotations needed
+}
+
+fn with_known_index_type() {
+    let bad_idx = 0u8;
+    let _foo = [1, 2, 3][Into::<usize>::into(bad_idx)] as i32;
+}
+
+fn invalid_operator_with_ambiguous_index() {
+    let bad_idx = 0u8;
+    let _foo = true + [1, 2, 3][bad_idx.into()];
+    //~^ ERROR cannot add
+}
+
+fn mismatched_numeric_binop_with_ambiguous_index() {
+    let bad_idx = 0u8;
+    let _foo = 0u64 + [1i32, 2, 3][bad_idx.into()];
+    //~^ ERROR type annotations needed
+}
+
+fn shift_with_ambiguous_index() {
+    let bad_idx = 0u8;
+    let _foo = 1u32 << [0u8][bad_idx.into()];
+    //~^ ERROR type annotations needed
+}
+
+fn string_add_with_ambiguous_index() {
+    let bad_idx = 0u8;
+    let _foo = String::new() + [""][bad_idx.into()];
+    //~^ ERROR type annotations needed
+}
+
+fn main() {}

tests/ui/type-inference/index-expr-ambiguous-type.stderr

@@ -0,0 +1,70 @@
+error[E0282]: type annotations needed
+  --> $DIR/index-expr-ambiguous-type.rs:9:34
+   |
+LL |     let _foo = [1, 2, 3][bad_idx.into()] as i32;
+   |                                  ^^^^ cannot infer type
+
+error[E0284]: type annotations needed
+  --> $DIR/index-expr-ambiguous-type.rs:15:38
+   |
+LL |     let _foo = 0 + [1, 2, 3][bad_idx.into()];
+   |                                      ^^^^ cannot infer type
+   |
+   = note: cannot satisfy `<i32 as Add<_>>::Output == _`
+
+error[E0283]: type annotations needed
+  --> $DIR/index-expr-ambiguous-type.rs:21:34
+   |
+LL |     let _foo = [1, 2, 3][bad_idx.into()];
+   |                                  ^^^^
+   |
+   = note: the type must implement `From<u8>`
+   = note: required for `u8` to implement `Into<_>`
+help: try using a fully qualified path to specify the expected types
+   |
+LL -     let _foo = [1, 2, 3][bad_idx.into()];
+LL +     let _foo = [1, 2, 3][<u8 as Into<T>>::into(bad_idx)];
+   |
+
+error[E0369]: cannot add `_` to `bool`
+  --> $DIR/index-expr-ambiguous-type.rs:32:21
+   |
+LL |     let _foo = true + [1, 2, 3][bad_idx.into()];
+   |                ---- ^ ------------------------- _
+   |                |
+   |                bool
+
+error[E0284]: type annotations needed
+  --> $DIR/index-expr-ambiguous-type.rs:38:21
+   |
+LL |     let _foo = 0u64 + [1i32, 2, 3][bad_idx.into()];
+   |                     ^ cannot infer type
+   |
+   = note: cannot satisfy `<u64 as Add<_>>::Output == _`
+
+error[E0284]: type annotations needed
+  --> $DIR/index-expr-ambiguous-type.rs:44:38
+   |
+LL |     let _foo = 1u32 << [0u8][bad_idx.into()];
+   |                                      ^^^^ cannot infer type
+   |
+   = note: cannot satisfy `<u32 as Shl<_>>::Output == _`
+
+error[E0283]: type annotations needed
+  --> $DIR/index-expr-ambiguous-type.rs:50:45
+   |
+LL |     let _foo = String::new() + [""][bad_idx.into()];
+   |                                             ^^^^
+   |
+   = note: the type must implement `From<u8>`
+   = note: required for `u8` to implement `Into<_>`
+help: try using a fully qualified path to specify the expected types
+   |
+LL -     let _foo = String::new() + [""][bad_idx.into()];
+LL +     let _foo = String::new() + [""][<u8 as Into<T>>::into(bad_idx)];
+   |
+
+error: aborting due to 7 previous errors
+
+Some errors have detailed explanations: E0282, E0283, E0284, E0369.
+For more information about an error, try `rustc --explain E0282`.