From a793a63b735c0a18858ee97a8f1296d0b113b335 Mon Sep 17 00:00:00 2001
From: Ilya Gorbunov <ilya.gorbunov@jetbrains.com>
Date: Mon, 9 Nov 2015 23:13:27 +0300
Subject: [PATCH] Simplify message expression from lambda argument when
 converting assert to if.

---
 .../ConvertAssertToIfWithThrowIntention.kt    | 32 +++++++++++++------
 .../assertErrorOverloaded.kt                  |  2 +-
 .../convertAssertToIf/booleanCondition.kt     |  2 +-
 .../booleanConditionSimplified.kt             |  2 +-
 .../booleanConditionSimplified2.kt            |  2 +-
 .../booleanConditionWithVariables.kt          |  2 +-
 .../doNotShortenReferenceInsideMessage.kt     |  2 +-
 .../convertAssertToIf/dotQualifiedCall.kt     |  2 +-
 .../functionCallCondition.kt                  |  2 +-
 .../lambdaMessageInsideParentheses.kt         |  4 ++-
 .../lambdaMessageInsideParentheses.kt.after   |  2 +-
 .../lambdaMessageOutsideParentheses.kt        |  4 ++-
 .../lambdaMessageOutsideParentheses.kt.after  |  4 ++-
 ...aMultiStatementMessageInsideParentheses.kt |  7 ++++
 ...StatementMessageInsideParentheses.kt.after |  9 ++++++
 ...MultiStatementMessageOutsideParentheses.kt |  7 ++++
 ...tatementMessageOutsideParentheses.kt.after |  9 ++++++
 .../parenthesizedCondition.kt                 |  2 +-
 .../convertAssertToIf/simpleConvert.kt        |  2 +-
 .../convertAssertToIf/stringVariable.kt       |  2 +-
 .../intentions/IntentionTestGenerated.java    | 12 +++++++
 21 files changed, 88 insertions(+), 24 deletions(-)
 create mode 100644 idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageInsideParentheses.kt
 create mode 100644 idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageInsideParentheses.kt.after
 create mode 100644 idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageOutsideParentheses.kt
 create mode 100644 idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageOutsideParentheses.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertAssertToIfWithThrowIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertAssertToIfWithThrowIntention.kt
