Allow functions with a single argument of any type

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

@@ -35,17 +35,18 @@ import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
 import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.backend.js.utils.asString
 import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
 import org.jetbrains.kotlin.ir.builders.declarations.buildFun
 import org.jetbrains.kotlin.ir.builders.irBlockBody
 import org.jetbrains.kotlin.ir.builders.irCall
 import org.jetbrains.kotlin.ir.builders.irCallOp
-import org.jetbrains.kotlin.ir.builders.irFalse
 import org.jetbrains.kotlin.ir.builders.irReturn
 import org.jetbrains.kotlin.ir.builders.irString
 import org.jetbrains.kotlin.ir.builders.parent
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
 import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.declarations.path
 import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrConst
@@ -53,10 +54,12 @@ import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
 import org.jetbrains.kotlin.ir.expressions.IrStringConcatenation
 import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionExpressionImpl
+import org.jetbrains.kotlin.ir.symbols.IrTypeParameterSymbol
 import org.jetbrains.kotlin.ir.types.IrSimpleType
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.IrTypeArgument
 import org.jetbrains.kotlin.ir.types.IrTypeProjection
+import org.jetbrains.kotlin.ir.types.classifierOrNull
 import org.jetbrains.kotlin.ir.types.getClass
 import org.jetbrains.kotlin.ir.types.isBoolean
 import org.jetbrains.kotlin.ir.types.isSubtypeOf
