From 4b4eabad3c3ce45812b08692b78d5c0d2e0d4dc2 Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Fri, 27 Jan 2023 11:25:34 +0100
Subject: [PATCH] Record method resolution for try expressions

---
 crates/hir-ty/src/infer.rs      | 26 +++++++++++++-------------
 crates/hir-ty/src/infer/expr.rs | 20 ++++++++++++++++----
 2 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/crates/hir-ty/src/infer.rs b/crates/hir-ty/src/infer.rs
index 4402c75947..571b3e9686 100644
--- a/crates/hir-ty/src/infer.rs
+++ b/crates/hir-ty/src/infer.rs
@@ -938,19 +938,24 @@ impl<'a> InferenceContext<'a> {
         self.db.trait_data(trait_).associated_type_by_name(&name![Item])
     }
 
-    fn resolve_ops_try_ok(&self) -> Option<TypeAliasId> {
-        let trait_ = self.resolve_lang_item(LangItem::Try)?.as_trait()?;
+    fn resolve_output_on(&self, trait_: TraitId) -> Option<TypeAliasId> {
         self.db.trait_data(trait_).associated_type_by_name(&name![Output])
     }
 
+    fn resolve_lang_trait(&self, lang: LangItem) -> Option<TraitId> {
+        self.resolve_lang_item(lang)?.as_trait()
+    }
+
+    fn resolve_ops_try_output(&self) -> Option<TypeAliasId> {
+        self.resolve_output_on(self.resolve_lang_trait(LangItem::Try)?)
+    }
+
     fn resolve_ops_neg_output(&self) -> Option<TypeAliasId> {
-        let trait_ = self.resolve_lang_item(LangItem::Neg)?.as_trait()?;
-        self.db.trait_data(trait_).associated_type_by_name(&name![Output])
+        self.resolve_output_on(self.resolve_lang_trait(LangItem::Neg)?)
     }
 
     fn resolve_ops_not_output(&self) -> Option<TypeAliasId> {
-        let trait_ = self.resolve_lang_item(LangItem::Not)?.as_trait()?;
-        self.db.trait_data(trait_).associated_type_by_name(&name![Output])
+        self.resolve_output_on(self.resolve_lang_trait(LangItem::Not)?)
     }
 
     fn resolve_future_future_output(&self) -> Option<TypeAliasId> {
@@ -960,7 +965,7 @@ impl<'a> InferenceContext<'a> {
             .lookup(self.db.upcast())
             .container
         else { return None };
-        self.db.trait_data(trait_).associated_type_by_name(&name![Output])
+        self.resolve_output_on(trait_)
     }
 
     fn resolve_boxed_box(&self) -> Option<AdtId> {
@@ -998,13 +1003,8 @@ impl<'a> InferenceContext<'a> {
         Some(struct_.into())
     }
 
-    fn resolve_ops_index(&self) -> Option<TraitId> {
-        self.resolve_lang_item(LangItem::Index)?.as_trait()
-    }
-
     fn resolve_ops_index_output(&self) -> Option<TypeAliasId> {
-        let trait_ = self.resolve_ops_index()?;
-        self.db.trait_data(trait_).associated_type_by_name(&name![Output])
+        self.resolve_output_on(self.resolve_lang_trait(LangItem::Index)?)
     }
 
     fn resolve_va_list(&self) -> Option<AdtId> {
diff --git a/crates/hir-ty/src/infer/expr.rs b/crates/hir-ty/src/infer/expr.rs
index 1f8fa4a20f..7b9bf0c5cf 100644
--- a/crates/hir-ty/src/infer/expr.rs
+++ b/crates/hir-ty/src/infer/expr.rs
@@ -13,6 +13,7 @@ use hir_def::{
         ArithOp, Array, BinaryOp, ClosureKind, Expr, ExprId, LabelId, Literal, Statement, UnaryOp,
     },
     generics::TypeOrConstParamData,
+    lang_item::LangItem,
     path::{GenericArg, GenericArgs},
     resolver::resolver_for_expr,
     ConstParamId, FieldId, ItemContainerId, Lookup,
@@ -157,7 +158,8 @@ impl<'a> InferenceContext<'a> {
                 }
 
                 // The ok-ish type that is expected from the last expression
-                let ok_ty = self.resolve_associated_type(try_ty.clone(), self.resolve_ops_try_ok());
+                let ok_ty =
+                    self.resolve_associated_type(try_ty.clone(), self.resolve_ops_try_output());
 
                 self.with_breakable_ctx(BreakableKind::Block, ok_ty.clone(), None, |this| {
                     this.infer_expr(*body, &Expectation::has_type(ok_ty));
@@ -593,9 +595,19 @@ impl<'a> InferenceContext<'a> {
                 self.resolve_associated_type(inner_ty, self.resolve_future_future_output())
             }
             Expr::Try { expr } => {
-                // FIXME: Note down method resolution her
                 let inner_ty = self.infer_expr_inner(*expr, &Expectation::none());
-                self.resolve_associated_type(inner_ty, self.resolve_ops_try_ok())
+                if let Some(trait_) = self.resolve_lang_trait(LangItem::Try) {
+                    if let Some(func) = self.db.trait_data(trait_).method_by_name(&name!(branch)) {
+                        let subst = TyBuilder::subst_for_def(self.db, trait_, None)
+                            .push(inner_ty.clone())
+                            .build();
+                        self.write_method_resolution(tgt_expr, func, subst.clone());
+                    }
+                    let try_output = self.resolve_output_on(trait_);
+                    self.resolve_associated_type(inner_ty, try_output)
+                } else {
+                    self.err_ty()
+                }
             }
             Expr::Cast { expr, type_ref } => {
                 // FIXME: propagate the "castable to" expectation (and find a test case that shows this is necessary)
@@ -744,7 +756,7 @@ impl<'a> InferenceContext<'a> {
                 let base_ty = self.infer_expr_inner(*base, &Expectation::none());
                 let index_ty = self.infer_expr(*index, &Expectation::none());
 
-                if let Some(index_trait) = self.resolve_ops_index() {
+                if let Some(index_trait) = self.resolve_lang_trait(LangItem::Index) {
                     let canonicalized = self.canonicalize(base_ty.clone());
                     let receiver_adjustments = method_resolution::resolve_indexing_op(
                         self.db,