From b6803af74612c8cf01745fdb698228bec17ecbad Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <Mikhail.Glukhikh@jetbrains.com>
Date: Thu, 23 Mar 2017 18:57:19 +0300
Subject: [PATCH] KT-6159: generate Unit while code inlining when needed

---
 .../kotlin/idea/codeInliner/ReplacementPerformer.kt   | 11 ++++++-----
 .../inline/function/EmptyFunction.kt.after            |  4 ++--
 .../inline/function/UnitReturnType.kt.after           |  2 ++
 .../inline/function/UnitReturnType2.kt.after          |  3 ++-
 4 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/codeInliner/ReplacementPerformer.kt b/idea/src/org/jetbrains/kotlin/idea/codeInliner/ReplacementPerformer.kt
index 22c6b9476c2..b515eebc179 100644
--- a/idea/src/org/jetbrains/kotlin/idea/codeInliner/ReplacementPerformer.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/codeInliner/ReplacementPerformer.kt
@@ -25,7 +25,6 @@ import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
 import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsExpression
-import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import java.util.*
 
 internal abstract class ReplacementPerformer<TElement : KtElement>(
@@ -85,14 +84,16 @@ internal class ExpressionReplacementPerformer(
             elementToBeReplaced.replace(codeToInline.mainExpression!!)
         }
         else {
-            val bindingContext = elementToBeReplaced.analyze(BodyResolveMode.FULL)
-            val canDropElementToBeReplaced = !elementToBeReplaced.isUsedAsExpression(bindingContext)
+            // NB: Unit is never used as expression
+            val stub = elementToBeReplaced.replace(psiFactory.createExpression("0")) as KtExpression
+            val bindingContext = stub.analyze()
+            val canDropElementToBeReplaced = !stub.isUsedAsExpression(bindingContext)
             if (canDropElementToBeReplaced) {
-                elementToBeReplaced.delete()
+                stub.delete()
                 null
             }
             else {
-                elementToBeReplaced.replace(psiFactory.createExpression("Unit"))
+                stub.replace(psiFactory.createExpression("Unit"))
             }
         }
 
diff --git a/idea/testData/refactoring/inline/function/EmptyFunction.kt.after b/idea/testData/refactoring/inline/function/EmptyFunction.kt.after
index f56a3ed7bb9..c14e42f7ef1 100644
--- a/idea/testData/refactoring/inline/function/EmptyFunction.kt.after
+++ b/idea/testData/refactoring/inline/function/EmptyFunction.kt.after
@@ -2,9 +2,9 @@ fun <T> doIt(p: () -> T): T = TODO()
 
 fun g(p: String?) {
 
-    p?.let { }
+    p?.let { Unit }
 }
 
 fun h() = Unit
 
-fun x() = doIt { }
+fun x() = doIt { Unit }
diff --git a/idea/testData/refactoring/inline/function/UnitReturnType.kt.after b/idea/testData/refactoring/inline/function/UnitReturnType.kt.after
index 2dfff0d52b7..cba7e7e8117 100644
--- a/idea/testData/refactoring/inline/function/UnitReturnType.kt.after
+++ b/idea/testData/refactoring/inline/function/UnitReturnType.kt.after
@@ -9,6 +9,7 @@ class C {
         p?.let {
             println(3)
             println(4)
+            Unit
         }
 
         if (other != null) {
@@ -25,5 +26,6 @@ class C {
     fun x() = doIt {
         println(9)
         println(10)
+        Unit
     }
 }
diff --git a/idea/testData/refactoring/inline/function/UnitReturnType2.kt.after b/idea/testData/refactoring/inline/function/UnitReturnType2.kt.after
index 841bdb85d56..5fcded0ffa7 100644
--- a/idea/testData/refactoring/inline/function/UnitReturnType2.kt.after
+++ b/idea/testData/refactoring/inline/function/UnitReturnType2.kt.after
@@ -9,6 +9,7 @@ fun g(p: String?) {
     p?.let {
         println(3)
         nonUnit(4)
+        Unit
     }
 }
 
@@ -20,5 +21,5 @@ fun h() {
 fun x() = doIt {
     println(7)
     nonUnit(8)
-    // Unit should be here
+    Unit
 }