@@ -99,23 +102,25 @@ class PowerAssertCallTransformer(
       .replace("\r\n", "\n") // https://youtrack.jetbrains.com/issue/KT-41888
 
     irFile.transformChildrenVoid()
+//    println(irFile.dumpKotlinLike())
   }
 
   override fun visitCall(expression: IrCall): IrExpression {
-    val fqName = expression.symbol.owner.kotlinFqName
+    val function = expression.symbol.owner
-    if (functions.none { fqName == it })
+    val fqName = function.kotlinFqName
+    if (function.valueParameters.isEmpty() || functions.none { fqName == it })
       return super.visitCall(expression)
 
     // Find a valid delegate function or do not translate
-    val delegate = findDelegate(fqName) ?: run {
+    val delegate = findDelegate(function) ?: run {
+      val valueType = function.valueParameters[0].type.asString()
       messageCollector.warn(
         expression,
-        "Unable to find overload for function $fqName callable as $fqName(Boolean, String) or $fqName(Boolean, () -> String) for power-assert transformation"
+        "Unable to find overload for function $fqName callable as $fqName($valueType, String) or $fqName($valueType, () -> String) for power-assert transformation"
       )
       return super.visitCall(expression)
     }
 
-    val function = expression.symbol.owner
     val assertionArgument = expression.getValueArgument(0)!!
     val messageArgument = if (function.valueParameters.size == 2) expression.getValueArgument(1) else null
 
@@ -130,7 +135,7 @@ class PowerAssertCallTransformer(
       at(expression)
 
       val generator = object : DiagramGenerator() {
-        override fun IrBuilderWithScope.buildAssertThrow(variables: List<IrTemporaryVariable>): IrExpression {
+        override fun IrBuilderWithScope.buildCall(argument: IrExpression, variables: List<IrTemporaryVariable>): IrExpression {
 
           val lambda = messageArgument?.asSimpleLambda()
           val title = when {
@@ -142,18 +147,18 @@ class PowerAssertCallTransformer(
               irCallOp(invoke.symbol, invoke.returnType, messageArgument)
             }
             // TODO what should the default message be?
-            else -> irString("Assertion failed")
+            assertionArgument.type.isBoolean() -> irString("Assertion failed")
+            else -> null
           }
 
-          val prefix = title.deepCopyWithSymbols(parent)
+          val prefix = title?.deepCopyWithSymbols(parent)
           val diagram = irDiagram(file, fileSource, prefix, expression, variables)
-          return delegate.buildCall(this, expression, irFalse(), diagram)
+          return delegate.buildCall(this, expression, argument, diagram)
         }
       }
 
 //      println(root.dump())
-
+      return generator.buildExpression(this, root)
-      return generator.buildAssert(this, root)
 //        .also { println(expression.dump()) }
 //        .also { println(it.dump()) }
 //        .also { println(expression.dumpKotlinLike()) }
@@ -165,20 +170,22 @@ class PowerAssertCallTransformer(
     fun buildCall(builder: IrBuilderWithScope, original: IrCall, argument: IrExpression, message: IrExpression): IrExpression
   }
 
-  private fun findDelegate(fqName: FqName): FunctionDelegate? {
+  private fun findDelegate(function: IrFunction): FunctionDelegate? {
-    return context.referenceFunctions(fqName)
+    if (function.valueParameters.isEmpty()) return null
+
+    return context.referenceFunctions(function.kotlinFqName)
       .mapNotNull { overload ->
         // TODO allow other signatures than (Boolean, String) and (Boolean, () -> String)
         val parameters = overload.owner.valueParameters
         if (parameters.size != 2) return@mapNotNull null
-        if (!parameters[0].type.isBoolean()) return@mapNotNull null
+        if (!function.valueParameters[0].type.isAssignableTo(parameters[0].type)) return@mapNotNull null
 
         val messageParameter = parameters.last()
         return@mapNotNull when {
           isStringSupertype(messageParameter.type) -> {
             object : FunctionDelegate {
               override fun buildCall(builder: IrBuilderWithScope, original: IrCall, argument: IrExpression, message: IrExpression): IrExpression = with(builder) {
-                irCall(overload, type = overload.owner.returnType).apply {
+                irCall(overload, type = original.type).apply {
                   dispatchReceiver = original.dispatchReceiver?.deepCopyWithSymbols(parent)
                   extensionReceiver = original.extensionReceiver?.deepCopyWithSymbols(parent)
                   for (i in 0 until original.typeArgumentsCount) {
@@ -207,7 +214,7 @@ class PowerAssertCallTransformer(
                   parent = scope.parent
                 }
                 val expression = IrFunctionExpressionImpl(original.startOffset, original.endOffset, messageParameter.type, lambda, IrStatementOrigin.LAMBDA)
-                irCall(overload, type = overload.owner.returnType).apply {
+                irCall(overload, type = original.type).apply {
                   dispatchReceiver = original.dispatchReceiver?.deepCopyWithSymbols(parent)
                   extensionReceiver = original.extensionReceiver?.deepCopyWithSymbols(parent)
                   for (i in 0 until original.typeArgumentsCount) {
@@ -236,6 +243,12 @@ class PowerAssertCallTransformer(
   private fun isStringSupertype(type: IrType): Boolean =
     context.irBuiltIns.stringType.isSubtypeOf(type, context.irBuiltIns)
 
+  private fun IrType.isAssignableTo(type: IrType): Boolean =
+    isSubtypeOf(type, context.irBuiltIns) ||
+      ((type.classifierOrNull as? IrTypeParameterSymbol)?.owner?.superTypes?.all {
+        isSubtypeOf(it, context.irBuiltIns)
+      } ?: false)
+
   private fun MessageCollector.info(expression: IrElement, message: String) {
     report(expression, CompilerMessageSeverity.INFO, message)
   }

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

@@ -16,45 +16,50 @@
 
 package com.bnorm.power.diagram
 
-import org.jetbrains.kotlin.backend.common.lower.irIfThen
-import org.jetbrains.kotlin.backend.common.lower.irNot
 import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
 import org.jetbrains.kotlin.ir.builders.IrStatementsBuilder
 import org.jetbrains.kotlin.ir.builders.irBlock
+import org.jetbrains.kotlin.ir.builders.irFalse
+import org.jetbrains.kotlin.ir.builders.irIfThenElse
+import org.jetbrains.kotlin.ir.builders.irTrue
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
 
 abstract class DiagramGenerator {
-  abstract fun IrBuilderWithScope.buildAssertThrow(variables: List<IrTemporaryVariable>): IrExpression
+  abstract fun IrBuilderWithScope.buildCall(
+    argument: IrExpression,
+    variables: List<IrTemporaryVariable>
+  ): IrExpression
 
-  fun buildAssert(
+  fun buildExpression(
     builder: IrBuilderWithScope,
     root: Node
   ): IrExpression {
     return builder.irBlock {
-      buildAssert(root, listOf()) { subStack ->
+      buildExpression(root, listOf()) { argument, subStack ->
-        buildAssertThrow(subStack)
+        buildCall(argument, subStack)
       }
     }
   }
 
-  private fun IrStatementsBuilder<*>.buildAssert(
+  private fun IrStatementsBuilder<*>.buildExpression(
     node: Node,
     variables: List<IrTemporaryVariable>,
-    thenPart: IrStatementsBuilder<*>.(List<IrTemporaryVariable>) -> IrExpression
+    call: IrStatementsBuilder<*>.(IrExpression, List<IrTemporaryVariable>) -> IrExpression
-  ): List<IrTemporaryVariable> = when (node) {
+  ) {
-    is ExpressionNode -> add(node, variables, thenPart)
+    when (node) {
-    is AndNode -> chain(node, variables, thenPart)
+      is ExpressionNode -> add(node, variables, call)
-    is OrNode -> nest(node, 0, variables)
+      is AndNode -> nest(node, 0, variables, call)
-    else -> TODO("Unknown node type=$node")
+      is OrNode -> nest(node, 0, variables, call)
+      else -> TODO("Unknown node type=$node")
+    }
   }
 
   /**
-   * TODO - this is how the following function should behave
    * ```
    * val result = call(1 + 2 + 3)
    * ```
-   * Transformed to
+   * Transforms to
    * ```
    * val result = run {
    *   val tmp0 = 1 + 2
@@ -66,67 +71,63 @@ abstract class DiagramGenerator {
   private fun IrStatementsBuilder<*>.add(
     node: ExpressionNode,
     variables: List<IrTemporaryVariable>,
-    thenPart: IrStatementsBuilder<*>.(List<IrTemporaryVariable>) -> IrExpression
+    call: IrStatementsBuilder<*>.(IrExpression, List<IrTemporaryVariable>) -> IrExpression
-  ): List<IrTemporaryVariable> {
+  ) {
     val head = node.expressions.first().deepCopyWithSymbols(scope.getLocalDeclarationParent())
     val expressions = (buildTree(head) as ExpressionNode).expressions
     val transformer = IrTemporaryExtractionTransformer(this, expressions.toSet())
     val transformed = expressions.first().transform(transformer, null)
 
-    if (transformed.type == context.irBuiltIns.booleanType) {
+    +call(transformed, variables + transformer.variables)
-      // TODO irIfThenElse to support property nesting of ANDs & ORs
-      +irIfThen(irNot(transformed), thenPart(variables + transformer.variables))
-    } else {
-      +thenPart(variables + transformer.variables)
-    }
-
-    return transformer.variables
   }
 
   /**
-   * TODO - this is how the following function should behave
    * ```
    * val result = call(1 == 1 && 2 == 2)
    * ```
-   * Transformed to
+   * Transforms to
    * ```
    * val result = run {
    *   val tmp0 = 1 == 1
    *   if (tmp0) {
    *     val tmp1 = 2 == 2
-   *     if (tmp1) call(true, <diagram>)
+   *     call(tmp1, <diagram>)
-   *     else call(false, <diagram>)
    *   }
    *   else call(false, <diagram>)
    * }
    * ```
    */
-  private fun IrStatementsBuilder<*>.chain(
+  private fun IrStatementsBuilder<*>.nest(
     node: AndNode,
+    index: Int,
     variables: List<IrTemporaryVariable>,
-    thenPart: IrStatementsBuilder<*>.(List<IrTemporaryVariable>) -> IrExpression
+    call: IrStatementsBuilder<*>.(IrExpression, List<IrTemporaryVariable>) -> IrExpression
-  ): List<IrTemporaryVariable> {
+  ) {
-    val newVariables = mutableListOf<IrTemporaryVariable>()
+    val children = node.children
-    for (child in node.children) {
+    val child = children[index]
-      newVariables.addAll(buildAssert(child, variables + newVariables, thenPart))
+    buildExpression(child, variables) { argument, newVariables ->
+      if (index + 1 == children.size) call(argument, newVariables) // last expression, result is false
+      else irIfThenElse(
+        context.irBuiltIns.anyType,
+        argument,
+        irBlock { nest(node, index + 1, newVariables, call) }, // more expressions, continue nesting
+        call(irFalse(), newVariables), // short-circuit result to false
+      )
     }
-    return newVariables
   }
 
   /**
-   * TODO - this is how the following function should behave
    * ```
    * val result = call(1 == 1 || 2 == 2)
    * ```
-   * Transformed to
+   * Transforms to
    * ```
    * val result = run {
    *   val tmp0 = 1 == 1
    *   if (tmp0) call(true, <diagram>)
    *   else {
    *     val tmp1 = 2 == 2
-   *     if (tmp1) call(true, <diagram>)
+   *     call(tmp1, <diagram>)
-   *     else call(false, <diagram>)
    *   }
    * }
    * ```
@@ -134,14 +135,19 @@ abstract class DiagramGenerator {
   private fun IrStatementsBuilder<*>.nest(
     node: OrNode,
     index: Int,
-    variables: List<IrTemporaryVariable>
+    variables: List<IrTemporaryVariable>,
-  ): List<IrTemporaryVariable> {
+    call: IrStatementsBuilder<*>.(IrExpression, List<IrTemporaryVariable>) -> IrExpression
+  ) {
     val children = node.children
     val child = children[index]
-    buildAssert(child, variables) { newVariables ->
+    buildExpression(child, variables) { argument, newVariables ->
-      if (index + 1 == children.size) buildAssertThrow(newVariables)
+      if (index + 1 == children.size) call(argument, newVariables) // last expression, result is false
-      else irBlock { nest(node, index + 1, newVariables) }
+      else irIfThenElse(
+        context.irBuiltIns.anyType,
+        argument,
+        call(irTrue(), newVariables), // short-circuit result to true
+        irBlock { nest(node, index + 1, newVariables, call) }, // more expressions, continue nesting
+      )
     }
-    return emptyList()
   }
 }

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

@@ -59,7 +59,7 @@ fun IrBuilderWithScope.irDiagram(
 
       addArgument(
         irString {
-          appendLine()
+          if (row != 0 || prefix != null) appendLine()
           append(rowSource)
           if (indentations.isNotEmpty()) {
             appendLine()

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

@@ -0,0 +1,45 @@
+package com.bnorm.power
+
+import org.jetbrains.kotlin.name.FqName
+import org.junit.Test
+import kotlin.test.assertEquals
+
+class AssertBooleanTest {
+  @Test
+  fun `test assertTrue transformation`() {
+    assertMessage(
+      """
+import kotlin.test.assertTrue
+
+fun main() {
+  assertTrue(1 != 1)
+}""",
+      """
+Assertion failed
+assertTrue(1 != 1)
+             |
+             false
+""".trimIndent(),
+      PowerAssertComponentRegistrar(setOf(FqName("kotlin.test.assertTrue")))
+    )
+  }
+
+  @Test
+  fun `test assertFalse transformation`() {
+    assertMessage(
+      """
+import kotlin.test.assertFalse
+
+fun main() {
+  assertFalse(1 == 1)
+}""",
+      """
+Assertion failed
+assertFalse(1 == 1)
+              |
+              true
+""".trimIndent(),
+      PowerAssertComponentRegistrar(setOf(FqName("kotlin.test.assertFalse")))
+    )
+  }
+}

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

@@ -0,0 +1,80 @@
+package com.bnorm.power
+
+import com.tschuchort.compiletesting.KotlinCompilation
+import com.tschuchort.compiletesting.SourceFile
+import org.jetbrains.kotlin.name.FqName
+import org.junit.Test
+import java.io.ByteArrayOutputStream
+import java.io.PrintStream
+import java.lang.reflect.InvocationTargetException
+import kotlin.test.assertEquals
+
+class DebugFunctionTest {
+  @Test
+  fun `debug function transformation`() {
+    val actual = executeMainDebug("""
+      dbg(1 + 2 + 3)
+    """.trimIndent())
+    assertEquals(
+      """
+      dbg(1 + 2 + 3)
+            |   |
+            |   6
+            3
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+  @Test
+  fun `debug function transformation with message`() {
+    val actual = executeMainDebug("""
+      dbg(1 + 2 + 3, "Message:")
+    """.trimIndent())
+    assertEquals(
+      """
+      Message:
+      dbg(1 + 2 + 3, "Message:")
+            |   |
+            |   6
+            3
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+}
+
+fun executeMainDebug(mainBody: String): String {
+  val file = SourceFile.kotlin(
+    name = "main.kt",
+    contents = """
+fun <T> dbg(value: T): T = value
+
+fun <T> dbg(value: T, msg: String): T {
+    println(msg)
+    return value
+}
+
+fun main() {
+  $mainBody
+}
+""",
+    trimIndent = false
+  )
+
+  val result = compile(listOf(file), PowerAssertComponentRegistrar(setOf(FqName("dbg"))))
+  assertEquals(KotlinCompilation.ExitCode.OK, result.exitCode)
+
+  val kClazz = result.classLoader.loadClass("MainKt")
+  val main = kClazz.declaredMethods.single { it.name == "main" && it.parameterCount == 0 }
+  val prevOut = System.out
+  try {
+    val out = ByteArrayOutputStream()
+    System.setOut(PrintStream(out))
+    main.invoke(null)
+    return out.toString(Charsets.UTF_8)
+  } catch (t: InvocationTargetException) {
+    throw t.cause!!
+  } finally {
+    System.setOut(prevOut)
+  }
+}

sample/build.gradle.kts

@@ -62,6 +62,7 @@ configure<com.bnorm.power.PowerAssertGradleExtension> {
     "kotlin.test.assertTrue",
     "kotlin.require",
     "com.bnorm.power.AssertScope.assert",
-    "com.bnorm.power.assert"
+    "com.bnorm.power.assert",
+    "com.bnorm.power.dbg"
   )
 }

sample/src/commonMain/kotlin/com/bnorm/power/Debug.kt

@@ -0,0 +1,10 @@
+package com.bnorm.power
+
+val debugLog = StringBuilder()
+
+fun <T> dbg(value: T): T = value
+
+fun <T> dbg(value: T, msg: String): T {
+  debugLog.appendLine(msg)
+  return value
+}

sample/src/commonTest/kotlin/com/bnorm/power/PowerAssertTest.kt

@@ -16,6 +16,7 @@
 
 package com.bnorm.power
 
+import kotlin.test.AfterTest
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith
@@ -23,6 +24,11 @@ import kotlin.test.assertTrue
 
 class PowerAssertTest {
 
+  @AfterTest
+  fun cleanup() {
+    debugLog.clear()
+  }
+
   @Test
   fun assertTrue() {
     val error = assertFailsWith<AssertionError> { assertTrue(Person.UNKNOWN.size == 1) }
@@ -111,4 +117,35 @@ class PowerAssertTest {
       """.trimIndent(),
     )
   }
+
+  @Test
+  fun dbgTest() {
+    val name = "Jane"
+    val greeting = dbg("Hello, $name")
+    assert(greeting == "Hello, Jane")
+    assertEquals(
+      actual = debugLog.toString().trim(),
+      expected = """
+        dbg("Hello, ${"$"}name")
+            |        |
+            |        Jane
+            Hello, Jane
+      """.trimIndent())
+  }
+
+  @Test
+  fun dbgMessageTest() {
+    val name = "Jane"
+    val greeting = dbg("Hello, $name", "Greeting:")
+    assert(greeting == "Hello, Jane")
+    assertEquals(
+      actual = debugLog.toString().trim(),
+      expected = """
+        Greeting:
+        dbg("Hello, ${"$"}name", "Greeting:")
+            |        |
+            |        Jane
+            Hello, Jane
+      """.trimIndent())
+  }
 }