From 319611b7382fc4c84170519dade68f4f558a44b1 Mon Sep 17 00:00:00 2001
From: unexge <unexge@gmail.com>
Date: Sat, 29 Oct 2022 23:44:34 +0100
Subject: [PATCH] Record diverging match arms in `InferenceResult`

---
 crates/hir-ty/src/infer.rs        |  2 ++
 crates/hir-ty/src/infer/expr.rs   |  5 +++++
 crates/hir/src/semantics.rs       |  8 ++++++++
 crates/hir/src/source_analyzer.rs | 15 +++++++++++++++
 4 files changed, 30 insertions(+)

diff --git a/crates/hir-ty/src/infer.rs b/crates/hir-ty/src/infer.rs
index 31e56dec62593..05776e192193a 100644
--- a/crates/hir-ty/src/infer.rs
+++ b/crates/hir-ty/src/infer.rs
@@ -333,6 +333,8 @@ pub struct InferenceResult {
     assoc_resolutions: FxHashMap<ExprOrPatId, AssocItemId>,
     pub diagnostics: Vec<InferenceDiagnostic>,
     pub type_of_expr: ArenaMap<ExprId, Ty>,
+    /// For each match expr, record diverging arm's expr.
+    pub diverging_arms: FxHashMap<ExprId, Vec<ExprId>>,
     /// For each pattern record the type it resolves to.
     ///
     /// **Note**: When a pattern type is resolved it may still contain
diff --git a/crates/hir-ty/src/infer/expr.rs b/crates/hir-ty/src/infer/expr.rs
index f56108b26c45b..3d2e091a0f6ad 100644
--- a/crates/hir-ty/src/infer/expr.rs
+++ b/crates/hir-ty/src/infer/expr.rs
@@ -375,6 +375,7 @@ impl<'a> InferenceContext<'a> {
 
                 let matchee_diverges = self.diverges;
                 let mut all_arms_diverge = Diverges::Always;
+                let mut diverging_arms = Vec::new();
 
                 for arm in arms.iter() {
                     self.diverges = Diverges::Maybe;
@@ -387,11 +388,15 @@ impl<'a> InferenceContext<'a> {
                     }
 
                     let arm_ty = self.infer_expr_inner(arm.expr, &expected);
+                    if self.diverges.is_always() {
+                        diverging_arms.push(arm.expr);
+                    }
                     all_arms_diverge &= self.diverges;
                     coerce.coerce(self, Some(arm.expr), &arm_ty);
                 }
 
                 self.diverges = matchee_diverges | all_arms_diverge;
+                self.result.diverging_arms.insert(tgt_expr, diverging_arms);
 
                 coerce.complete()
             }
diff --git a/crates/hir/src/semantics.rs b/crates/hir/src/semantics.rs
index 119ec3210e175..2f835657d37a1 100644
--- a/crates/hir/src/semantics.rs
+++ b/crates/hir/src/semantics.rs
@@ -481,6 +481,10 @@ impl<'db, DB: HirDatabase> Semantics<'db, DB> {
     pub fn is_unsafe_ident_pat(&self, ident_pat: &ast::IdentPat) -> bool {
         self.imp.is_unsafe_ident_pat(ident_pat)
     }
+
+    pub fn is_diverging_match_arm(&self, match_arm: &ast::MatchArm) -> Option<bool> {
+        self.imp.is_diverging_match_arm(match_arm)
+    }
 }
 
 impl<'db> SemanticsImpl<'db> {
@@ -1421,6 +1425,10 @@ impl<'db> SemanticsImpl<'db> {
             .map(|ty| ty.original.is_packed(self.db))
             .unwrap_or(false)
     }
+
+    fn is_diverging_match_arm(&self, match_arm: &ast::MatchArm) -> Option<bool> {
+        self.analyze(match_arm.syntax())?.is_diverging_match_arm(self.db, match_arm)
+    }
 }
 
 fn macro_call_to_macro_id(
diff --git a/crates/hir/src/source_analyzer.rs b/crates/hir/src/source_analyzer.rs
index f86c571005367..2e61946a738de 100644
--- a/crates/hir/src/source_analyzer.rs
+++ b/crates/hir/src/source_analyzer.rs
@@ -782,6 +782,21 @@ impl SourceAnalyzer {
         false
     }
 
+    pub(crate) fn is_diverging_match_arm(
+        &self,
+        db: &dyn HirDatabase,
+        match_arm: &ast::MatchArm,
+    ) -> Option<bool> {
+        let infer = self.infer.as_ref()?;
+        let match_expr = match_arm.syntax().ancestors().find_map(ast::MatchExpr::cast)?;
+        let match_id = self.expr_id(db, &match_expr.into())?;
+        let diverging_arms = infer.diverging_arms.get(&match_id)?;
+        let match_arm_expr = match_arm.expr()?;
+        let match_arm_expr_id = self.expr_id(db, &match_arm_expr)?;
+
+        Some(diverging_arms.contains(&match_arm_expr_id))
+    }
+
     fn resolve_impl_method_or_trait_def(
         &self,
         db: &dyn HirDatabase,