fix: check_expr_if to point to a more accurate location of the compilation error in some cases

Author: sgasho

compiler/rustc_hir_typeck/src/_if.rs

@@ -0,0 +1,320 @@
+use rustc_errors::Diag;
+use rustc_hir::{self as hir, HirId};
+use rustc_infer::traits;
+use rustc_middle::ty::{Ty, TypeVisitableExt};
+use rustc_span::{ErrorGuaranteed, Span};
+use smallvec::SmallVec;
+
+use crate::coercion::{CoerceMany, DynamicCoerceMany};
+use crate::{Diverges, Expectation, FnCtxt, bug};
+
+#[derive(Clone, Debug)]
+struct BranchBody<'tcx> {
+    expr: &'tcx hir::Expr<'tcx>,
+    ty: Ty<'tcx>,
+    diverges: Diverges,
+    span: Span,
+}
+
+#[derive(Clone, Debug)]
+struct IfBranch<'tcx> {
+    if_expr: &'tcx hir::Expr<'tcx>,
+    cond_diverges: Diverges,
+    body: BranchBody<'tcx>,
+}
+
+#[derive(Default, Debug)]
+struct IfChain<'tcx> {
+    branches: SmallVec<[IfBranch<'tcx>; 4]>,
+    cond_error: Option<ErrorGuaranteed>,
+}
+
+enum IfChainTail<'tcx> {
+    FinalElse(BranchBody<'tcx>),
+    Missing(&'tcx hir::Expr<'tcx>),
+}
+
+const RECENT_BRANCH_HISTORY_LIMIT: usize = 5;
+
+#[derive(Default)]
+struct RecentBranchTypeHistory<'tcx> {
+    entries: SmallVec<[(Ty<'tcx>, Span); 4]>,
+}
+
+impl<'tcx> RecentBranchTypeHistory<'tcx> {
+    fn diagnostic_snapshot(&self) -> SmallVec<[(Ty<'tcx>, Span); 4]> {
+        self.entries.clone()
+    }
+
+    fn record(&mut self, ty: Ty<'tcx>, span: Span) {
+        if ty.is_never() {
+            return;
+        }
+
+        self.entries.push((ty, span));
+        if self.entries.len() > RECENT_BRANCH_HISTORY_LIMIT {
+            self.entries.remove(0);
+        }
+    }
+}
+
+impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
+    pub(crate) fn check_expr_if(
+        &self,
+        expr_id: HirId,
+        sp: Span,
+        orig_expected: Expectation<'tcx>,
+    ) -> Ty<'tcx> {
+        let root_if_expr = self.tcx.hir_expect_expr(expr_id);
+        let expected = orig_expected.try_structurally_resolve_and_adjust_for_branches(self, sp);
+
+        let initial_diverges = self.diverges.get();
+
+        let (chain, tail) = self.collect_if_chain(root_if_expr, expected);
+
+        let terminal_else_expr = match &tail {
+            IfChainTail::FinalElse(else_branch) => Some(else_branch.expr),
+            IfChainTail::Missing(last_if_expr) => Some(*last_if_expr),
+        };
+
+        let coerce_to_ty = expected.coercion_target_type(self, sp);
+        let mut coerce: DynamicCoerceMany<'_> = CoerceMany::new(coerce_to_ty);
+
+        let tail_defines_return_position_impl_trait =
+            self.return_position_impl_trait_from_match_expectation(orig_expected);
+        let mut recent_branch_types = RecentBranchTypeHistory::default();
+
+        for (idx, branch) in chain.branches.iter().enumerate() {
+            if idx > 0 {
+                let merged_ty = coerce.merged_ty();
+                self.ensure_if_branch_type(branch.if_expr.hir_id, merged_ty);
+            }
+
+            let branch_body = &branch.body;
+            let skip_coercion = matches!(tail, IfChainTail::Missing(_))
+                && chain.branches.len() > 1
+                && idx == chain.branches.len() - 1;
+            if skip_coercion {
+                recent_branch_types.record(branch_body.ty, branch_body.span);
+                continue;
+            }
+            let next_else_expr =
+                chain.branches.get(idx + 1).map(|next| next.if_expr).or(terminal_else_expr);
+            let mut branch_cause = if let Some(next_else_expr) = next_else_expr {
+                let prev_branch =
+                    if idx > 0 { chain.branches.get(idx - 1).unwrap_or(branch) } else { branch };
+                self.if_cause(
+                    prev_branch.if_expr.hir_id,
+                    next_else_expr,
+                    tail_defines_return_position_impl_trait,
+                )
+            } else {
+                self.misc(branch_body.span)
+            };
+            let diag_info = recent_branch_types.diagnostic_snapshot();
+            let cause_span =
+                if idx == 0 { Some(root_if_expr.span) } else { Some(branch_body.span) };
+
+            self.coerce_if_arm(
+                &mut coerce,
+                &mut branch_cause,
+                branch_body.expr,
+                branch_body.ty,
+                cause_span,
+                branch_body.span,
+                diag_info,
+            );
+
+            recent_branch_types.record(branch_body.ty, branch_body.span);
+        }
+
+        let mut tail_diverges = match tail {
+            IfChainTail::FinalElse(else_branch) => {
+                let mut else_cause = self.if_cause(
+                    expr_id,
+                    else_branch.expr,
+                    tail_defines_return_position_impl_trait,
+                );
+                let diag_info = recent_branch_types.diagnostic_snapshot();
+                self.coerce_if_arm(
+                    &mut coerce,
+                    &mut else_cause,
+                    else_branch.expr,
+                    else_branch.ty,
+                    None,
+                    else_branch.span,
+                    diag_info,
+                );
+                recent_branch_types.record(else_branch.ty, else_branch.span);
+
+                else_branch.diverges
+            }
+            IfChainTail::Missing(last_if_expr) => {
+                let hir::ExprKind::If(tail_cond, tail_then, _) = last_if_expr.kind else {
+                    bug!("expected `if` expression, found {:#?}", last_if_expr);
+                };
+                self.if_fallback_coercion(last_if_expr.span, tail_cond, tail_then, &mut coerce);
+                Diverges::Maybe
+            }
+        };
+
+        for branch in chain.branches.iter().rev() {
+            tail_diverges = branch.cond_diverges | (branch.body.diverges & tail_diverges);
+        }
+        self.diverges.set(initial_diverges | tail_diverges);
+
+        let result_ty = coerce.complete(self);
+
+        let final_ty = if let Some(guar) = chain.cond_error {
+            Ty::new_error(self.tcx, guar)
+        } else {
+            result_ty
+        };
+
+        for branch in chain.branches.iter().skip(1) {
+            self.overwrite_if_branch_type(branch.if_expr.hir_id, final_ty);
+        }
+        if let Err(guar) = final_ty.error_reported() {
+            self.set_tainted_by_errors(guar);
+        }
+
+        final_ty
+    }
+
+    fn coerce_if_arm(
+        &self,
+        coerce: &mut DynamicCoerceMany<'tcx>,
+        cause: &mut traits::ObligationCause<'tcx>,
+        expr: &'tcx hir::Expr<'tcx>,
+        ty: Ty<'tcx>,
+        cause_span: Option<Span>,
+        body_span: Span,
+        prior_branches: SmallVec<[(Ty<'tcx>, Span); 4]>,
+    ) {
+        if let Some(span) = cause_span {
+            cause.span = span;
+        }
+        coerce.coerce_inner(
+            self,
+            cause,
+            Some(expr),
+            ty,
+            move |err| self.explain_if_branch_mismatch(err, body_span, &prior_branches),
+            false,
+        );
+    }
+
+    fn check_if_condition(
+        &self,
+        cond_expr: &'tcx hir::Expr<'tcx>,
+        then_span: Span,
+    ) -> (Ty<'tcx>, Diverges) {
+        let cond_ty = self.check_expr_has_type_or_error(cond_expr, self.tcx.types.bool, |_| {});
+        self.warn_if_unreachable(
+            cond_expr.hir_id,
+            then_span,
+            "block in `if` or `while` expression",
+        );
+        let cond_diverges = self.take_diverges();
+        (cond_ty, cond_diverges)
+    }
+
+    fn collect_if_chain(
+        &self,
+        mut current_if: &'tcx hir::Expr<'tcx>,
+        expected: Expectation<'tcx>,
+    ) -> (IfChain<'tcx>, IfChainTail<'tcx>) {
+        let mut chain: IfChain<'tcx> = IfChain::default();
+
+        loop {
+            let Some(else_expr) = self.collect_if_branch(current_if, expected, &mut chain) else {
+                return (chain, IfChainTail::Missing(current_if));
+            };
+
+            if let hir::ExprKind::If(..) = else_expr.kind {
+                current_if = else_expr;
+                continue;
+            }
+
+            return (chain, IfChainTail::FinalElse(self.collect_final_else(else_expr, expected)));
+        }
+    }
+
+    fn collect_if_branch(
+        &self,
+        if_expr: &'tcx hir::Expr<'tcx>,
+        expected: Expectation<'tcx>,
+        chain: &mut IfChain<'tcx>,
+    ) -> Option<&'tcx hir::Expr<'tcx>> {
+        let hir::ExprKind::If(cond_expr, then_expr, opt_else_expr) = if_expr.kind else {
+            bug!("expected `if` expression, found {:#?}", if_expr);
+        };
+
+        let (cond_ty, cond_diverges) = self.check_if_condition(cond_expr, then_expr.span);
+        if let Err(guar) = cond_ty.error_reported() {
+            chain.cond_error.get_or_insert(guar);
+        }
+        let branch_body = self.check_branch_body(then_expr, expected);
+
+        chain.branches.push(IfBranch { if_expr, cond_diverges, body: branch_body });
+
+        opt_else_expr
+    }
+
+    fn collect_final_else(
+        &self,
+        else_expr: &'tcx hir::Expr<'tcx>,
+        expected: Expectation<'tcx>,
+    ) -> BranchBody<'tcx> {
+        self.check_branch_body(else_expr, expected)
+    }
+
+    fn reset_diverges_to_maybe(&self) {
+        self.diverges.set(Diverges::Maybe);
+    }
+
+    fn take_diverges(&self) -> Diverges {
+        let diverges = self.diverges.get();
+        self.reset_diverges_to_maybe();
+        diverges
+    }
+
+    fn ensure_if_branch_type(&self, hir_id: HirId, ty: Ty<'tcx>) {
+        let mut typeck = self.typeck_results.borrow_mut();
+        let mut node_ty = typeck.node_types_mut();
+        node_ty.entry(hir_id).or_insert(ty);
+    }
+
+    fn overwrite_if_branch_type(&self, hir_id: HirId, ty: Ty<'tcx>) {
+        let mut typeck = self.typeck_results.borrow_mut();
+        let mut node_ty = typeck.node_types_mut();
+        node_ty.insert(hir_id, ty);
+    }
+
+    fn check_branch_body(
+        &self,
+        expr: &'tcx hir::Expr<'tcx>,
+        expected: Expectation<'tcx>,
+    ) -> BranchBody<'tcx> {
+        self.reset_diverges_to_maybe();
+        let ty = self.check_expr_with_expectation(expr, expected);
+        let diverges = self.take_diverges();
+        let span = self.find_block_span_from_hir_id(expr.hir_id);
+        BranchBody { expr, ty, diverges, span }
+    }
+
+    fn explain_if_branch_mismatch(
+        &self,
+        err: &mut Diag<'_>,
+        branch_span: Span,
+        prior_branches: &[(Ty<'tcx>, Span)],
+    ) {
+        let Some(&(.., prior_span)) =
+            prior_branches.iter().rev().find(|&&(_, span)| span != branch_span)
+        else {
+            return;
+        };
+
+        err.span_label(prior_span, "expected because of this");
+    }
+}