Distinguish between expected and final type in CoerceMany

Author: Veykril

crates/hir-ty/src/infer.rs

@@ -66,8 +66,10 @@ pub(crate) fn infer_query(db: &dyn HirDatabase, def: DefWithBodyId) -> Arc<Infer
     let mut ctx = InferenceContext::new(db, def, &body, resolver);
 
     match def {
+        DefWithBodyId::FunctionId(f) => {
+            ctx.collect_fn(f);
+        }
         DefWithBodyId::ConstId(c) => ctx.collect_const(&db.const_data(c)),
-        DefWithBodyId::FunctionId(f) => ctx.collect_fn(f),
         DefWithBodyId::StaticId(s) => ctx.collect_static(&db.static_data(s)),
         DefWithBodyId::VariantId(v) => {
             ctx.return_ty = TyBuilder::builtin(match db.enum_data(v.parent).variant_body_type() {
@@ -392,9 +394,12 @@ pub(crate) struct InferenceContext<'a> {
     /// currently within one.
     ///
     /// We might consider using a nested inference context for checking
-    /// closures, but currently this is the only field that will change there,
-    /// so it doesn't make sense.
+    /// closures so we can swap all shared things out at once.
     return_ty: Ty,
+    /// If `Some`, this stores coercion information for returned
+    /// expressions. If `None`, this is in a context where return is
+    /// inappropriate, such as a const expression.
+    return_coercion: Option<CoerceMany>,
     /// The resume type and the yield type, respectively, of the generator being inferred.
     resume_yield_tys: Option<(Ty, Ty)>,
     diverges: Diverges,
@@ -462,6 +467,7 @@ impl<'a> InferenceContext<'a> {
             trait_env,
             return_ty: TyKind::Error.intern(Interner), // set in collect_* calls
             resume_yield_tys: None,
+            return_coercion: None,
             db,
             owner,
             body,
@@ -595,10 +601,19 @@ impl<'a> InferenceContext<'a> {
         };
 
         self.return_ty = self.normalize_associated_types_in(return_ty);
+        self.return_coercion = Some(CoerceMany::new(self.return_ty.clone()));
     }
 
     fn infer_body(&mut self) {
-        self.infer_expr_coerce(self.body.body_expr, &Expectation::has_type(self.return_ty.clone()));
+        match self.return_coercion {
+            Some(_) => self.infer_return(self.body.body_expr),
+            None => {
+                _ = self.infer_expr_coerce(
+                    self.body.body_expr,
+                    &Expectation::has_type(self.return_ty.clone()),
+                )
+            }
+        }
     }
 
     fn write_expr_ty(&mut self, expr: ExprId, ty: Ty) {

crates/hir-ty/src/infer/coerce.rs

@@ -50,11 +50,44 @@ fn success(
 #[derive(Clone, Debug)]
 pub(super) struct CoerceMany {
     expected_ty: Ty,
+    final_ty: Option<Ty>,
 }
 
 impl CoerceMany {
     pub(super) fn new(expected: Ty) -> Self {
-        CoerceMany { expected_ty: expected }
+        CoerceMany { expected_ty: expected, final_ty: None }
+    }
+
+    /// Returns the "expected type" with which this coercion was
+    /// constructed. This represents the "downward propagated" type
+    /// that was given to us at the start of typing whatever construct
+    /// we are typing (e.g., the match expression).
+    ///
+    /// Typically, this is used as the expected type when
+    /// type-checking each of the alternative expressions whose types
+    /// we are trying to merge.
+    pub(super) fn expected_ty(&self) -> Ty {
+        self.expected_ty.clone()
+    }
+
+    /// Returns the current "merged type", representing our best-guess
+    /// at the LUB of the expressions we've seen so far (if any). This
+    /// isn't *final* until you call `self.complete()`, which will return
+    /// the merged type.
+    pub(super) fn merged_ty(&self) -> Ty {
+        self.final_ty.clone().unwrap_or_else(|| self.expected_ty.clone())
+    }
+
+    pub(super) fn complete(self, ctx: &mut InferenceContext<'_>) -> Ty {
+        if let Some(final_ty) = self.final_ty {
+            final_ty
+        } else {
+            ctx.result.standard_types.never.clone()
+        }
+    }
+
+    pub(super) fn coerce_forced_unit(&mut self, ctx: &mut InferenceContext<'_>) {
+        self.coerce(ctx, None, &ctx.result.standard_types.unit.clone())
     }
 
     /// Merge two types from different branches, with possible coercion.
@@ -76,51 +109,46 @@ impl CoerceMany {
         // Special case: two function types. Try to coerce both to
         // pointers to have a chance at getting a match. See
         // https://github.com/rust-lang/rust/blob/7b805396bf46dce972692a6846ce2ad8481c5f85/src/librustc_typeck/check/coercion.rs#L877-L916
-        let sig = match (self.expected_ty.kind(Interner), expr_ty.kind(Interner)) {
+        let sig = match (self.merged_ty().kind(Interner), expr_ty.kind(Interner)) {
             (TyKind::FnDef(..) | TyKind::Closure(..), TyKind::FnDef(..) | TyKind::Closure(..)) => {
                 // FIXME: we're ignoring safety here. To be more correct, if we have one FnDef and one Closure,
                 // we should be coercing the closure to a fn pointer of the safety of the FnDef
                 cov_mark::hit!(coerce_fn_reification);
                 let sig =
-                    self.expected_ty.callable_sig(ctx.db).expect("FnDef without callable sig");
+                    self.merged_ty().callable_sig(ctx.db).expect("FnDef without callable sig");
                 Some(sig)
             }
             _ => None,
         };
         if let Some(sig) = sig {
             let target_ty = TyKind::Function(sig.to_fn_ptr()).intern(Interner);
-            let result1 = ctx.table.coerce_inner(self.expected_ty.clone(), &target_ty);
+            let result1 = ctx.table.coerce_inner(self.merged_ty(), &target_ty);
             let result2 = ctx.table.coerce_inner(expr_ty.clone(), &target_ty);
             if let (Ok(result1), Ok(result2)) = (result1, result2) {
                 ctx.table.register_infer_ok(result1);
                 ctx.table.register_infer_ok(result2);
-                return self.expected_ty = target_ty;
+                return self.final_ty = Some(target_ty);
             }
         }
 
         // It might not seem like it, but order is important here: If the expected
         // type is a type variable and the new one is `!`, trying it the other
         // way around first would mean we make the type variable `!`, instead of
         // just marking it as possibly diverging.
-        if ctx.coerce(expr, &expr_ty, &self.expected_ty).is_ok() {
-            /* self.expected_ty is already correct */
-        } else if ctx.coerce(expr, &self.expected_ty, &expr_ty).is_ok() {
-            self.expected_ty = expr_ty;
+        if let Ok(res) = ctx.coerce(expr, &expr_ty, &self.merged_ty()) {
+            self.final_ty = Some(res);
+        } else if let Ok(res) = ctx.coerce(expr, &self.merged_ty(), &expr_ty) {
+            self.final_ty = Some(res);
         } else {
             if let Some(id) = expr {
                 ctx.result.type_mismatches.insert(
                     id.into(),
-                    TypeMismatch { expected: self.expected_ty.clone(), actual: expr_ty },
+                    TypeMismatch { expected: self.merged_ty().clone(), actual: expr_ty.clone() },
                 );
             }
             cov_mark::hit!(coerce_merge_fail_fallback);
-            /* self.expected_ty is already correct */
         }
     }
-
-    pub(super) fn complete(self) -> Ty {
-        self.expected_ty
-    }
 }
 
 pub fn could_coerce(