Support message lambda as a local variable

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/PowerAssertCallTransformer.kt

@@ -40,11 +40,13 @@ import org.jetbrains.kotlin.ir.builders.irString
 import org.jetbrains.kotlin.ir.builders.parent
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrValueParameter
 import org.jetbrains.kotlin.ir.declarations.path
 import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrConst
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrFunctionExpression
+import org.jetbrains.kotlin.ir.expressions.IrGetValue
 import org.jetbrains.kotlin.ir.expressions.IrStringConcatenation
 import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
 import org.jetbrains.kotlin.ir.symbols.IrTypeParameterSymbol
@@ -137,7 +139,8 @@ class PowerAssertCallTransformer(
     variables: List<IrTemporaryVariable> = listOf()
   ): IrExpression {
     if (index >= roots.size) {
-      val prefix = buildMessagePrefix(messageArgument, roots, original)?.deepCopyWithSymbols(parent)
+      val prefix = buildMessagePrefix(messageArgument, delegate.messageParameter, roots, original)
+        ?.deepCopyWithSymbols(parent)
       val diagram = irDiagramString(file, fileSource, prefix, original, variables)
       return delegate.buildCall(this, original, arguments, diagram)
     } else {
@@ -156,20 +159,30 @@ class PowerAssertCallTransformer(
 
   private fun DeclarationIrBuilder.buildMessagePrefix(
     messageArgument: IrExpression?,
+    messageParameter: IrValueParameter,
     roots: List<Node?>,
     original: IrCall
   ): IrExpression? {
     return when {
       messageArgument is IrConst<*> -> messageArgument
       messageArgument is IrStringConcatenation -> messageArgument
-      messageArgument is IrFunctionExpression -> {
+      messageArgument is IrGetValue -> {
-        val invoke = messageArgument.type.classOrNull!!.functions
+        if (messageArgument.type.isAssignableTo(context.irBuiltIns.stringType)) {
-          .filter { it.owner.name.toString() == "invoke" }
+          return messageArgument
+        } else {
+          val invoke = messageParameter.type.classOrNull!!.functions
+            .filter { !it.owner.isFakeOverride } // TODO best way to find single access method?
+            .single()
+          irCall(invoke).apply { dispatchReceiver = messageArgument }
+        }
+      }
+      // Kotlin Lambda or SAMs conversion lambda
+      messageArgument is IrFunctionExpression || messageArgument is IrTypeOperatorCall -> {
+        val invoke = messageParameter.type.classOrNull!!.functions
+          .filter { !it.owner.isFakeOverride } // TODO best way to find single access method?
           .single()
         irCall(invoke).apply { dispatchReceiver = messageArgument }
       }
-      // SAMs conversion can be unwrapped to IrFunctionExpression
-      messageArgument is IrTypeOperatorCall -> buildMessagePrefix(messageArgument.argument, roots, original)
       // TODO what should the default message be?
       roots.size == 1 && original.getValueArgument(0)!!.type.isBoolean() -> irString("Assertion failed")
       else -> null
@@ -199,7 +212,7 @@ class PowerAssertCallTransformer(
 
         val messageParameter = parameters.last()
         return@mapNotNull when {
-          isStringSupertype(messageParameter.type) -> SimpleFunctionDelegate(overload)
+          isStringSupertype(messageParameter.type) -> SimpleFunctionDelegate(overload, messageParameter)
           isStringFunction(messageParameter.type) -> LambdaFunctionDelegate(overload, messageParameter)
           isStringJavaSupplierFunction(messageParameter.type) ->
             SamConversionLambdaFunctionDelegate(overload, messageParameter)

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/FunctionDelegate.kt

@@ -20,6 +20,8 @@ import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
 import org.jetbrains.kotlin.ir.builders.irCall
 import org.jetbrains.kotlin.ir.builders.parent
 import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrValueParameter
+import org.jetbrains.kotlin.ir.declarations.IrVariable
 import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
@@ -27,6 +29,7 @@ import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
 
 interface FunctionDelegate {
   val function: IrFunction
+  val messageParameter: IrValueParameter
 
   fun buildCall(
     builder: IrBuilderWithScope,

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/LambdaFunctionDelegate.kt

@@ -26,7 +26,7 @@ import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 
 class LambdaFunctionDelegate(
   private val overload: IrSimpleFunctionSymbol,
-  private val messageParameter: IrValueParameter
+  override val messageParameter: IrValueParameter,
 ) : FunctionDelegate {
   override val function = overload.owner
 

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/SamConversionLambdaFunctionDelegate.kt

@@ -28,7 +28,7 @@ import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 
 class SamConversionLambdaFunctionDelegate(
   private val overload: IrSimpleFunctionSymbol,
-  private val messageParameter: IrValueParameter
+  override val messageParameter: IrValueParameter,
 ) : FunctionDelegate {
   override val function = overload.owner
 

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/SimpleFunctionDelegate.kt

@@ -17,12 +17,14 @@
 package com.bnorm.power.delegate
 
 import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
+import org.jetbrains.kotlin.ir.declarations.IrValueParameter
 import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 
 class SimpleFunctionDelegate(
-  private val overload: IrSimpleFunctionSymbol
+  private val overload: IrSimpleFunctionSymbol,
+  override val messageParameter: IrValueParameter,
 ) : FunctionDelegate {
   override val function = overload.owner
 

kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/AssertTest.kt

@@ -77,6 +77,23 @@ assert(1 == 2) { "Not equal" }
     )
   }
 
+  @Test
+  fun customLocalVariableMessage() {
+    assertMessage(
+      """
+fun main() {
+  val lambda = { "Not equal" }
+  assert(1 == 2, lambda) 
+}""",
+      """
+Not equal
+assert(1 == 2, lambda)
+         |
+         false
+""".trimIndent()
+    )
+  }
+
   @Test
   fun booleanExpressionsShortCircuit() {
     assertMessage(

kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/Junit5AssertionsTest.kt

@@ -58,6 +58,26 @@ class Junit5AssertionsTest {
     )
   }
 
+  @Test
+  fun `test JUnit5 Assertions#assertTrue transformation with local variable message`() {
+    assertMessage(
+      """
+      import org.junit.jupiter.api.Assertions.assertTrue
+      
+      fun main() {
+        val message = "Message:"
+        assertTrue(1 != 1, message)
+      }""",
+      """
+      Message:
+      assertTrue(1 != 1, message)
+                   |
+                   false ==> expected: <true> but was: <false>
+      """.trimIndent(),
+      PowerAssertComponentRegistrar(setOf(FqName("org.junit.jupiter.api.Assertions.assertTrue")))
+    )
+  }
+
   @Test
   fun `test JUnit5 Assertions#assertTrue transformation with message supplier`() {
     assertMessage(
@@ -77,6 +97,27 @@ class Junit5AssertionsTest {
     )
   }
 
+  @Test
+  fun `test JUnit5 Assertions#assertTrue transformation with local variable message supplier`() {
+    assertMessage(
+      """
+      import java.util.function.Supplier
+      import org.junit.jupiter.api.Assertions.assertTrue
+      
+      fun main() {
+        val supplier = Supplier { "Message:" }
+        assertTrue(1 != 1, supplier)
+      }""",
+      """
+      Message:
+      assertTrue(1 != 1, supplier)
+                   |
+                   false ==> expected: <true> but was: <false>
+      """.trimIndent(),
+      PowerAssertComponentRegistrar(setOf(FqName("org.junit.jupiter.api.Assertions.assertTrue")))
+    )
+  }
+
   @Test
   fun `test JUnit5 Assertions#assertFalse transformation`() {
     assertMessage(
@@ -115,6 +156,26 @@ class Junit5AssertionsTest {
     )
   }
 
+  @Test
+  fun `test JUnit5 Assertions#assertFalse transformation with local variable message`() {
+    assertMessage(
+      """
+      import org.junit.jupiter.api.Assertions.assertFalse
+      
+      fun main() {
+        val message = "Message:"
+        assertFalse(1 == 1, message)
+      }""",
+      """
+      Message:
+      assertFalse(1 == 1, message)
+                    |
+                    true ==> expected: <false> but was: <true>
+      """.trimIndent(),
+      PowerAssertComponentRegistrar(setOf(FqName("org.junit.jupiter.api.Assertions.assertFalse")))
+    )
+  }
+
   @Test
   fun `test JUnit5 Assertions#assertFalse transformation with message supplier`() {
     assertMessage(
@@ -134,6 +195,27 @@ class Junit5AssertionsTest {
     )
   }
 
+  @Test
+  fun `test JUnit5 Assertions#assertFalse transformation with local variable message supplier`() {
+    assertMessage(
+      """
+      import java.util.function.Supplier
+      import org.junit.jupiter.api.Assertions.assertFalse
+      
+      fun main() {
+        val supplier = Supplier { "Message:" }
+        assertFalse(1 == 1, supplier)
+      }""",
+      """
+      Message:
+      assertFalse(1 == 1, supplier)
+                    |
+                    true ==> expected: <false> but was: <true>
+      """.trimIndent(),
+      PowerAssertComponentRegistrar(setOf(FqName("org.junit.jupiter.api.Assertions.assertFalse")))
+    )
+  }
+
   @Test
   fun `test JUnit5 Assertions#assertEquals transformation`() {
     assertMessage(
@@ -177,6 +259,29 @@ class Junit5AssertionsTest {
     )
   }
 
+  @Test
+  fun `test JUnit5 Assertions#assertEquals transformation with local variable message`() {
+    assertMessage(
+      """
+      import org.junit.jupiter.api.Assertions.assertEquals
+      
+      fun main() {
+        val greeting = "Hello"
+        val name = "World"
+        val message = "Message:"
+        assertEquals(greeting, name, message)
+      }""",
+      """
+      Message:
+      assertEquals(greeting, name, message)
+                   |         |
+                   |         World
+                   Hello ==> expected: <Hello> but was: <World>
+      """.trimIndent(),
+      PowerAssertComponentRegistrar(setOf(FqName("org.junit.jupiter.api.Assertions.assertEquals")))
+    )
+  }
+
   @Test
   fun `test JUnit5 Assertions#assertEquals transformation with message supplier`() {
     assertMessage(
@@ -198,4 +303,28 @@ class Junit5AssertionsTest {
       PowerAssertComponentRegistrar(setOf(FqName("org.junit.jupiter.api.Assertions.assertEquals")))
     )
   }
+
+  @Test
+  fun `test JUnit5 Assertions#assertEquals transformation with local variable message supplier`() {
+    assertMessage(
+      """
+      import java.util.function.Supplier
+      import org.junit.jupiter.api.Assertions.assertEquals
+      
+      fun main() {
+        val greeting = "Hello"
+        val name = "World"
+        val supplier = Supplier { "Message:" }
+        assertEquals(greeting, name, supplier)
+      }""",
+      """
+      Message:
+      assertEquals(greeting, name, supplier)
+                   |         |
+                   |         World
+                   Hello ==> expected: <Hello> but was: <World>
+      """.trimIndent(),
+      PowerAssertComponentRegistrar(setOf(FqName("org.junit.jupiter.api.Assertions.assertEquals")))
+    )
+  }
 }

kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/KotlinTestAssertTest.kt

@@ -58,6 +58,26 @@ class KotlinTestAssertTest {
     )
   }
 
+  @Test
+  fun `test assertTrue transformation with local variable message`() {
+    assertMessage(
+      """
+      import kotlin.test.assertTrue
+      
+      fun main() {
+        val message = "Message:"
+        assertTrue(1 != 1, message)
+      }""",
+      """
+      Message:
+      assertTrue(1 != 1, message)
+                   |
+                   false
+      """.trimIndent(),
+      PowerAssertComponentRegistrar(setOf(FqName("kotlin.test.assertTrue")))
+    )
+  }
+
   @Test
   fun `test assertFalse transformation`() {
     assertMessage(
@@ -96,6 +116,26 @@ class KotlinTestAssertTest {
     )
   }
 
+  @Test
+  fun `test assertFalse transformation with local variable message`() {
+    assertMessage(
+      """
+      import kotlin.test.assertFalse
+      
+      fun main() {
+        val message = "Message:"
+        assertFalse(1 == 1, message)
+      }""",
+      """
+      Message:
+      assertFalse(1 == 1, message)
+                    |
+                    true
+      """.trimIndent(),
+      PowerAssertComponentRegistrar(setOf(FqName("kotlin.test.assertFalse")))
+    )
+  }
+
   @Test
   fun `test assertEquals transformation`() {
     assertMessage(
@@ -139,6 +179,29 @@ class KotlinTestAssertTest {
     )
   }
 
+  @Test
+  fun `test assertEquals transformation with local variable message`() {
+    assertMessage(
+      """
+      import kotlin.test.assertEquals
+      
+      fun main() {
+        val greeting = "Hello"
+        val name = "World"
+        val message = "Message:"
+        assertEquals(greeting, name, message)
+      }""",
+      """
+      Message:
+      assertEquals(greeting, name, message)
+                   |         |
+                   |         World
+                   Hello ==> expected: <Hello> but was: <World>
+      """.trimIndent(),
+      PowerAssertComponentRegistrar(setOf(FqName("kotlin.test.assertEquals")))
+    )
+  }
+
   @Test
   fun `test assertNotNull transformation`() {
     assertMessage(
@@ -177,4 +240,25 @@ class KotlinTestAssertTest {
       PowerAssertComponentRegistrar(setOf(FqName("kotlin.test.assertNotNull")))
     )
   }
+
+  @Test
+  fun `test assertNotNull transformation with local variable message`() {
+    assertMessage(
+      """
+      import kotlin.test.assertNotNull
+      
+      fun main() {
+        val name: String? = null
+        val message = "Message:"
+        assertNotNull(name, message)
+      }""",
+      """
+      Message:
+      assertNotNull(name, message)
+                    |
+                    null ==> expected: not <null>
+      """.trimIndent(),
+      PowerAssertComponentRegistrar(setOf(FqName("kotlin.test.assertNotNull")))
+    )
+  }
 }