From 8e35b7a29cfeba7a2cf0b4c178400c3c8ccbda98 Mon Sep 17 00:00:00 2001
From: Dmitry Gridin <dmitry.gridin@jetbrains.com>
Date: Wed, 22 Apr 2020 21:15:03 +0700
Subject: [PATCH] KotlinIntroduceVariableHandler: fix for declarations with
 expression body

#KT-38449 Fixed
---
 .../kotlin/idea/codeInliner/ReplacementPerformer.kt   |  4 +++-
 .../idea/intentions/ConvertToBlockBodyIntention.kt    | 11 ++++-------
 .../intentions/UsePropertyAccessSyntaxIntention.kt    |  2 +-
 .../CreateLocalVariableActionFactory.kt               |  2 +-
 .../idea/refactoring/introduce/introduceUtil.kt       |  4 ++--
 .../KotlinIntroduceVariableHandler.kt                 |  8 +++++++-
 .../refactoring/introduceVariable/kt38449_int.kt      |  7 +++++++
 .../introduceVariable/kt38449_int.kt.after            | 10 ++++++++++
 .../refactoring/introduceVariable/kt38449_unit.kt     |  7 +++++++
 .../introduceVariable/kt38449_unit.kt.after           | 10 ++++++++++
 .../introduce/ExtractionTestGenerated.java            | 10 ++++++++++
 11 files changed, 62 insertions(+), 13 deletions(-)
 create mode 100644 idea/testData/refactoring/introduceVariable/kt38449_int.kt
 create mode 100644 idea/testData/refactoring/introduceVariable/kt38449_int.kt.after
 create mode 100644 idea/testData/refactoring/introduceVariable/kt38449_unit.kt
 create mode 100644 idea/testData/refactoring/introduceVariable/kt38449_unit.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/codeInliner/ReplacementPerformer.kt b/idea/src/org/jetbrains/kotlin/idea/codeInliner/ReplacementPerformer.kt
