Simplify message expression from lambda argument when converting assert to if.

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)
     }

idea/testData/intentions/convertAssertToIf/assertErrorOverloaded.kt

@@ -1,7 +1,7 @@
 package a
 
 fun foo() {
-    <caret>assert(true, "text")
+    <caret>assert(true, { "text" })
 }
 
 class AssertionError

idea/testData/intentions/convertAssertToIf/booleanCondition.kt

@@ -1,4 +1,4 @@
 // WITH_RUNTIME
 fun foo() {
-    <caret>assert(true && false, "text")
+    <caret>assert(true && false, { "text" })
 }

idea/testData/intentions/convertAssertToIf/booleanConditionSimplified.kt

@@ -1,4 +1,4 @@
 // WITH_RUNTIME
 fun foo() {
-    <caret>assert(1 > 0, "text")
+    <caret>assert(1 > 0) { "text" }
 }

idea/testData/intentions/convertAssertToIf/booleanConditionSimplified2.kt

@@ -1,4 +1,4 @@
 // WITH_RUNTIME
 fun foo() {
-    <caret>assert(0 != 1, "text")
+    <caret>assert(0 != 1) { "text" }
 }

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" }
 }

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

idea/testData/intentions/convertAssertToIf/dotQualifiedCall.kt

@@ -1,4 +1,4 @@
 // WITH_RUNTIME
 fun foo() {
-    kotlin.<caret>assert(true, "text")
+    kotlin.<caret>assert(true) {"text"}
 }

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

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
+    })
 }

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)
     }
 }

idea/testData/intentions/convertAssertToIf/lambdaMessageOutsideParentheses.kt

@@ -1,4 +1,6 @@
 // WITH_RUNTIME
 fun foo() {
-    <caret>assert(true) { "text" }
+    <caret>assert(true) {
+        return
+    }
 }

idea/testData/intentions/convertAssertToIf/lambdaMessageOutsideParentheses.kt.after

@@ -1,6 +1,8 @@
 // WITH_RUNTIME
 fun foo() {
     if (!true) {
-        throw AssertionError({ "text" }())
+        throw AssertionError({
+            return
+        }())
     }
 }

idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageInsideParentheses.kt

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+fun foo() {
+    <caret>assert(true, {
+        val value = 1
+        "text and $value"
+    })
+}

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"
+        }())
+    }
+}

idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageOutsideParentheses.kt

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+fun foo() {
+    <caret>assert(true) {
+        if (false) return
+        "text"
+    }
+}

idea/testData/intentions/convertAssertToIf/lambdaMultiStatementMessageOutsideParentheses.kt.after

@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+fun foo() {
+    if (!true) {
+        throw AssertionError({
+            if (false) return
+            "text"
+        }())
+    }
+}

idea/testData/intentions/convertAssertToIf/parenthesizedCondition.kt

@@ -1,4 +1,4 @@
 // WITH_RUNTIME
 fun foo() {
-    <caret>assert((true && false), "text")
+    <caret>assert((true && false)) { "text" }
 }

idea/testData/intentions/convertAssertToIf/simpleConvert.kt

@@ -1,4 +1,4 @@
 // WITH_RUNTIME
 fun foo() {
-    <caret>assert(true, "text")
+    <caret>assert(true) { "text" }
 }

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 }
 }

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");