From 3ee4e6c54c0bb9586c607dd8890d58fb21c95259 Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Tue, 25 May 2021 14:59:54 +0200
Subject: Fix type inference not working for new Try trait

---
 crates/hir_ty/src/infer.rs        |  5 ++++-
 crates/hir_ty/src/tests/traits.rs | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

(limited to 'crates/hir_ty')

diff --git a/crates/hir_ty/src/infer.rs b/crates/hir_ty/src/infer.rs
index db3c937ff..edb65622f 100644
--- a/crates/hir_ty/src/infer.rs
+++ b/crates/hir_ty/src/infer.rs
@@ -580,7 +580,10 @@ impl<'a> InferenceContext<'a> {
     fn resolve_ops_try_ok(&self) -> Option<TypeAliasId> {
         let path = path![core::ops::Try];
         let trait_ = self.resolver.resolve_known_trait(self.db.upcast(), &path)?;
-        self.db.trait_data(trait_).associated_type_by_name(&name![Ok])
+        let trait_data = self.db.trait_data(trait_);
+        trait_data
+            .associated_type_by_name(&name![Ok])
+            .or_else(|| trait_data.associated_type_by_name(&name![Output]))
     }
 
     fn resolve_ops_neg_output(&self) -> Option<TypeAliasId> {
diff --git a/crates/hir_ty/src/tests/traits.rs b/crates/hir_ty/src/tests/traits.rs
index a5a2df54c..71905baeb 100644
--- a/crates/hir_ty/src/tests/traits.rs
+++ b/crates/hir_ty/src/tests/traits.rs
@@ -160,6 +160,43 @@ mod result {
     );
 }
 
+#[test]
+fn infer_tryv2() {
+    check_types(
+        r#"
+//- /main.rs crate:main deps:core
+fn test() {
+    let r: Result<i32, u64> = Result::Ok(1);
+    let v = r?;
+    v;
+} //^ i32
+
+//- /core.rs crate:core
+#[prelude_import] use ops::*;
+mod ops {
+    trait Try {
+        type Output;
+        type Residual;
+    }
+}
+
+#[prelude_import] use result::*;
+mod result {
+    enum Infallible {}
+    enum Result<O, E> {
+        Ok(O),
+        Err(E)
+    }
+
+    impl<O, E> crate::ops::Try for Result<O, E> {
+        type Output = O;
+        type Error = Result<Infallible, E>;
+    }
+}
+"#,
+    );
+}
+
 #[test]
 fn infer_for_loop() {
     check_types(
-- 
cgit v1.2.3