index ce307afdefd..344e1ab61c2 100644
--- a/idea/src/org/jetbrains/kotlin/idea/codeInliner/ReplacementPerformer.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/codeInliner/ReplacementPerformer.kt
@@ -265,7 +265,9 @@ internal class ExpressionReplacementPerformer(
     private fun <TElement : KtElement> withElementToBeReplacedPreserved(action: () -> TElement): TElement {
         elementToBeReplaced.putCopyableUserData(ELEMENT_TO_BE_REPLACED_KEY, Unit)
         val result = action()
-        elementToBeReplaced = result.findDescendantOfType { it.getCopyableUserData(ELEMENT_TO_BE_REPLACED_KEY) != null }!!
+        elementToBeReplaced = result.findDescendantOfType { it.getCopyableUserData(ELEMENT_TO_BE_REPLACED_KEY) != null }
+            ?: error("Element `elementToBeReplaced` not found")
+
         elementToBeReplaced.putCopyableUserData(ELEMENT_TO_BE_REPLACED_KEY, null)
         return result
     }
diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToBlockBodyIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToBlockBodyIntention.kt
index 40f899af88f..466970db89e 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToBlockBodyIntention.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToBlockBodyIntention.kt
@@ -44,11 +44,11 @@ class ConvertToBlockBodyIntention : SelfTargetingIntention<KtDeclarationWithBody
     override fun allowCaretInsideElement(element: PsiElement) = element !is KtDeclaration && super.allowCaretInsideElement(element)
 
     override fun applyTo(element: KtDeclarationWithBody, editor: Editor?) {
-        convert(element)
+        convert(element, true)
     }
 
     companion object {
-        fun convert(declaration: KtDeclarationWithBody): KtDeclarationWithBody {
+        fun convert(declaration: KtDeclarationWithBody, withReformat: Boolean = false): KtDeclarationWithBody {
             val body = declaration.bodyExpression!!
 
             fun generateBody(returnsValue: Boolean): KtExpression {
@@ -86,13 +86,10 @@ class ConvertToBlockBodyIntention : SelfTargetingIntention<KtDeclarationWithBody
 
             declaration.equalsToken!!.delete()
             val replaced = body.replace(newBody)
-            declaration.containingKtFile.adjustLineIndent(replaced.startOffset, replaced.endOffset)
+            if (withReformat) declaration.containingKtFile.adjustLineIndent(replaced.startOffset, replaced.endOffset)
             return declaration
         }
 
-        private fun KtNamedFunction.returnType(): KotlinType? {
-            val descriptor = resolveToDescriptorIfAny() ?: return null
-            return descriptor.returnType
-        }
+        private fun KtNamedFunction.returnType(): KotlinType? = resolveToDescriptorIfAny()?.returnType
     }
 }
diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/UsePropertyAccessSyntaxIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/UsePropertyAccessSyntaxIntention.kt
index a9f452959df..9d89731e1ed 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/UsePropertyAccessSyntaxIntention.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/UsePropertyAccessSyntaxIntention.kt
@@ -273,7 +273,7 @@ class UsePropertyAccessSyntaxIntention : SelfTargetingOffsetIndependentIntention
         val callParent = call.parent
         var callToConvert = callExpression
         if (callParent is KtDeclarationWithBody && call == callParent.bodyExpression) {
-            ConvertToBlockBodyIntention.convert(callParent)
+            ConvertToBlockBodyIntention.convert(callParent, true)
             val firstStatement = callParent.bodyBlockExpression?.statements?.first()
             callToConvert = (firstStatement as? KtQualifiedExpression)?.selectorExpression as? KtCallExpression
                 ?: firstStatement as? KtCallExpression
diff --git a/idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/createVariable/CreateLocalVariableActionFactory.kt b/idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/createVariable/CreateLocalVariableActionFactory.kt
index 27a68cc4e16..b5a5164bae3 100644
--- a/idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/createVariable/CreateLocalVariableActionFactory.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/createVariable/CreateLocalVariableActionFactory.kt
@@ -46,7 +46,7 @@ object CreateLocalVariableActionFactory : KotlinSingleIntentionActionFactory() {
 
                 val actualContainer = when (container) {
                     is KtBlockExpression -> container
-                    else -> ConvertToBlockBodyIntention.convert(container as KtDeclarationWithBody).bodyExpression!!
+                    else -> ConvertToBlockBodyIntention.convert(container as KtDeclarationWithBody, true).bodyExpression!!
                 } as KtBlockExpression
 
                 if (actualContainer != container) {
diff --git a/idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceUtil.kt b/idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceUtil.kt
index 16a847d12ad..41b654dd36c 100644
--- a/idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceUtil.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceUtil.kt
@@ -164,8 +164,8 @@ fun PsiElement.findExpressionByCopyableDataAndClearIt(key: Key<Boolean>): KtExpr
     return result
 }
 
-fun PsiElement.findElementByCopyableDataAndClearIt(key: Key<Boolean>): PsiElement {
-    val result = findDescendantOfType<PsiElement> { it.getCopyableUserData(key) != null }!!
+fun PsiElement.findElementByCopyableDataAndClearIt(key: Key<Boolean>): PsiElement? {
+    val result = findDescendantOfType<PsiElement> { it.getCopyableUserData(key) != null } ?: return null
     result.putCopyableUserData(key, null)
     return result
 }
diff --git a/idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceVariable/KotlinIntroduceVariableHandler.kt b/idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceVariable/KotlinIntroduceVariableHandler.kt
index 0414f1c7a92..d61fa48bad1 100644
--- a/idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceVariable/KotlinIntroduceVariableHandler.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceVariable/KotlinIntroduceVariableHandler.kt
@@ -311,7 +311,13 @@ object KotlinIntroduceVariableHandler : RefactoringActionHandler {
                 } ?: newReplace
             }
 
-            runRefactoring(isVar, newExpression ?: return, newCommonContainer, newCommonParent, newAllReplaces)
+            runRefactoring(
+                isVar,
+                newExpression ?: return,
+                newCommonContainer,
+                newCommonParent ?: return,
+                newAllReplaces
+            )
         }
     }
 
diff --git a/idea/testData/refactoring/introduceVariable/kt38449_int.kt b/idea/testData/refactoring/introduceVariable/kt38449_int.kt
new file mode 100644
index 00000000000..78bfb704f19
--- /dev/null
+++ b/idea/testData/refactoring/introduceVariable/kt38449_int.kt
@@ -0,0 +1,7 @@
+fun b(body: () -> Int) = body()
+
+class A {
+    fun test() = b {
+        <selection>24</selection>
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/refactoring/introduceVariable/kt38449_int.kt.after b/idea/testData/refactoring/introduceVariable/kt38449_int.kt.after
new file mode 100644
index 00000000000..56fb261db50
--- /dev/null
+++ b/idea/testData/refactoring/introduceVariable/kt38449_int.kt.after
@@ -0,0 +1,10 @@
+fun b(body: () -> Int) = body()
+
+class A {
+    fun test(): Int {
+        val i = 24
+        return b {
+            i
+        }
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/refactoring/introduceVariable/kt38449_unit.kt b/idea/testData/refactoring/introduceVariable/kt38449_unit.kt
new file mode 100644
index 00000000000..7a36b4d9335
--- /dev/null
+++ b/idea/testData/refactoring/introduceVariable/kt38449_unit.kt
@@ -0,0 +1,7 @@
+fun b(body: () -> Unit): Unit = body()
+
+class A {
+    fun test() = b {
+        <selection>24</selection>
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/refactoring/introduceVariable/kt38449_unit.kt.after b/idea/testData/refactoring/introduceVariable/kt38449_unit.kt.after
new file mode 100644
index 00000000000..beaeac60499
--- /dev/null
+++ b/idea/testData/refactoring/introduceVariable/kt38449_unit.kt.after
@@ -0,0 +1,10 @@
+fun b(body: () -> Unit): Unit = body()
+
+class A {
+    fun test() {
+        val i = 24
+        b {
+            i
+        }
+    }
+}
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/refactoring/introduce/ExtractionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/refactoring/introduce/ExtractionTestGenerated.java
index 36365bb372c..8c0b02dd72b 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/refactoring/introduce/ExtractionTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/refactoring/introduce/ExtractionTestGenerated.java
@@ -220,6 +220,16 @@ public class ExtractionTestGenerated extends AbstractExtractionTest {
             runTest("idea/testData/refactoring/introduceVariable/kt21530_withParam.kt");
         }
 
+        @TestMetadata("kt38449_int.kt")
+        public void testKt38449_int() throws Exception {
+            runTest("idea/testData/refactoring/introduceVariable/kt38449_int.kt");
+        }
+
+        @TestMetadata("kt38449_unit.kt")
+        public void testKt38449_unit() throws Exception {
+            runTest("idea/testData/refactoring/introduceVariable/kt38449_unit.kt");
+        }
+
         @TestMetadata("LoopRange.kt")
         public void testLoopRange() throws Exception {
             runTest("idea/testData/refactoring/introduceVariable/LoopRange.kt");