From 9a1253c5c75b4fd16c6cd37c9e6ba41241e384ff Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Tue, 8 May 2018 16:24:50 +0300
Subject: [PATCH] Specify type explicitly: suggest also types from overridden

So #KT-22092 Fixed
---
 .../SpecifyTypeExplicitlyIntention.kt         | 37 +++++++++++++++----
 .../KotlinInplaceVariableIntroducer.java      |  2 +-
 .../specifyTypeExplicitly/overriddenAsNull.kt |  9 +++++
 .../overriddenAsNull.kt.after                 |  9 +++++
 .../intentions/IntentionTestGenerated.java    |  5 +++
 5 files changed, 54 insertions(+), 8 deletions(-)
 create mode 100644 idea/testData/intentions/specifyTypeExplicitly/overriddenAsNull.kt
 create mode 100644 idea/testData/intentions/specifyTypeExplicitly/overriddenAsNull.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/SpecifyTypeExplicitlyIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/SpecifyTypeExplicitlyIntention.kt
index 77358c8b2b5..40011d9a360 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/SpecifyTypeExplicitlyIntention.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/SpecifyTypeExplicitlyIntention.kt
@@ -119,7 +119,11 @@ class SpecifyTypeExplicitlyIntention : SelfTargetingRangeIntention<KtCallableDec
             return type ?: ErrorUtils.createErrorType("null type")
         }
 
-        fun createTypeExpressionForTemplate(exprType: KotlinType, contextElement: KtElement): Expression? {
+        fun createTypeExpressionForTemplate(
+            exprType: KotlinType,
+            contextElement: KtDeclaration,
+            useTypesFromOverridden: Boolean = false
+        ): Expression? {
             val resolutionFacade = contextElement.getResolutionFacade()
             val bindingContext = resolutionFacade.analyze(contextElement, BodyResolveMode.PARTIAL)
             val scope = contextElement.getResolutionScope(bindingContext, resolutionFacade)
@@ -133,13 +137,24 @@ class SpecifyTypeExplicitlyIntention : SelfTargetingRangeIntention<KtCallableDec
                 }
             }
 
