Support dispatch infix functions

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

@@ -24,6 +24,7 @@ import com.bnorm.power.diagram.IrTemporaryVariable
 import com.bnorm.power.diagram.Node
 import com.bnorm.power.diagram.SourceFile
 import com.bnorm.power.diagram.buildDiagramNesting
+import com.bnorm.power.diagram.buildDiagramNestingNullable
 import com.bnorm.power.diagram.buildTree
 import com.bnorm.power.diagram.irDiagramString
 import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
@@ -99,6 +100,8 @@ class PowerAssertCallTransformer(
       return super.visitCall(expression)
     }
 
+    val dispatchRoot =
+      if (expression.symbol.owner.isInfix) expression.dispatchReceiver?.let { buildTree(it) } else null
     val extensionRoot =
       if (expression.symbol.owner.isInfix) expression.extensionReceiver?.let { buildTree(it) } else null
     val messageArgument: IrExpression?
@@ -116,14 +119,21 @@ class PowerAssertCallTransformer(
     }
 
     // If all roots are null, there are no transformable parameters
-    if (extensionRoot == null && roots.all { it == null }) {
+    if (dispatchRoot == null && extensionRoot == null && roots.all { it == null }) {
       messageCollector.info(expression, "Expression is constant and will not be power-assert transformed")
       return super.visitCall(expression)
     }
 
     val symbol = currentScope!!.scope.scopeOwnerSymbol
     val builder = DeclarationIrBuilder(context, symbol, expression.startOffset, expression.endOffset)
-    return builder.diagram(expression, delegate, messageArgument, roots, extensionRoot)
+    return builder.diagram(
+      call = expression,
+      delegate = delegate,
+      messageArgument = messageArgument,
+      roots = roots,
+      dispatchRoot = dispatchRoot,
+      extensionRoot = extensionRoot
+    )
   }
 
   private fun DeclarationIrBuilder.diagram(
@@ -131,11 +141,13 @@ class PowerAssertCallTransformer(
     delegate: FunctionDelegate,
     messageArgument: IrExpression?,
     roots: List<Node?>,
+    dispatchRoot: Node? = null,
     extensionRoot: Node? = null
   ): IrExpression {
     fun recursive(
       index: Int,
-      extension: IrExpression? = null,
+      dispatch: IrExpression?,
+      extension: IrExpression?,
       arguments: List<IrExpression?>,
       variables: List<IrTemporaryVariable>
     ): IrExpression {
@@ -143,27 +155,25 @@ class PowerAssertCallTransformer(
         val prefix = buildMessagePrefix(messageArgument, delegate.messageParameter, roots, call)
           ?.deepCopyWithSymbols(parent)
         val diagram = irDiagramString(sourceFile, prefix, call, variables)
-        return delegate.buildCall(this, call, extension, arguments, diagram)
+        return delegate.buildCall(this, call, dispatch, extension, arguments, diagram)
       } else {
         val root = roots[index]
         if (root == null) {
           val newArguments = arguments + call.getValueArgument(index)
-          return recursive(index + 1, extension, newArguments, variables)
+          return recursive(index + 1, dispatch, extension, newArguments, variables)
         } else {
-          return buildDiagramNesting(root) { argument, newVariables ->
+          return buildDiagramNesting(root, variables) { argument, newVariables ->
             val newArguments = arguments + argument
-            recursive(index + 1, extension, newArguments, variables + newVariables)
+            recursive(index + 1, dispatch, extension, newArguments, newVariables)
           }
         }
       }
     }
 
-    return if (extensionRoot != null) {
+    return buildDiagramNestingNullable(dispatchRoot) { dispatch, newVariables ->
-      buildDiagramNesting(extensionRoot) { extension, newVariables ->
+      buildDiagramNestingNullable(extensionRoot, newVariables) { extension, newVariables ->
-        recursive(0, extension, emptyList(), newVariables)
+        recursive(0, dispatch, extension, emptyList(), newVariables)
       }
-    } else {
-      recursive(0, null, emptyList(), emptyList())
     }
   }
 
@@ -216,7 +226,7 @@ class PowerAssertCallTransformer(
 
     return possible.mapNotNull { overload ->
       // Dispatch receivers must always match exactly
-      if (function.dispatchReceiverParameter != overload.owner.dispatchReceiverParameter) {
+      if (function.dispatchReceiverParameter?.type != overload.owner.dispatchReceiverParameter?.type) {
         return@mapNotNull null
       }
 

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

@@ -33,6 +33,7 @@ interface FunctionDelegate {
   fun buildCall(
     builder: IrBuilderWithScope,
     original: IrCall,
+    dispatchReceiver: IrExpression?,
     extensionReceiver: IrExpression?,
     valueArguments: List<IrExpression?>,
     messageArgument: IrExpression
@@ -41,12 +42,13 @@ interface FunctionDelegate {
   fun IrBuilderWithScope.irCallCopy(
     overload: IrSimpleFunctionSymbol,
     original: IrCall,
+    dispatchReceiver: IrExpression?,
     extensionReceiver: IrExpression?,
     valueArguments: List<IrExpression?>,
     messageArgument: IrExpression
   ): IrExpression {
     return irCall(overload, type = original.type).apply {
-      dispatchReceiver = original.dispatchReceiver?.deepCopyWithSymbols(parent)
+      this.dispatchReceiver = original.dispatchReceiver?.deepCopyWithSymbols(parent)
       this.extensionReceiver = (extensionReceiver ?: original.extensionReceiver)?.deepCopyWithSymbols(parent)
       for (i in 0 until original.typeArgumentsCount) {
         putTypeArgument(i, original.getTypeArgument(i))

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

@@ -33,6 +33,7 @@ class LambdaFunctionDelegate(
   override fun buildCall(
     builder: IrBuilderWithScope,
     original: IrCall,
+    dispatchReceiver: IrExpression?,
     extensionReceiver: IrExpression?,
     valueArguments: List<IrExpression?>,
     messageArgument: IrExpression
@@ -40,6 +41,13 @@ class LambdaFunctionDelegate(
     val expression = irLambda(context.irBuiltIns.stringType, messageParameter.type) {
       +irReturn(messageArgument)
     }
-    irCallCopy(overload, original, extensionReceiver, valueArguments, expression)
+    irCallCopy(
+      overload = overload,
+      original = original,
+      dispatchReceiver = dispatchReceiver,
+      extensionReceiver = extensionReceiver,
+      valueArguments = valueArguments,
+      messageArgument = expression
+    )
   }
 }

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

@@ -20,11 +20,9 @@ import com.bnorm.power.irLambda
 import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
 import org.jetbrains.kotlin.ir.builders.irReturn
 import org.jetbrains.kotlin.ir.builders.irSamConversion
-import org.jetbrains.kotlin.ir.builders.typeOperator
 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.expressions.IrTypeOperator
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 
 class SamConversionLambdaFunctionDelegate(
@@ -36,6 +34,7 @@ class SamConversionLambdaFunctionDelegate(
   override fun buildCall(
     builder: IrBuilderWithScope,
     original: IrCall,
+    dispatchReceiver: IrExpression?,
     extensionReceiver: IrExpression?,
     valueArguments: List<IrExpression?>,
     messageArgument: IrExpression
@@ -44,6 +43,13 @@ class SamConversionLambdaFunctionDelegate(
       +irReturn(messageArgument)
     }
     val expression = irSamConversion(lambda, messageParameter.type)
-    irCallCopy(overload, original, extensionReceiver, valueArguments, expression)
+    irCallCopy(
+      overload = overload,
+      original = original,
+      dispatchReceiver = dispatchReceiver,
+      extensionReceiver = extensionReceiver,
+      valueArguments = valueArguments,
+      messageArgument = expression
+    )
   }
 }

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

@@ -31,8 +31,16 @@ class SimpleFunctionDelegate(
   override fun buildCall(
     builder: IrBuilderWithScope,
     original: IrCall,
+    dispatchReceiver: IrExpression?,
     extensionReceiver: IrExpression?,
     valueArguments: List<IrExpression?>,
     messageArgument: IrExpression
-  ): IrExpression = builder.irCallCopy(overload, original, extensionReceiver, valueArguments, messageArgument)
+  ): IrExpression = builder.irCallCopy(
+    overload = overload,
+    original = original,
+    dispatchReceiver = dispatchReceiver,
+    extensionReceiver = extensionReceiver,
+    valueArguments = valueArguments,
+    messageArgument = messageArgument
+  )
 }

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

@@ -26,13 +26,22 @@ import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
 
 fun IrBuilderWithScope.buildDiagramNesting(
   root: Node,
+  variables: List<IrTemporaryVariable> = emptyList(),
   call: IrBuilderWithScope.(IrExpression, List<IrTemporaryVariable>) -> IrExpression
 ): IrExpression {
-  return buildExpression(root, listOf()) { argument, subStack ->
+  return buildExpression(root, variables) { argument, subStack ->
     call(argument, subStack)
   }
 }
 
+fun IrBuilderWithScope.buildDiagramNestingNullable(
+  root: Node?,
+  variables: List<IrTemporaryVariable> = emptyList(),
+  call: IrBuilderWithScope.(IrExpression?, List<IrTemporaryVariable>) -> IrExpression
+): IrExpression {
+  return if (root != null) buildDiagramNesting(root, variables, call) else call(null, variables)
+}
+
 private fun IrBuilderWithScope.buildExpression(
   node: Node,
   variables: List<IrTemporaryVariable>,

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

@@ -26,8 +26,8 @@ import kotlin.test.fail
 
 class InfixFunctionTest {
   @Test
-  fun `infix function call includes receiver`() {
+  fun `extension infix function call includes receiver`() {
-    val actual = execute(
+    val actual = runExtensionInfix(
       """
       (1 + 1) mustEqual (2 + 4)
       """.trimIndent()
@@ -44,8 +44,8 @@ class InfixFunctionTest {
   }
 
   @Test
-  fun `infix function call with constant receiver`() {
+  fun `extension infix function call with constant receiver`() {
-    val actual = execute(
+    val actual = runExtensionInfix(
       """
       1 mustEqual (2 + 4)
       """.trimIndent()
@@ -61,8 +61,8 @@ class InfixFunctionTest {
   }
 
   @Test
-  fun `infix function call with constant parameter`() {
+  fun `extension infix function call with constant parameter`() {
-    val actual = execute(
+    val actual = runExtensionInfix(
       """
       (1 + 1) mustEqual 6
       """.trimIndent()
@@ -78,8 +78,8 @@ class InfixFunctionTest {
   }
 
   @Test
-  fun `infix function call with only constants`() {
+  fun `extension infix function call with only constants`() {
-    val actual = execute(
+    val actual = runExtensionInfix(
       """
       2 mustEqual 6
       """.trimIndent()
@@ -93,8 +93,8 @@ class InfixFunctionTest {
   }
 
   @Test
-  fun `non-infix function call includes receiver`() {
+  fun `extension non-infix function call includes receiver`() {
-    val actual = execute(
+    val actual = runExtensionInfix(
       """
       (1 + 1).mustEqual(2 + 4)
       """.trimIndent()
@@ -111,8 +111,8 @@ class InfixFunctionTest {
   }
 
   @Test
-  fun `non-infix function call with constant receiver`() {
+  fun `extension non-infix function call with constant receiver`() {
-    val actual = execute(
+    val actual = runExtensionInfix(
       """
       1.mustEqual(2 + 4)
       """.trimIndent()
@@ -128,8 +128,8 @@ class InfixFunctionTest {
   }
 
   @Test
-  fun `non-infix function call with constant parameter`() {
+  fun `extension non-infix function call with constant parameter`() {
-    val actual = execute(
+    val actual = runExtensionInfix(
       """
       (1 + 1).mustEqual(6)
       """.trimIndent()
@@ -145,8 +145,8 @@ class InfixFunctionTest {
   }
 
   @Test
-  fun `non-infix function call with only constants`() {
+  fun `extension non-infix function call with only constants`() {
-    val actual = execute(
+    val actual = runExtensionInfix(
       """
       2.mustEqual(6)
       """.trimIndent()
@@ -159,24 +159,199 @@ class InfixFunctionTest {
     )
   }
 
-  private fun execute(mainBody: String): String {
+  @Test
-    val file = SourceFile.kotlin(
+  fun `dispatch infix function call includes receiver`() {
-      name = "main.kt",
+    val actual = runDispatchInfix(
-      contents = """
+      """
-infix fun <V> V.mustEqual(expected: V): Unit = assert(this == expected)
+      Wrapper(1 + 1) mustEqual (2 + 4)
-
+      """.trimIndent()
-fun <V> V.mustEqual(expected: V, message: () -> String): Unit =
-	assert(this == expected, message)
-
-fun main() {
-  $mainBody
-}
-""",
-      trimIndent = false
     )
+    assertEquals(
+      """
+      Wrapper(1 + 1) mustEqual (2 + 4)
+      |         |                 |
+      |         |                 6
+      |         2
+      Wrapper
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
 
-    val result = compile(listOf(file), PowerAssertComponentRegistrar(setOf(FqName("mustEqual"))))
+  @Test
-    assertEquals(KotlinCompilation.ExitCode.OK, result.exitCode)
+  fun `dispatch infix function call with constant receiver`() {
+    val actual = runDispatchInfix(
+      """
+      Wrapper(1) mustEqual (2 + 4)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Wrapper(1) mustEqual (2 + 4)
+      |                       |
+      |                       6
+      Wrapper
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `dispatch infix function call with constant parameter`() {
+    val actual = runDispatchInfix(
+      """
+      Wrapper(1 + 1) mustEqual 6
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Wrapper(1 + 1) mustEqual 6
+      |         |
+      |         2
+      Wrapper
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `dispatch infix function call with only constants`() {
+    val actual = runDispatchInfix(
+      """
+      Wrapper(2) mustEqual 6
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Wrapper(2) mustEqual 6
+      |
+      Wrapper
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `dispatch non-infix function call includes receiver`() {
+    val actual = runDispatchInfix(
+      """
+      Wrapper(1 + 1).mustEqual(2 + 4)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Wrapper(1 + 1).mustEqual(2 + 4)
+      |         |                |
+      |         |                6
+      |         2
+      Wrapper
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `dispatch non-infix function call with constant receiver`() {
+    val actual = runDispatchInfix(
+      """
+      Wrapper(1).mustEqual(2 + 4)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Wrapper(1).mustEqual(2 + 4)
+      |                      |
+      |                      6
+      Wrapper
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `dispatch non-infix function call with constant parameter`() {
+    val actual = runDispatchInfix(
+      """
+      Wrapper(1 + 1).mustEqual(6)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Wrapper(1 + 1).mustEqual(6)
+      |         |
+      |         2
+      Wrapper
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `dispatch non-infix function call with only constants`() {
+    val actual = runDispatchInfix(
+      """
+      Wrapper(2).mustEqual(6)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Wrapper(2).mustEqual(6)
+      |
+      Wrapper
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  private fun runExtensionInfix(mainBody: String): String {
+    return run(
+      SourceFile.kotlin(
+        name = "main.kt",
+        contents = """
+        infix fun <V> V.mustEqual(expected: V): Unit = assert(this == expected)
+        
+        fun <V> V.mustEqual(expected: V, message: () -> String): Unit =
+          assert(this == expected, message)
+        
+        fun main() {
+          $mainBody
+        }
+        """.trimIndent(),
+        trimIndent = false
+      ),
+      setOf(FqName("mustEqual"))
+    )
+  }
+
+  private fun runDispatchInfix(mainBody: String): String {
+    return run(
+      SourceFile.kotlin(
+        name = "main.kt",
+        contents = """
+        class Wrapper<V>(
+          private val value: V
+        ) {
+          infix fun mustEqual(expected: V): Unit = assert(value == expected)
+        
+          fun mustEqual(expected: V, message: () -> String): Unit =
+            assert(value == expected, message)
+        
+          override fun toString() = "Wrapper"
+        }
+        
+        fun main() {
+          $mainBody
+        }
+        """.trimIndent(),
+        trimIndent = false
+      ),
+      setOf(FqName("Wrapper.mustEqual"))
+    )
+  }
+
+  private fun run(file: SourceFile, fqNames: Set<FqName>): String {
+    val result = compile(listOf(file), PowerAssertComponentRegistrar(fqNames))
+    assertEquals(KotlinCompilation.ExitCode.OK, result.exitCode, "Failed with messages: " + result.messages)
 
     val kClazz = result.classLoader.loadClass("MainKt")
     val main = kClazz.declaredMethods.single { it.name == "main" && it.parameterCount == 0 }