Repository: bnorm/kotlin-power-assert
Pull request #41 (2021-04-23 04:49:59)

First attempt at multi-parameter function support


Patch:
diff --git a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/PowerAssertCallTransformer.kt b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/PowerAssertCallTransformer.kt
index e9367f5..a7ea4f0 100644
--- a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/PowerAssertCallTransformer.kt
+++ b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/PowerAssertCallTransformer.kt
@@ -16,6 +16,11 @@
 
 package com.bnorm.power
 
+import com.bnorm.power.delegate.FunctionDelegate
+import com.bnorm.power.delegate.LambdaFunctionDelegate
+import com.bnorm.power.delegate.SimpleFunctionDelegate
+import com.bnorm.power.diagram.IrTemporaryVariable
+import com.bnorm.power.diagram.Node
 import com.bnorm.power.diagram.buildDiagramNesting
 import com.bnorm.power.diagram.buildTree
 import com.bnorm.power.diagram.info
@@ -30,34 +35,25 @@ import org.jetbrains.kotlin.backend.common.lower.DeclarationIrBuilder
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageLocation
 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.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
 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.classOrNull
 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
 import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
@@ -65,7 +61,6 @@ import org.jetbrains.kotlin.ir.util.functions
 import org.jetbrains.kotlin.ir.util.isFunctionOrKFunction
 import org.jetbrains.kotlin.ir.util.kotlinFqName
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
 class PowerAssertCallTransformer(
@@ -78,146 +73,128 @@ class PowerAssertCallTransformer(
   override fun visitCall(expression: IrCall): IrExpression {
     val function = expression.symbol.owner
     val fqName = function.kotlinFqName
-    if (function.valueParameters.isEmpty() || functions.none { fqName == it })
+    if (function.valueParameters.isEmpty() || functions.none { fqName == it }) {
       return super.visitCall(expression)
+    }
 
     // Find a valid delegate function or do not translate
-    val delegate = findDelegate(function) ?: run {
-      val valueType = function.valueParameters[0].type.asString()
+    // TODO better way to determine which delegate to actually use
+    val delegates = findDelegates(function)
+    val delegate = delegates.maxByOrNull { it.function.valueParameters.size }
+    if (delegate == null) {
+      val valueTypesTruncated = function.valueParameters.subList(0, function.valueParameters.size - 1)
+        .joinToString("") { it.type.asString() + ", " }
+      val valueTypesAll = function.valueParameters.joinToString("") { it.type.asString() + ", " }
       messageCollector.warn(
-        expression,
-        "Unable to find overload for function $fqName callable as $fqName($valueType, String) or $fqName($valueType, () -> String) for power-assert transformation"
+        expression = expression,
+        message = """
+          |Unable to find overload of function $fqName for power-assert transformation callable as:
+          | - $fqName(${valueTypesTruncated}String)
+          | - $fqName($valueTypesTruncated() -> String)
+          | - $fqName(${valueTypesAll}String)
+          | - $fqName($valueTypesAll() -> String)
+        """.trimMargin()
       )
       return super.visitCall(expression)
     }
 
-    // TODO - support more arguments by currying?
-    val assertionArgument = expression.getValueArgument(0)!!
-    val messageArgument = if (function.valueParameters.size == 2) expression.getValueArgument(1) else null
+    val messageArgument: IrExpression?
+    val roots: List<Node?>
+    if (delegate.function.valueParameters.size == function.valueParameters.size) {
+      messageArgument = expression.getValueArgument(expression.valueArgumentsCount - 1)
+      roots = (0 until expression.valueArgumentsCount - 1)
+        .map { index -> expression.getValueArgument(index) }
+        .map { arg -> arg?.let { buildTree(it) } }
+    } else {
+      messageArgument = null
+      roots = (0 until expression.valueArgumentsCount)
+        .map { index -> expression.getValueArgument(index) }
+        .map { arg -> arg?.let { buildTree(it) } }
+    }
 
-    // If the tree does not contain any children, the expression is not transformable
-    val root = buildTree(assertionArgument) ?: run {
+    // If all roots are null, there are no transformable parameters
+    if (roots.all { it == null }) {
       messageCollector.info(expression, "Expression is constant and will not be power-assert transformed")
       return super.visitCall(expression)
     }
-//      println(root.dump())
 
     val symbol = currentScope!!.scope.scopeOwnerSymbol
     val builder = DeclarationIrBuilder(context, symbol, expression.startOffset, expression.endOffset)
-    return builder.buildDiagramNesting(root) { argument, variables ->
-      val lambda = messageArgument?.asSimpleLambda()
-      val title = when {
-        messageArgument is IrConst<*> -> messageArgument
-        messageArgument is IrStringConcatenation -> messageArgument
-        lambda != null -> lambda.deepCopyWithSymbols(parent).inline(parent)
-          .transform(ReturnableBlockTransformer(context, symbol), null)
-        messageArgument != null -> {
-          val invoke =
-            messageArgument.type.getClass()!!.functions.single { it.name == OperatorNameConventions.INVOKE }
-          irCallOp(invoke.symbol, invoke.returnType, messageArgument)
-        }
-        // TODO what should the default message be?
-        assertionArgument.type.isBoolean() -> irString("Assertion failed")
-        else -> null
-      }
-
-      val prefix = title?.deepCopyWithSymbols(parent)
-      val diagram = irDiagramString(file, fileSource, prefix, expression, variables)
-      delegate.buildCall(this, expression, argument, diagram)
-    }
+    return builder.diagram(expression, delegate, messageArgument, roots)
 //        .also { println(expression.dump()) }
 //        .also { println(it.dump()) }
 //        .also { println(expression.dumpKotlinLike()) }
 //        .also { println(it.dumpKotlinLike()) }
   }
 
-  private interface FunctionDelegate {
-    fun buildCall(
-      builder: IrBuilderWithScope,
-      original: IrCall,
-      argument: IrExpression,
-      message: IrExpression
-    ): IrExpression
+  private fun DeclarationIrBuilder.diagram(
+    original: IrCall,
+    delegate: FunctionDelegate,
+    messageArgument: IrExpression?,
+    roots: List<Node?>,
+    index: Int = 0,
+    arguments: List<IrExpression?> = listOf(),
+    variables: List<IrTemporaryVariable> = listOf()
+  ): IrExpression {
+    if (index >= roots.size) {
+      val prefix = buildMessagePrefix(messageArgument, roots, original)?.deepCopyWithSymbols(parent)
+      val diagram = irDiagramString(file, fileSource, prefix, original, variables)
+      return delegate.buildCall(this, original, arguments, diagram)
+    } else {
+      val root = roots[index]
+      if (root == null) {
+        val newArguments = arguments + original.getValueArgument(index)
+        return diagram(original, delegate, messageArgument, roots, index + 1, newArguments, variables)
+      } else {
+        return buildDiagramNesting(root) { argument, newVariables ->
+          val newArguments = arguments + argument
+          diagram(original, delegate, messageArgument, roots, index + 1, newArguments, variables + newVariables)
+        }
+      }
+    }
+  }
+
+  private fun DeclarationIrBuilder.buildMessagePrefix(
+    messageArgument: IrExpression?,
+    roots: List<Node?>,
+    original: IrCall
+  ): IrExpression? {
+    val lambda = messageArgument?.asSimpleLambda()
+    return when {
+      messageArgument is IrConst<*> -> messageArgument
+      messageArgument is IrStringConcatenation -> messageArgument
+      lambda != null -> lambda.deepCopyWithSymbols(parent).inline(parent)
+        .transform(ReturnableBlockTransformer(context, scope.scopeOwnerSymbol), null)
+      messageArgument != null -> {
+        val invoke = messageArgument.type.classOrNull!!.owner.functions
+          .single { it.name == OperatorNameConventions.INVOKE }
+        irCallOp(invoke.symbol, invoke.returnType, messageArgument)
+      }
+      // TODO what should the default message be?
+      roots.size == 1 && original.getValueArgument(0)!!.type.isBoolean() -> irString("Assertion failed")
+      else -> null
+    }
   }
 
-  private fun findDelegate(function: IrFunction): FunctionDelegate? {
-    if (function.valueParameters.isEmpty()) return null
+  private fun findDelegates(function: IrFunction): List<FunctionDelegate> {
+    val values = function.valueParameters
+    if (values.isEmpty()) return emptyList()
 
     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 (!function.valueParameters[0].type.isAssignableTo(parameters[0].type)) return@mapNotNull null
+        if (parameters.size !in values.size..values.size + 1) return@mapNotNull null
+        if (!parameters.zip(values).all { (param, value) -> value.type.isAssignableTo(param.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 = original.type).apply {
-                  dispatchReceiver = original.dispatchReceiver?.deepCopyWithSymbols(parent)
-                  extensionReceiver = original.extensionReceiver?.deepCopyWithSymbols(parent)
-                  for (i in 0 until original.typeArgumentsCount) {
-                    putTypeArgument(i, original.getTypeArgument(i))
-                  }
-                  putValueArgument(0, argument)
-                  putValueArgument(1, message)
-                }
-              }
-            }
-          }
-          isStringFunction(messageParameter.type) -> {
-            object : FunctionDelegate {
-              override fun buildCall(
-                builder: IrBuilderWithScope,
-                original: IrCall,
-                argument: IrExpression,
-                message: IrExpression
-              ): IrExpression = with(builder) {
-                val scope = this
-                val lambda = builder.context.irFactory.buildFun {
-                  name = Name.special("<anonymous>")
-                  returnType = context.irBuiltIns.stringType
-                  visibility = DescriptorVisibilities.LOCAL
-                  origin = IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
-                }.apply {
-                  val bodyBuilder = DeclarationIrBuilder(this@PowerAssertCallTransformer.context, symbol)
-                  body = bodyBuilder.irBlockBody {
-                    +irReturn(message)
-                  }
-                  parent = scope.parent
-                }
-                val expression = IrFunctionExpressionImpl(
-                  original.startOffset,
-                  original.endOffset,
-                  messageParameter.type,
-                  lambda,
-                  IrStatementOrigin.LAMBDA
-                )
-                irCall(overload, type = original.type).apply {
-                  dispatchReceiver = original.dispatchReceiver?.deepCopyWithSymbols(parent)
-                  extensionReceiver = original.extensionReceiver?.deepCopyWithSymbols(parent)
-                  for (i in 0 until original.typeArgumentsCount) {
-                    putTypeArgument(i, original.getTypeArgument(i))
-                  }
-                  putValueArgument(0, argument)
-                  putValueArgument(1, expression)
-                }
-              }
-            }
-          }
-          else -> {
-            null
-          }
+          isStringSupertype(messageParameter.type) -> SimpleFunctionDelegate(overload)
+          isStringFunction(messageParameter.type) -> LambdaFunctionDelegate(overload, messageParameter)
+          else -> null
         }
       }
-      .singleOrNull()
   }
 
   private fun isStringFunction(type: IrType): Boolean =
diff --git a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/FunctionDelegate.kt b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/FunctionDelegate.kt
new file mode 100644
index 0000000..956527f
--- /dev/null
+++ b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/FunctionDelegate.kt
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2021 Brian Norman
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.bnorm.power.delegate
+
+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.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
+import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
+
+interface FunctionDelegate {
+  val function: IrFunction
+
+  fun buildCall(
+    builder: IrBuilderWithScope,
+    original: IrCall,
+    arguments: List<IrExpression?>,
+    message: IrExpression
+  ): IrExpression
+
+  fun IrBuilderWithScope.irCallCopy(
+    overload: IrSimpleFunctionSymbol,
+    original: IrCall,
+    arguments: List<IrExpression?>,
+    expression: IrExpression
+  ): IrExpression {
+    return irCall(overload, type = original.type).apply {
+      dispatchReceiver = original.dispatchReceiver?.deepCopyWithSymbols(parent)
+      extensionReceiver = original.extensionReceiver?.deepCopyWithSymbols(parent)
+      for (i in 0 until original.typeArgumentsCount) {
+        putTypeArgument(i, original.getTypeArgument(i))
+      }
+      for ((i, argument) in arguments.withIndex()) {
+        putValueArgument(i, argument?.deepCopyWithSymbols(parent))
+      }
+      putValueArgument(arguments.size, expression)
+    }
+  }
+}
diff --git a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/LambdaFunctionDelegate.kt b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/LambdaFunctionDelegate.kt
new file mode 100644
index 0000000..d18bab6
--- /dev/null
+++ b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/LambdaFunctionDelegate.kt
@@ -0,0 +1,44 @@
+/*
+ * Copyright (C) 2021 Brian Norman
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.bnorm.power.delegate
+
+import com.bnorm.power.irLambda
+import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
+import org.jetbrains.kotlin.ir.builders.irReturn
+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 LambdaFunctionDelegate(
+  private val overload: IrSimpleFunctionSymbol,
+  private val messageParameter: IrValueParameter
+) : FunctionDelegate {
+  override val function = overload.owner
+
+  override fun buildCall(
+    builder: IrBuilderWithScope,
+    original: IrCall,
+    arguments: List<IrExpression?>,
+    message: IrExpression
+  ): IrExpression = with(builder) {
+    val expression = irLambda(context.irBuiltIns.stringType, messageParameter.type) {
+      +irReturn(message)
+    }
+    irCallCopy(overload, original, arguments, expression)
+  }
+}
diff --git a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/SimpleFunctionDelegate.kt b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/SimpleFunctionDelegate.kt
new file mode 100644
index 0000000..186da24
--- /dev/null
+++ b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/SimpleFunctionDelegate.kt
@@ -0,0 +1,35 @@
+/*
+ * Copyright (C) 2021 Brian Norman
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.bnorm.power.delegate
+
+import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
+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
+) : FunctionDelegate {
+  override val function = overload.owner
+
+  override fun buildCall(
+    builder: IrBuilderWithScope,
+    original: IrCall,
+    arguments: List<IrExpression?>,
+    message: IrExpression
+  ): IrExpression = builder.irCallCopy(overload, original, arguments, message)
+}
diff --git a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/IrTemporaryVariable.kt b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/IrTemporaryVariable.kt
index 84218e1..27f3ac3 100644
--- a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/IrTemporaryVariable.kt
+++ b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/IrTemporaryVariable.kt
@@ -1,3 +1,19 @@
+/*
+ * Copyright (C) 2021 Brian Norman
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.bnorm.power.diagram
 
 import org.jetbrains.kotlin.ir.builders.IrStatementsBuilder
diff --git a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/irUtils.kt b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/irUtils.kt
index 05043ab..4ee01c8 100644
--- a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/irUtils.kt
+++ b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/irUtils.kt
@@ -1,7 +1,58 @@
+/*
+ * Copyright (C) 2020 Brian Norman
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.bnorm.power
 
+import org.jetbrains.kotlin.backend.common.lower.DeclarationIrBuilder
+import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
+import org.jetbrains.kotlin.ir.builders.IrBlockBodyBuilder
 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.irString
+import org.jetbrains.kotlin.ir.builders.parent
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.expressions.IrFunctionExpression
+import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionExpressionImpl
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.name.Name
 
 fun IrBuilderWithScope.irString(builderAction: StringBuilder.() -> Unit) =
   irString(buildString { builderAction() })
+
+fun IrBuilderWithScope.irLambda(
+  returnType: IrType,
+  lambdaType: IrType,
+  startOffset: Int = this.startOffset,
+  endOffset: Int = this.endOffset,
+  block: IrBlockBodyBuilder.() -> Unit
+): IrFunctionExpression {
+  val scope = this
+  val lambda = context.irFactory.buildFun {
+    name = Name.special("<anonymous>")
+    this.returnType = returnType
+    visibility = DescriptorVisibilities.LOCAL
+    origin = IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
+  }.apply {
+    val bodyBuilder = DeclarationIrBuilder(context, symbol)
+    body = bodyBuilder.irBlockBody {
+      block()
+    }
+    parent = scope.parent
+  }
+  return IrFunctionExpressionImpl(startOffset, endOffset, lambdaType, lambda, IrStatementOrigin.LAMBDA)
+}


Test patch:
diff --git a/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/ArithmeticExpressionTest.kt b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/ArithmeticExpressionTest.kt
index 8f10802..1b9c9c3 100644
--- a/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/ArithmeticExpressionTest.kt
+++ b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/ArithmeticExpressionTest.kt
@@ -1,3 +1,19 @@
+/*
+ * Copyright (C) 2020 Brian Norman
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.bnorm.power
 
 import org.junit.Test
diff --git a/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/AssertBooleanTest.kt b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/AssertBooleanTest.kt
deleted file mode 100644
index ade3f26..0000000
--- a/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/AssertBooleanTest.kt
+++ /dev/null
@@ -1,44 +0,0 @@
-package com.bnorm.power
-
-import org.jetbrains.kotlin.name.FqName
-import org.junit.Test
-
-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")))
-    )
-  }
-}
diff --git a/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/AssertLibraryTest.kt b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/AssertLibraryTest.kt
new file mode 100644
index 0000000..4f67a68
--- /dev/null
+++ b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/AssertLibraryTest.kt
@@ -0,0 +1,180 @@
+/*
+ * Copyright (C) 2021 Brian Norman
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.bnorm.power
+
+import org.jetbrains.kotlin.name.FqName
+import org.junit.Test
+
+class AssertLibraryTest {
+  @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 assertTrue transformation with message`() {
+    assertMessage(
+      """
+      import kotlin.test.assertTrue
+      
+      fun main() {
+        assertTrue(1 != 1, "Message:")
+      }""",
+      """
+      Message:
+      assertTrue(1 != 1, "Message:")
+                   |
+                   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")))
+    )
+  }
+
+  @Test
+  fun `test assertFalse transformation with message`() {
+    assertMessage(
+      """
+      import kotlin.test.assertFalse
+      
+      fun main() {
+        assertFalse(1 == 1, "Message:")
+      }""",
+      """
+      Message:
+      assertFalse(1 == 1, "Message:")
+                    |
+                    true
+      """.trimIndent(),
+      PowerAssertComponentRegistrar(setOf(FqName("kotlin.test.assertFalse")))
+    )
+  }
+
+  @Test
+  fun `test assertEquals transformation`() {
+    assertMessage(
+      """
+      import kotlin.test.assertEquals
+      
+      fun main() {
+        val greeting = "Hello"
+        val name = "World"
+        assertEquals(greeting, name)
+      }""",
+      """
+      assertEquals(greeting, name)
+                   |         |
+                   |         World
+                   Hello expected:<[Hello]> but was:<[World]>
+      """.trimIndent(),
+      PowerAssertComponentRegistrar(setOf(FqName("kotlin.test.assertEquals")))
+    )
+  }
+
+  @Test
+  fun `test assertEquals transformation with message`() {
+    assertMessage(
+      """
+      import kotlin.test.assertEquals
+      
+      fun main() {
+        val greeting = "Hello"
+        val name = "World"
+        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(
+      """
+      import kotlin.test.assertNotNull
+      
+      fun main() {
+        val name: String? = null
+        assertNotNull(name)
+      }""",
+      """
+      assertNotNull(name)
+                    |
+                    null
+      """.trimIndent(),
+      PowerAssertComponentRegistrar(setOf(FqName("kotlin.test.assertNotNull")))
+    )
+  }
+
+  @Test
+  fun `test assertNotNull transformation with message`() {
+    assertMessage(
+      """
+      import kotlin.test.assertNotNull
+      
+      fun main() {
+        val name: String? = null
+        assertNotNull(name, "Message:")
+      }""",
+      """
+      Message:
+      assertNotNull(name, "Message:")
+                    |
+                    null
+      """.trimIndent(),
+      PowerAssertComponentRegistrar(setOf(FqName("kotlin.test.assertNotNull")))
+    )
+  }
+}
diff --git a/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/DebugFunctionTest.kt b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/DebugFunctionTest.kt
index b77dad7..d720e3e 100644
--- a/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/DebugFunctionTest.kt
+++ b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/DebugFunctionTest.kt
@@ -1,3 +1,19 @@
+/*
+ * Copyright (C) 2021 Brian Norman
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.bnorm.power
 
 import com.tschuchort.compiletesting.KotlinCompilation
@@ -27,6 +43,7 @@ class DebugFunctionTest {
       actual.trim()
     )
   }
+
   @Test
   fun `debug function transformation with message`() {
     val actual = executeMainDebug(
@@ -47,7 +64,7 @@ class DebugFunctionTest {
   }
 }
 
-fun executeMainDebug(mainBody: String): String {
+private fun executeMainDebug(mainBody: String): String {
   val file = SourceFile.kotlin(
     name = "main.kt",
     contents = """
diff --git a/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/MultiParameterFunctionTest.kt b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/MultiParameterFunctionTest.kt
new file mode 100644
index 0000000..1a2ce2b
--- /dev/null
+++ b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/MultiParameterFunctionTest.kt
@@ -0,0 +1,170 @@
+/*
+ * Copyright (C) 2021 Brian Norman
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+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 MultiParameterFunctionTest {
+  @Test
+  fun `debug multi-parameter function transformation`() {
+    val actual = executeMainDebug(
+      """
+      val operation = "sum"
+      dbg(operation, 1 + 2 + 3)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      sum=6
+      dbg(operation, 1 + 2 + 3)
+          |            |   |
+          |            |   6
+          |            3
+          sum
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `debug multi-parameter function transformation with complex booleans`() {
+    val actual = executeMainDebug(
+      """
+      val greeting: String? = null
+      val name: String? = null
+      dbg(
+        key = greeting != null && greeting.length == 5,
+        value = name == null || name.length == 5
+      )
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      false=true
+      dbg(
+        key = greeting != null && greeting.length == 5,
+              |        |
+              |        false
+              null
+        value = name == null || name.length == 5
+                |    |
+                |    true
+                null
+      )
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `debug multi-parameter function transformation with message`() {
+    val actual = executeMainDebug(
+      """
+      val operation = "sum"
+      dbg(operation, 1 + 2 + 3, "Message:")
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      sum=6
+      Message:
+      dbg(operation, 1 + 2 + 3, "Message:")
+          |            |   |
+          |            |   6
+          |            3
+          sum
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `debug multi-parameter function transformation with message and complex booleans`() {
+    val actual = executeMainDebug(
+      """
+      val greeting: String? = null
+      val name: String? = null
+      dbg(
+        key = greeting != null && greeting.length == 5,
+        value = name == null || name.length == 5,
+        msg = "Message:"
+      )
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      false=true
+      Message:
+      dbg(
+        key = greeting != null && greeting.length == 5,
+              |        |
+              |        false
+              null
+        value = name == null || name.length == 5,
+                |    |
+                |    true
+                null
+        msg = "Message:"
+      )
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+}
+
+private fun executeMainDebug(mainBody: String): String {
+  val file = SourceFile.kotlin(
+    name = "main.kt",
+    contents = """
+fun <T> dbg(key: Any, value: T): T = value
+
+fun <T> dbg(key: Any, value: T, msg: String): T {
+    println(key.toString() + "=" + value + "\n" + 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("UTF-8")
+  } catch (t: InvocationTargetException) {
+    throw t.cause!!
+  } finally {
+    System.setOut(prevOut)
+  }
+}
diff --git a/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/RegexMatchTest.kt b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/RegexMatchTest.kt
index ac40b07..8207d6b 100644
--- a/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/RegexMatchTest.kt
+++ b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/RegexMatchTest.kt
@@ -1,3 +1,19 @@
+/*
+ * Copyright (C) 2021 Brian Norman
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.bnorm.power
 
 import org.junit.Test
diff --git a/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/compiler.kt b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/compiler.kt
index 5726ed1..cfc28e0 100644
--- a/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/compiler.kt
+++ b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/compiler.kt
@@ -1,3 +1,19 @@
+/*
+ * Copyright (C) 2020 Brian Norman
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.bnorm.power
 
 import com.tschuchort.compiletesting.KotlinCompilation