index b8fb94b4679..a81bd3eeb47 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertAssertToIfWithThrowIntention.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertAssertToIfWithThrowIntention.kt
@@ -23,8 +23,12 @@ import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.util.ShortenReferences
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
+import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsExpression
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
+import org.jetbrains.kotlin.utils.addToStdlib.check
 
 public class ConvertAssertToIfWithThrowIntention : SelfTargetingIntention<KtCallExpression>(javaClass(), "Replace 'assert' with 'if' statement"), LowPriorityAction {
     override fun isApplicableTo(element: KtCallExpression, caretOffset: Int): Boolean {
@@ -44,17 +48,21 @@ public class ConvertAssertToIfWithThrowIntention : SelfTargetingIntention<KtCall
     override fun applyTo(element: KtCallExpression, editor: Editor) {
         val args = element.getValueArguments()
         val conditionText = args[0]?.getArgumentExpression()?.getText() ?: return
-        val functionLiteral = element.getFunctionLiteralArguments().singleOrNull()
-        val messageIsFunction = messageIsFunction(element)
-
+        val functionLiteralArgument = element.functionLiteralArguments.singleOrNull()
+        val bindingContext = element.analyze(BodyResolveMode.PARTIAL)
         val psiFactory = KtPsiFactory(element)
 
-        val messageExpr = when {
-            args.size() == 2 -> args[1]?.getArgumentExpression() ?: return
-            functionLiteral != null -> functionLiteral!!
-            else -> psiFactory.createExpression("\"Assertion failed\"")
+        val messageFunctionExpr = when {
+            args.size == 2 -> args[1]?.getArgumentExpression() ?: return
+            functionLiteralArgument != null -> functionLiteralArgument.getFunctionLiteral()
+            else -> null
         }
 
+        val extractedMessageSingleExpr = (messageFunctionExpr as? KtFunctionLiteralExpression)?.let { extractMessageSingleExpression(it, bindingContext) }
+
+        val messageIsFunction = extractedMessageSingleExpr == null && messageIsFunction(element, bindingContext)
+        val messageExpr = extractedMessageSingleExpr ?: messageFunctionExpr ?: psiFactory.createExpression("\"Assertion failed\"")
+
         val ifExpression = replaceWithIfThenThrowExpression(element)
 
         // shorten java.lang.AssertionError
@@ -85,8 +93,14 @@ public class ConvertAssertToIfWithThrowIntention : SelfTargetingIntention<KtCall
         simplifyConditionIfPossible(ifExpression)
     }
 
-    private fun messageIsFunction(callExpr: KtCallExpression): Boolean {
-        val resolvedCall = callExpr.getResolvedCall(callExpr.analyze()) ?: return false
+    private fun extractMessageSingleExpression(functionLiteral: KtFunctionLiteralExpression, bindingContext: BindingContext): KtExpression? {
+        return functionLiteral.bodyExpression?.statements?.singleOrNull()?.let { singleStatement ->
+            singleStatement.check { it.isUsedAsExpression(bindingContext) }
+        }
+    }
+
+    private fun messageIsFunction(callExpr: KtCallExpression, bindingContext: BindingContext): Boolean {
+        val resolvedCall = callExpr.getResolvedCall(bindingContext) ?: return false
         val valParameters = resolvedCall.getResultingDescriptor().getValueParameters()
         return valParameters.size() > 1 && !KotlinBuiltIns.isAny(valParameters[1].type)
     }
diff --git a/idea/testData/intentions/convertAssertToIf/assertErrorOverloaded.kt b/idea/testData/intentions/convertAssertToIf/assertErrorOverloaded.kt
index 017f4c65b5c..26c165825ee 100644
--- a/idea/testData/intentions/convertAssertToIf/assertErrorOverloaded.kt
+++ b/idea/testData/intentions/convertAssertToIf/assertErrorOverloaded.kt
@@ -1,7 +1,7 @@
 package a
 
 fun foo() {
-    <caret>assert(true, "text")
+    <caret>assert(true, { "text" })
 }
 
 class AssertionError
diff --git a/idea/testData/intentions/convertAssertToIf/booleanCondition.kt b/idea/testData/intentions/convertAssertToIf/booleanCondition.kt
index a5e360eb3e3..abee093a6cc 100644
--- a/idea/testData/intentions/convertAssertToIf/booleanCondition.kt
+++ b/idea/testData/intentions/convertAssertToIf/booleanCondition.kt
@@ -1,4 +1,4 @@
 // WITH_RUNTIME
 fun foo() {
-    <caret>assert(true && false, "text")
+    <caret>assert(true && false, { "text" })
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/booleanConditionSimplified.kt b/idea/testData/intentions/convertAssertToIf/booleanConditionSimplified.kt
index f03a44b1399..7cdbfba9a17 100644
--- a/idea/testData/intentions/convertAssertToIf/booleanConditionSimplified.kt
+++ b/idea/testData/intentions/convertAssertToIf/booleanConditionSimplified.kt
@@ -1,4 +1,4 @@
 // WITH_RUNTIME
 fun foo() {
-    <caret>assert(1 > 0, "text")
+    <caret>assert(1 > 0) { "text" }
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/booleanConditionSimplified2.kt b/idea/testData/intentions/convertAssertToIf/booleanConditionSimplified2.kt
index 59135636428..e7b6acde349 100644
--- a/idea/testData/intentions/convertAssertToIf/booleanConditionSimplified2.kt
+++ b/idea/testData/intentions/convertAssertToIf/booleanConditionSimplified2.kt
@@ -1,4 +1,4 @@
 // WITH_RUNTIME
 fun foo() {
-    <caret>assert(0 != 1, "text")
+    <caret>assert(0 != 1) { "text" }
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/booleanConditionWithVariables.kt b/idea/testData/intentions/convertAssertToIf/booleanConditionWithVariables.kt
index a1c1864233c..4e706b88e2f 100644
--- a/idea/testData/intentions/convertAssertToIf/booleanConditionWithVariables.kt
+++ b/idea/testData/intentions/convertAssertToIf/booleanConditionWithVariables.kt
@@ -2,5 +2,5 @@
 fun foo() {
     val x = true
     val y = false
-    <caret>assert(x || y, "text")
+    <caret>assert(x || y) { "text" }
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/doNotShortenReferenceInsideMessage.kt b/idea/testData/intentions/convertAssertToIf/doNotShortenReferenceInsideMessage.kt
index 88fbd9192ac..1067fade341 100644
--- a/idea/testData/intentions/convertAssertToIf/doNotShortenReferenceInsideMessage.kt
+++ b/idea/testData/intentions/convertAssertToIf/doNotShortenReferenceInsideMessage.kt
@@ -1,5 +1,5 @@
 fun main(args: Array<String>) {
-    asse<caret>rt(false, "mess" as kotlin.String)
+    asse<caret>rt(false) { "mess" as kotlin.String }
 }
 
 // WITH_RUNTIME
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/dotQualifiedCall.kt b/idea/testData/intentions/convertAssertToIf/dotQualifiedCall.kt
index 1dbc43f76d1..0f7f5ecad34 100644
--- a/idea/testData/intentions/convertAssertToIf/dotQualifiedCall.kt
+++ b/idea/testData/intentions/convertAssertToIf/dotQualifiedCall.kt
@@ -1,4 +1,4 @@
 // WITH_RUNTIME
 fun foo() {
-    kotlin.<caret>assert(true, "text")
+    kotlin.<caret>assert(true) {"text"}
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/functionCallCondition.kt b/idea/testData/intentions/convertAssertToIf/functionCallCondition.kt
index f2fc11fbb32..702dd734733 100644
--- a/idea/testData/intentions/convertAssertToIf/functionCallCondition.kt
+++ b/idea/testData/intentions/convertAssertToIf/functionCallCondition.kt
@@ -1,6 +1,6 @@
 // WITH_RUNTIME
 fun foo() {
-    <caret>assert(bar(), "text")
+    <caret>assert(bar()) { "text" }
 }
 
 fun bar(): Boolean = true
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/lambdaMessageInsideParentheses.kt b/idea/testData/intentions/convertAssertToIf/lambdaMessageInsideParentheses.kt
index eb6f3011d9a..31ac7316873 100644
--- a/idea/testData/intentions/convertAssertToIf/lambdaMessageInsideParentheses.kt
+++ b/idea/testData/intentions/convertAssertToIf/lambdaMessageInsideParentheses.kt
@@ -1,4 +1,6 @@
 // WITH_RUNTIME
 fun foo() {
-    <caret>assert(true, { "text" })
+    <caret>assert(true, {
+        if (true) "text" else return
+    })
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/lambdaMessageInsideParentheses.kt.after b/idea/testData/intentions/convertAssertToIf/lambdaMessageInsideParentheses.kt.after
index e6ab611f942..42022406dfe 100644
--- a/idea/testData/intentions/convertAssertToIf/lambdaMessageInsideParentheses.kt.after
+++ b/idea/testData/intentions/convertAssertToIf/lambdaMessageInsideParentheses.kt.after
@@ -1,6 +1,6 @@
 // WITH_RUNTIME
 fun foo() {
     if (!true) {
-        throw AssertionError({ "text" }())
+        throw AssertionError(if (true) "text" else return)
     }
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/lambdaMessageOutsideParentheses.kt b/idea/testData/intentions/convertAssertToIf/lambdaMessageOutsideParentheses.kt
index 0939757b318..edac359413b 100644
--- a/idea/testData/intentions/convertAssertToIf/lambdaMessageOutsideParentheses.kt
+++ b/idea/testData/intentions/convertAssertToIf/lambdaMessageOutsideParentheses.kt
@@ -1,4 +1,6 @@
 // WITH_RUNTIME
 fun foo() {
-    <caret>assert(true) { "text" }
+    <caret>assert(true) {
+        return
+    }
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/lambdaMessageOutsideParentheses.kt.after b/idea/testData/intentions/convertAssertToIf/lambdaMessageOutsideParentheses.kt.after
index e6ab611f942..861c3701762 100644
--- a/idea/testData/intentions/convertAssertToIf/lambdaMessageOutsideParentheses.kt.after
+++ b/idea/testData/intentions/convertAssertToIf/lambdaMessageOutsideParentheses.kt.after
@@ -1,6 +1,8 @@
 // WITH_RUNTIME
 fun foo() {
     if (!true) {
-        throw AssertionError({ "text" }())
+        throw AssertionError({
+            return
+        }())
     }
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageInsideParentheses.kt b/idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageInsideParentheses.kt
new file mode 100644
index 00000000000..0bf302eef5c
--- /dev/null
+++ b/idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageInsideParentheses.kt
@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+fun foo() {
+    <caret>assert(true, {
+        val value = 1
+        "text and $value"
+    })
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageInsideParentheses.kt.after b/idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageInsideParentheses.kt.after
new file mode 100644
index 00000000000..fdeac466960
--- /dev/null
+++ b/idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageInsideParentheses.kt.after
@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+fun foo() {
+    if (!true) {
+        throw AssertionError({
+            val value = 1
+            "text and $value"
+        }())
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageOutsideParentheses.kt b/idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageOutsideParentheses.kt
new file mode 100644
index 00000000000..dfddb0b1a8e
--- /dev/null
+++ b/idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageOutsideParentheses.kt
@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+fun foo() {
+    <caret>assert(true) {
+        if (false) return
+        "text"
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageOutsideParentheses.kt.after b/idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageOutsideParentheses.kt.after
new file mode 100644
index 00000000000..33c4a0667a6
--- /dev/null
+++ b/idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageOutsideParentheses.kt.after
@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+fun foo() {
+    if (!true) {
+        throw AssertionError({
+            if (false) return
+            "text"
+        }())
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/parenthesizedCondition.kt b/idea/testData/intentions/convertAssertToIf/parenthesizedCondition.kt
index 2fe0413f4c0..d61786133b2 100644
--- a/idea/testData/intentions/convertAssertToIf/parenthesizedCondition.kt
+++ b/idea/testData/intentions/convertAssertToIf/parenthesizedCondition.kt
@@ -1,4 +1,4 @@
 // WITH_RUNTIME
 fun foo() {
-    <caret>assert((true && false), "text")
+    <caret>assert((true && false)) { "text" }
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/simpleConvert.kt b/idea/testData/intentions/convertAssertToIf/simpleConvert.kt
index cab30283641..0939757b318 100644
--- a/idea/testData/intentions/convertAssertToIf/simpleConvert.kt
+++ b/idea/testData/intentions/convertAssertToIf/simpleConvert.kt
@@ -1,4 +1,4 @@
 // WITH_RUNTIME
 fun foo() {
-    <caret>assert(true, "text")
+    <caret>assert(true) { "text" }
 }
\ No newline at end of file
diff --git a/idea/testData/intentions/convertAssertToIf/stringVariable.kt b/idea/testData/intentions/convertAssertToIf/stringVariable.kt
index 13aae70c434..2b83ee77e35 100644
--- a/idea/testData/intentions/convertAssertToIf/stringVariable.kt
+++ b/idea/testData/intentions/convertAssertToIf/stringVariable.kt
@@ -1,5 +1,5 @@
 // WITH_RUNTIME
 fun foo() {
     val f = "text"
-    <caret>assert(true, f)
+    <caret>assert(true) { f }
 }
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
index c0570ce4cf7..2c486976ba2 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
@@ -3088,6 +3088,18 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("lambdaMultiStatementMessageInsideParentheses.kt")
+        public void testLambdaMultiStatementMessageInsideParentheses() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageInsideParentheses.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("lambdaMultiStatementMessageOutsideParentheses.kt")
+        public void testLambdaMultiStatementMessageOutsideParentheses() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageOutsideParentheses.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("lambdaVariable.kt")
         public void testLambdaVariable() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertAssertToIf/lambdaVariable.kt");