TY: Fix type inference for await

intellij-rust · Sep 12, 2019 · 189e219 · 189e219
1 parent 413fd6e
commit 189e219
Show file tree

Hide file tree

Showing 5 changed files with 76 additions and 17 deletions.
diff --git a/src/main/kotlin/org/rust/lang/core/completion/RsKeywordCompletionContributor.kt b/src/main/kotlin/org/rust/lang/core/completion/RsKeywordCompletionContributor.kt
@@ -22,7 +22,7 @@ import org.rust.lang.core.psi.*
 import org.rust.lang.core.psi.RsElementTypes.*
 import org.rust.lang.core.psi.ext.*
 import org.rust.lang.core.resolve.ImplLookup
-import org.rust.lang.core.types.infer.lookupFutureOutputTy
+import org.rust.lang.core.types.TraitRef
 import org.rust.lang.core.types.ty.Ty
 import org.rust.lang.core.types.ty.TyUnknown
 import org.rust.lang.core.types.type
@@ -276,6 +276,13 @@ class RsKeywordCompletionContributor : CompletionContributor(), DumbAware {
         return baseInherentImplDeclarationPattern().and(statementBeginningPattern("pub"))
     }
 
+    private fun Ty.lookupFutureOutputTy(lookup: ImplLookup): Ty {
+        val futureTrait = lookup.items.Future ?: return TyUnknown
+        val outputType = futureTrait.findAssociatedType("Output") ?: return TyUnknown
+        val selection = lookup.selectProjectionStrict(TraitRef(this, futureTrait.withSubst()), outputType)
+        return selection.ok()?.value ?: TyUnknown
+    }
+
     companion object {
         @JvmField
         val CONDITION_KEYWORDS: List<String> = listOf("if", "match")

diff --git a/src/main/kotlin/org/rust/lang/core/types/infer/TypeInference.kt b/src/main/kotlin/org/rust/lang/core/types/infer/TypeInference.kt
@@ -770,17 +770,6 @@ data class TyWithObligations<out T>(
 fun <T> TyWithObligations<T>.withObligations(addObligations: List<Obligation>) =
     TyWithObligations(value, obligations + addObligations)
 
-fun Ty.lookupFutureOutputTy(lookup: ImplLookup): Ty {
-    if (this !is TyAnon) return TyUnknown
-    val futureTrait = lookup.items.Future ?: return TyUnknown
-    val outputType = futureTrait.findAssociatedType("Output") ?: return TyUnknown
-    return lookup.lookupAssocTypeInBounds(
-        getTraitBoundsTransitively().asSequence(),
-        futureTrait,
-        outputType
-    ) ?: TyUnknown
-}
-
 sealed class ResolvedPath {
     abstract val element: RsElement
 

diff --git a/src/main/kotlin/org/rust/lang/core/types/infer/TypeInferenceWalker.kt b/src/main/kotlin/org/rust/lang/core/types/infer/TypeInferenceWalker.kt
@@ -606,7 +606,7 @@ class RsTypeInferenceWalker(
         return methodType.retType
     }
 
-    private fun <T: AssocItemScopeEntryBase<E>, E> filterAssocItems(variants: List<T>, context: RsElement): List<T> {
+    private fun <T : AssocItemScopeEntryBase<E>, E> filterAssocItems(variants: List<T>, context: RsElement): List<T> {
         return variants.singleOrLet { list ->
             // 1. filter traits that are not imported
             TypeInferenceMarks.methodPickTraitScope.hit()
@@ -1288,6 +1288,13 @@ class RsTypeInferenceWalker(
 
     fun writePatFieldTy(psi: RsPatField, ty: Ty): Unit =
         ctx.writePatFieldTy(psi, ty)
+
+    private fun Ty.lookupFutureOutputTy(lookup: ImplLookup): Ty {
+        val futureTrait = lookup.items.Future ?: return TyUnknown
+        val outputType = futureTrait.findAssociatedType("Output") ?: return TyUnknown
+        val selection = lookup.selectProjection(TraitRef(this, futureTrait.withSubst()), outputType)
+        return selection.ok()?.register() ?: TyUnknown
+    }
 }
 
 private val RsSelfParameter.typeOfValue: Ty

diff --git a/src/test/kotlin/org/rust/lang/core/completion/RsKeywordCompletionContributorTest.kt b/src/test/kotlin/org/rust/lang/core/completion/RsKeywordCompletionContributorTest.kt
@@ -519,7 +519,7 @@ class RsKeywordCompletionContributorTest : RsCompletionTestBase() {
     """)
 
     @MockEdition(CargoWorkspace.Edition.EDITION_2015)
-    fun `test postfix await 2015`() = checkCompletion("await", """
+    fun `test postfix await 2015 (anon)`() = checkCompletion("await", """
         #[lang = "core::future::future::Future"]
         trait Future { type Output; }
         fn foo() -> impl Future<Output=i32> { unimplemented!() }
@@ -535,8 +535,29 @@ class RsKeywordCompletionContributorTest : RsCompletionTestBase() {
         }
     """)
 
+    @MockEdition(CargoWorkspace.Edition.EDITION_2015)
+    fun `test postfix await 2015 (adt)`() = checkCompletion("await", """
+        #[lang = "core::future::future::Future"]
+        trait Future { type Output; }
+        struct S;
+        impl Future for S { type Output = i32; }
+        fn foo() -> S { unimplemented!() }
+        fn main() {
+            foo()./*caret*/;
+        }
+    """, """
+        #[lang = "core::future::future::Future"]
+        trait Future { type Output; }
+        struct S;
+        impl Future for S { type Output = i32; }
+        fn foo() -> S { unimplemented!() }
+        fn main() {
+            foo()./*caret*/;
+        }
+    """)
+
     @MockEdition(CargoWorkspace.Edition.EDITION_2018)
-    fun `test postfix await 2018`() = checkCompletion("await", """
+    fun `test postfix await 2018 (anon)`() = checkCompletion("await", """
         #[lang = "core::future::future::Future"]
         trait Future { type Output; }
         fn foo() -> impl Future<Output=i32> { unimplemented!() }
@@ -552,6 +573,27 @@ class RsKeywordCompletionContributorTest : RsCompletionTestBase() {
         }
     """)
 
+    @MockEdition(CargoWorkspace.Edition.EDITION_2018)
+    fun `test postfix await 2018 (adt)`() = checkCompletion("await", """
+        #[lang = "core::future::future::Future"]
+        trait Future { type Output; }
+        struct S;
+        impl Future for S { type Output = i32; }
+        fn foo() -> S { unimplemented!() }
+        fn main() {
+            foo()./*caret*/;
+        }
+    """, """
+        #[lang = "core::future::future::Future"]
+        trait Future { type Output; }
+        struct S;
+        impl Future for S { type Output = i32; }
+        fn foo() -> S { unimplemented!() }
+        fn main() {
+            foo().await/*caret*/;
+        }
+    """)
+
     fun `test const parameter first`() = checkCompletion("const",
         "fn foo</*caret*/>() {}",
         "fn foo<const /*caret*/>() {}"

diff --git a/src/test/kotlin/org/rust/lang/core/type/RsExpressionTypeInferenceTest.kt b/src/test/kotlin/org/rust/lang/core/type/RsExpressionTypeInferenceTest.kt
@@ -514,7 +514,7 @@ class RsExpressionTypeInferenceTest : RsTypificationTestBase() {
     """)
 
     @MockEdition(CargoWorkspace.Edition.EDITION_2018)
-    fun `test await postfix 1`() = testExpr("""
+    fun `test await postfix 2018 (anon)`() = testExpr("""
         #[lang = "core::future::future::Future"]
         trait Future { type Output; }
         fn main() {
@@ -524,7 +524,21 @@ class RsExpressionTypeInferenceTest : RsTypificationTestBase() {
         }
     """)
 
-    fun `test await postfix 2`() = testExpr("""
+    @MockEdition(CargoWorkspace.Edition.EDITION_2018)
+    fun `test await postfix 2018 (adt)`() = testExpr("""
+        #[lang = "core::future::future::Future"]
+        trait Future { type Output; }
+        struct S;
+        impl Future for S { type Output = i32; }
+        fn foo() -> S { unimplemented!() }
+        fn main() {
+            let x = foo().await;
+            x;
+          //^ i32
+        }
+    """)
+
+    fun `test await postfix 2015`() = testExpr("""
         struct S { await: i32 }
         fn main() {
             let x = (S { await: 42 }).await;