Add type_mismatches to InferenceResult and use this in ok-wrapping code fix

3 files changed, 27 insertions, 8 deletions

diff --git a/crates/ra_hir/src/expr/validation.rs b/crates/ra_hir/src/expr/validation.rs
index f5e641557..0e7b08c54 100644
--- a/crates/ra_hir/src/expr/validation.rs
+++ b/crates/ra_hir/src/expr/validation.rs
@@ -102,14 +102,16 @@ impl<'a, 'b> ExprValidator<'a, 'b> {
     }
 
     fn validate_results_in_tail_expr(&mut self, id: ExprId, db: &impl HirDatabase) {
-        let expr_ty = &self.infer[id];
-        let func_ty = self.func.ty(db);
-        let func_sig = func_ty.callable_sig(db).unwrap();
-        let ret = func_sig.ret();
-        let ret = match ret {
+        let mismatch = match self.infer.type_mismatch_for_expr(id) {
+            Some(m) => m,
+            None => return,
+        };
+
+        let ret = match &mismatch.expected {
             Ty::Apply(t) => t,
             _ => return,
         };
+
         let ret_enum = match ret.ctor {
             TypeCtor::Adt(AdtDef::Enum(e)) => e,
             _ => return,
@@ -119,7 +121,7 @@ impl<'a, 'b> ExprValidator<'a, 'b> {
             return;
         }
         let params = &ret.parameters;
-        if params.len() == 2 && &params[0] == expr_ty {
+        if params.len() == 2 && &params[0] == &mismatch.actual {
             let source_map = self.func.body_source_map(db);
             let file_id = self.func.source(db).file_id;
             let parse = db.parse(file_id.original_file(db));
diff --git a/crates/ra_hir/src/ty.rs b/crates/ra_hir/src/ty.rs
index 0efd94cef..b54c80318 100644
--- a/crates/ra_hir/src/ty.rs
+++ b/crates/ra_hir/src/ty.rs
@@ -516,7 +516,7 @@ impl Ty {
         }
     }
 
-    pub fn callable_sig(&self, db: &impl HirDatabase) -> Option<FnSig> {
+    fn callable_sig(&self, db: &impl HirDatabase) -> Option<FnSig> {
         match self {
             Ty::Apply(a_ty) => match a_ty.ctor {
                 TypeCtor::FnPtr { .. } => Some(FnSig::from_fn_ptr_substs(&a_ty.parameters)),
diff --git a/crates/ra_hir/src/ty/infer.rs b/crates/ra_hir/src/ty/infer.rs
index b33de5687..d94e8154b 100644
--- a/crates/ra_hir/src/ty/infer.rs
+++ b/crates/ra_hir/src/ty/infer.rs
@@ -106,6 +106,13 @@ impl Default for BindingMode {
     }
 }
 
+/// A mismatch between an expected and an inferred type.
+#[derive(Clone, PartialEq, Eq, Debug, Hash)]
+pub struct TypeMismatch {
+    pub expected: Ty,
+    pub actual: Ty,
+}
+
 /// The result of type inference: A mapping from expressions and patterns to types.
 #[derive(Clone, PartialEq, Eq, Debug, Default)]
 pub struct InferenceResult {
@@ -120,6 +127,7 @@ pub struct InferenceResult {
     diagnostics: Vec<InferenceDiagnostic>,
     pub(super) type_of_expr: ArenaMap<ExprId, Ty>,
     pub(super) type_of_pat: ArenaMap<PatId, Ty>,
+    pub(super) type_mismatches: ArenaMap<ExprId, TypeMismatch>,
 }
 
 impl InferenceResult {
@@ -141,6 +149,9 @@ impl InferenceResult {
     pub fn assoc_resolutions_for_pat(&self, id: PatId) -> Option<ImplItem> {
         self.assoc_resolutions.get(&id.into()).copied()
     }
+    pub fn type_mismatch_for_expr(&self, expr: ExprId) -> Option<&TypeMismatch> {
+        self.type_mismatches.get(expr)
+    }
     pub(crate) fn add_diagnostics(
         &self,
         db: &impl HirDatabase,
@@ -1345,9 +1356,15 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
         };
         // use a new type variable if we got Ty::Unknown here
         let ty = self.insert_type_vars_shallow(ty);
-        self.unify(&ty, &expected.ty);
+        let could_unify = self.unify(&ty, &expected.ty);
         let ty = self.resolve_ty_as_possible(&mut vec![], ty);
         self.write_expr_ty(tgt_expr, ty.clone());
+        if !could_unify {
+            self.result.type_mismatches.insert(
+                tgt_expr,
+                TypeMismatch { expected: expected.ty.clone(), actual: ty.clone() },
+            );
+        }
         ty
     }