From 08c6a21ac15ad9d37d96e89a48ed29fd66cfbc88 Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Sun, 23 Oct 2016 01:52:45 +0300
Subject: [PATCH] Handling safe calls of multi-statement function

---
 .../kotlin/idea/replacement/introduceValue.kt | 30 +++++++++++++------
 .../inline/function/returnAtEnd/SafeCall.kt   | 11 +++++++
 .../function/returnAtEnd/SafeCall.kt.after    | 19 ++++++++++++
 .../inline/InlineTestGenerated.java           |  6 ++++
 4 files changed, 57 insertions(+), 9 deletions(-)
 create mode 100644 idea/testData/refactoring/inline/function/returnAtEnd/SafeCall.kt
 create mode 100644 idea/testData/refactoring/inline/function/returnAtEnd/SafeCall.kt.after

diff --git a/idea/idea-replacement-engine/src/org/jetbrains/kotlin/idea/replacement/introduceValue.kt b/idea/idea-replacement-engine/src/org/jetbrains/kotlin/idea/replacement/introduceValue.kt
index 73300fb7368..2000f48671a 100644
--- a/idea/idea-replacement-engine/src/org/jetbrains/kotlin/idea/replacement/introduceValue.kt
+++ b/idea/idea-replacement-engine/src/org/jetbrains/kotlin/idea/replacement/introduceValue.kt
@@ -101,17 +101,29 @@ internal fun MutableReplacementCode.introduceValue(
         }
     }
     else {
-        //TODO: handle mainExpression == null and statementsBefore!
+        val useIt = !isNameUsed("it")
+        val name = if (useIt) Name.identifier("it") else suggestName { !isNameUsed(it) }
+        replaceUsages(name)
 
-        mainExpression = if (!isNameUsed("it")) {
-            replaceUsages(Name.identifier("it"))
-            psiFactory.createExpressionByPattern("$0?.let { $1 }", value, mainExpression!!)
-        }
-        else {
-            val name = suggestName { !isNameUsed(it) }
-            replaceUsages(name)
-            psiFactory.createExpressionByPattern("$0?.let { $1 -> $2 }", value, name, mainExpression!!)
+        mainExpression = psiFactory.buildExpression {
+            appendExpression(value)
+            appendFixedText("?.let{")
+
+            if (!useIt) {
+                appendName(name)
+                appendFixedText("->")
+            }
+
+            for (statement in statementsBefore) {
+                appendExpression(statement)
+                appendFixedText("\n")
+            }
+
+            appendExpression(mainExpression!!) //TODO: mainExpression == null
+
+            appendFixedText("}")
         }
+        statementsBefore.clear() //TODO: postInsertActions!
     }
 }
 
diff --git a/idea/testData/refactoring/inline/function/returnAtEnd/SafeCall.kt b/idea/testData/refactoring/inline/function/returnAtEnd/SafeCall.kt
new file mode 100644
index 00000000000..66404cefa5d
--- /dev/null
+++ b/idea/testData/refactoring/inline/function/returnAtEnd/SafeCall.kt
@@ -0,0 +1,11 @@
+fun String.<caret>f(p: Int): Int {
+    println(p)
+    return hashCode() * p
+}
+
+fun f(s: String?) {
+    s?.f(1)
+    s?.substring(1)?.f(2)
+    val s1 = s?.f(3)
+    val s2 = s?.substring(1)?.f(4)
+}
\ No newline at end of file
diff --git a/idea/testData/refactoring/inline/function/returnAtEnd/SafeCall.kt.after b/idea/testData/refactoring/inline/function/returnAtEnd/SafeCall.kt.after
new file mode 100644
index 00000000000..51ac483f23c
--- /dev/null
+++ b/idea/testData/refactoring/inline/function/returnAtEnd/SafeCall.kt.after
@@ -0,0 +1,19 @@
+fun f(s: String?) {
+    if (s != null) {
+        println(1)
+        s.hashCode() * 1
+    }
+    val substring = s?.substring(1)
+    if (substring != null) {
+        println(2)
+        substring.hashCode() * 2
+    }
+    val s1 = s?.let {
+        println(3)
+        it.hashCode() * 3
+    }
+    val s2 = s?.substring(1)?.let {
+        println(4)
+        it.hashCode() * 4
+    }
+}
diff --git a/idea/tests/org/jetbrains/kotlin/idea/refactoring/inline/InlineTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/refactoring/inline/InlineTestGenerated.java
index be06e2481c9..a335446c4e7 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/refactoring/inline/InlineTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/refactoring/inline/InlineTestGenerated.java
@@ -138,6 +138,12 @@ public class InlineTestGenerated extends AbstractInlineTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("SafeCall.kt")
+            public void testSafeCall() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/inline/function/returnAtEnd/SafeCall.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("SingleStatement.kt")
             public void testSingleStatement() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/inline/function/returnAtEnd/SingleStatement.kt");