Inline refactoring: shouldn't lose return type information

#KT-26705 Fixed

idea/src/org/jetbrains/kotlin/idea/refactoring/inline/KotlinInlineFunctionHandler.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -15,16 +15,17 @@ import com.intellij.refactoring.RefactoringBundle
 import com.intellij.refactoring.util.CommonRefactoringUtil
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
+import org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.idea.KotlinLanguage
 import org.jetbrains.kotlin.idea.caches.resolve.analyzeWithContent
 import org.jetbrains.kotlin.idea.caches.resolve.unsafeResolveToDescriptor
 import org.jetbrains.kotlin.idea.codeInliner.CallableUsageReplacementStrategy
-import org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.types.typeUtil.isUnit
 
 class KotlinInlineFunctionHandler : InlineActionHandler() {
     override fun isEnabledForLanguage(language: Language) = language == KotlinLanguage.INSTANCE
@@ -47,10 +48,11 @@ class KotlinInlineFunctionHandler : InlineActionHandler() {
         }
 
         val descriptor = element.unsafeResolveToDescriptor() as SimpleFunctionDescriptor
+        val returnType = descriptor.returnType
         val codeToInline = buildCodeToInline(
             element,
-            descriptor.returnType,
+            returnType,
-            element.hasDeclaredReturnType(),
+            element.hasDeclaredReturnType() || (element.hasBlockBody() && returnType?.isUnit() == true),
             element.bodyExpression!!,
             element.hasBlockBody(),
             editor

idea/src/org/jetbrains/kotlin/idea/refactoring/inline/inlineUtils.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -119,20 +119,17 @@ internal fun buildCodeToInline(
     isBlockBody: Boolean,
     editor: Editor?
 ): CodeToInline? {
-    val bodyCopy = bodyOrInitializer.copied()
-
-    val expectedType = if (!isBlockBody && isReturnTypeExplicit)
-        returnType ?: TypeUtils.NO_EXPECTED_TYPE
-    else
-        TypeUtils.NO_EXPECTED_TYPE
-
     val scope by lazy { bodyOrInitializer.getResolutionScope() }
     fun analyzeExpressionInContext(expression: KtExpression): BindingContext = expression.analyzeInContext(
         scope = scope,
         contextExpression = bodyOrInitializer,
-        expectedType = expectedType
+        expectedType = if (isReturnTypeExplicit && (!isBlockBody || expression.parent is KtReturnExpression))
+            returnType ?: TypeUtils.NO_EXPECTED_TYPE
+        else
+            TypeUtils.NO_EXPECTED_TYPE
     )
 
+    val bodyCopy = bodyOrInitializer.copied()
     val descriptor = declaration.unsafeResolveToDescriptor()
     val builder = CodeToInlineBuilder(descriptor as CallableDescriptor, declaration.getResolutionFacade())
     if (isBlockBody) {
@@ -166,7 +163,7 @@ internal fun buildCodeToInline(
         }
 
         return builder.prepareCodeToInline(
-            lastReturn?.returnedExpression?.copied(),
+            lastReturn?.returnedExpression,
             statements.dropLast(returnStatements.size), ::analyzeExpressionInContext, reformat = true
         )
     } else {

idea/testData/refactoring/inline/function/returnAtEnd/ImplicitGenericTypeInReturn.kt

@@ -0,0 +1,9 @@
+fun <T> buildList(size: Int) = listOf<T>()
+
+fun f() {
+    buildMyList<caret>(5)
+}
+
+private fun buildMyList(count: Int): List<String> {
+    return buildList(size = count)
+}

idea/testData/refactoring/inline/function/returnAtEnd/ImplicitGenericTypeInReturn.kt.after

@@ -0,0 +1,6 @@
+fun <T> buildList(size: Int) = listOf<T>()
+
+fun f() {
+    buildList<String>(size = 5)
+}
+

idea/testData/refactoring/inline/function/returnAtEnd/ImplicitGenericTypeInReturnWithInvalidStatementBefore.kt

@@ -0,0 +1,10 @@
+fun <T> buildList(size: Int) = listOf<T>()
+
+fun f() {
+    buildMyList<caret>(5)
+}
+
+private fun buildMyList(count: Int): List<String> {
+    buildList(size = count)
+    return buildList(size = count)
+}

idea/testData/refactoring/inline/function/returnAtEnd/ImplicitGenericTypeInReturnWithInvalidStatementBefore.kt.after

@@ -0,0 +1,7 @@
+fun <T> buildList(size: Int) = listOf<T>()
+
+fun f() {
+    buildList(size = 5)
+    buildList<String>(size = 5)
+}
+

idea/testData/refactoring/inline/function/returnAtEnd/ImplicitGenericTypeInReturnWithOtherStatementBefore.kt

@@ -0,0 +1,10 @@
+fun <T> buildList(size: Int) = listOf<T>()
+
+fun f() {
+    buildMyList<caret>(5)
+}
+
+private fun buildMyList(count: Int): List<String> {
+    val b: List<Int> = buildList(size = count)
+    return buildList(size = count)
+}

idea/testData/refactoring/inline/function/returnAtEnd/ImplicitGenericTypeInReturnWithOtherStatementBefore.kt.after

@@ -0,0 +1,7 @@
+fun <T> buildList(size: Int) = listOf<T>()
+
+fun f() {
+    val b: List<Int> = buildList(size = 5)
+    buildList<String>(size = 5)
+}
+

idea/testData/refactoring/inline/function/returnAtEnd/ImplicitUnitGenericTypeInReturn.kt

@@ -0,0 +1,9 @@
+fun <T> buildType(size: Int): T = TODO()
+
+fun f() {
+    buildMyList<caret>(5)
+}
+
+private fun buildMyList(count: Int) {
+    return buildType(size = count)
+}

idea/testData/refactoring/inline/function/returnAtEnd/ImplicitUnitGenericTypeInReturn.kt.after

@@ -0,0 +1,6 @@
+fun <T> buildType(size: Int): T = TODO()
+
+fun f() {
+    buildType<Unit>(size = 5)
+}
+

idea/testData/refactoring/inline/function/returnAtEnd/Kt26705.kt

@@ -0,0 +1,11 @@
+fun <T> buildList(size: Int) = listOf<T>()
+
+class Test {
+    fun f() {
+        val list = buildMyList(5)
+    }
+
+    private fun buildMyList<caret>(count: Int): List<String> {
+        return buildList(size = count)
+    }
+}

idea/testData/refactoring/inline/function/returnAtEnd/Kt26705.kt.after

@@ -0,0 +1,8 @@
+fun <T> buildList(size: Int) = listOf<T>()
+
+class Test {
+    fun f() {
+        val list = buildList<String>(size = 5)
+    }
+
+}

idea/tests/org/jetbrains/kotlin/idea/refactoring/inline/InlineTestGenerated.java

@@ -463,6 +463,26 @@ public class InlineTestGenerated extends AbstractInlineTest {
                 runTest("idea/testData/refactoring/inline/function/returnAtEnd/GenericTypeArgument2.kt");
             }
 
+            @TestMetadata("ImplicitGenericTypeInReturn.kt")
+            public void testImplicitGenericTypeInReturn() throws Exception {
+                runTest("idea/testData/refactoring/inline/function/returnAtEnd/ImplicitGenericTypeInReturn.kt");
+            }
+
+            @TestMetadata("ImplicitGenericTypeInReturnWithInvalidStatementBefore.kt")
+            public void testImplicitGenericTypeInReturnWithInvalidStatementBefore() throws Exception {
+                runTest("idea/testData/refactoring/inline/function/returnAtEnd/ImplicitGenericTypeInReturnWithInvalidStatementBefore.kt");
+            }
+
+            @TestMetadata("ImplicitGenericTypeInReturnWithOtherStatementBefore.kt")
+            public void testImplicitGenericTypeInReturnWithOtherStatementBefore() throws Exception {
+                runTest("idea/testData/refactoring/inline/function/returnAtEnd/ImplicitGenericTypeInReturnWithOtherStatementBefore.kt");
+            }
+
+            @TestMetadata("ImplicitUnitGenericTypeInReturn.kt")
+            public void testImplicitUnitGenericTypeInReturn() throws Exception {
+                runTest("idea/testData/refactoring/inline/function/returnAtEnd/ImplicitUnitGenericTypeInReturn.kt");
+            }
+
             @TestMetadata("InnerFunction.kt")
             public void testInnerFunction() throws Exception {
                 runTest("idea/testData/refactoring/inline/function/returnAtEnd/InnerFunction.kt");
@@ -473,6 +493,11 @@ public class InlineTestGenerated extends AbstractInlineTest {
                 runTest("idea/testData/refactoring/inline/function/returnAtEnd/InnerFunction2.kt");
             }
 
+            @TestMetadata("Kt26705.kt")
+            public void testKt26705() throws Exception {
+                runTest("idea/testData/refactoring/inline/function/returnAtEnd/Kt26705.kt");
+            }
+
             @TestMetadata("MultipleStatements.kt")
             public void testMultipleStatements() throws Exception {
                 runTest("idea/testData/refactoring/inline/function/returnAtEnd/MultipleStatements.kt");