-            val types = with(exprType.getResolvableApproximations(scope, checkTypeParameters).toList()) {
-                when {
-                    exprType.isNullabilityFlexible() -> flatMap {
-                        listOf(TypeUtils.makeNotNullable(it), TypeUtils.makeNullable(it))
+            fun KotlinType.toResolvableApproximations(): List<KotlinType> =
+                with(getResolvableApproximations(scope, checkTypeParameters).toList()) {
+                    when {
+                        exprType.isNullabilityFlexible() -> flatMap {
+                            listOf(TypeUtils.makeNotNullable(it), TypeUtils.makeNullable(it))
+                        }
+                        else -> this
                     }
-                    else -> this
                 }
+
+            val overriddenTypes: List<KotlinType> = if (!useTypesFromOverridden) {
+                null
+            } else {
+                val declarationDescriptor = contextElement.resolveToDescriptorIfAny() as? CallableDescriptor
+                declarationDescriptor?.overriddenDescriptors?.mapNotNull { it.returnType }
+            } ?: emptyList()
+            val types = (listOf(exprType) + overriddenTypes).distinct().flatMap {
+                it.toResolvableApproximations()
             }.ifEmpty { return null }
 
             if (ApplicationManager.getApplication().isUnitTestMode) {
@@ -150,6 +165,14 @@ class SpecifyTypeExplicitlyIntention : SelfTargetingRangeIntention<KtCallableDec
                     val targetType = types.firstOrNull { it.isMarkedNullable } ?: types.first()
                     return TypeChooseValueExpression(listOf(targetType), targetType)
                 }
+                // This helps to be sure something except Nothing is suggested
+                if (contextElement.containingKtFile.findDescendantOfType<PsiComment>()?.takeIf {
+                        it.text == "// DO_NOT_CHOOSE_NOTHING"
+                    } != null) {
+                    val targetType = types.firstOrNull { !KotlinBuiltIns.isNothingOrNullableNothing(it) } ?: types.first()
+                    return TypeChooseValueExpression(listOf(targetType), targetType)
+                }
+
             }
 
             return TypeChooseValueExpression(types, types.first())
@@ -204,7 +227,7 @@ class SpecifyTypeExplicitlyIntention : SelfTargetingRangeIntention<KtCallableDec
             assert(!exprType.isError) { "Unexpected error type, should have been checked before: " + declaration.getElementTextWithContext() + ", type = " + exprType }
 
             val project = declaration.project
-            val expression = createTypeExpressionForTemplate(exprType, declaration) ?: return
+            val expression = createTypeExpressionForTemplate(exprType, declaration, useTypesFromOverridden = true) ?: return
 
             declaration.setType(KotlinBuiltIns.FQ_NAMES.any.asString())
 
diff --git a/idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceVariable/KotlinInplaceVariableIntroducer.java b/idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceVariable/KotlinInplaceVariableIntroducer.java
index 2a1b222797c..5bfaa6a6689 100644
--- a/idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceVariable/KotlinInplaceVariableIntroducer.java
+++ b/idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceVariable/KotlinInplaceVariableIntroducer.java
@@ -343,7 +343,7 @@ public class KotlinInplaceVariableIntroducer<D extends KtCallableDeclaration> ex
 
     protected void addTypeReferenceVariable(TemplateBuilderImpl builder) {
         KtTypeReference typeReference = myDeclaration.getTypeReference();
-        Expression expression = SpecifyTypeExplicitlyIntention.Companion.createTypeExpressionForTemplate(myExprType, myDeclaration);
+        Expression expression = SpecifyTypeExplicitlyIntention.Companion.createTypeExpressionForTemplate(myExprType, myDeclaration, false);
         if (typeReference != null && expression != null) {
             builder.replaceElement(typeReference, TYPE_REFERENCE_VARIABLE_NAME, expression, false);
         }
diff --git a/idea/testData/intentions/specifyTypeExplicitly/overriddenAsNull.kt b/idea/testData/intentions/specifyTypeExplicitly/overriddenAsNull.kt
new file mode 100644
index 00000000000..49facbbf905
--- /dev/null
+++ b/idea/testData/intentions/specifyTypeExplicitly/overriddenAsNull.kt
@@ -0,0 +1,9 @@
+// DO_NOT_CHOOSE_NOTHING
+
+interface I {
+    fun foo(): String?
+}
+
+class Test : I {
+    override fun foo()<caret> = null
+}
diff --git a/idea/testData/intentions/specifyTypeExplicitly/overriddenAsNull.kt.after b/idea/testData/intentions/specifyTypeExplicitly/overriddenAsNull.kt.after
new file mode 100644
index 00000000000..f8e1cbcf895
--- /dev/null
+++ b/idea/testData/intentions/specifyTypeExplicitly/overriddenAsNull.kt.after
@@ -0,0 +1,9 @@
+// DO_NOT_CHOOSE_NOTHING
+
+interface I {
+    fun foo(): String?
+}
+
+class Test : I {
+    override fun foo(): String? = null
+}
diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
index 33d332b923a..c26c78cd899 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
@@ -14574,6 +14574,11 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             runTest("idea/testData/intentions/specifyTypeExplicitly/loopParameter.kt");
         }
 
+        @TestMetadata("overriddenAsNull.kt")
+        public void testOverriddenAsNull() throws Exception {
+            runTest("idea/testData/intentions/specifyTypeExplicitly/overriddenAsNull.kt");
+        }
+
         @TestMetadata("overrideNotNullFunction.kt")
         public void testOverrideNotNullFunction() throws Exception {
             runTest("idea/testData/intentions/specifyTypeExplicitly/overrideNotNullFunction.kt");