From 3d0919b3f332958729acb051f8d3b85b2f622b85 Mon Sep 17 00:00:00 2001 From: Brian Norman Date: Mon, 1 Aug 2022 22:28:08 -0500 Subject: [PATCH] Support dispatch infix functions --- .../bnorm/power/PowerAssertCallTransformer.kt | 36 ++- .../bnorm/power/delegate/FunctionDelegate.kt | 4 +- .../power/delegate/LambdaFunctionDelegate.kt | 10 +- .../SamConversionLambdaFunctionDelegate.kt | 12 +- .../power/delegate/SimpleFunctionDelegate.kt | 10 +- .../com/bnorm/power/diagram/DiagramBuilder.kt | 11 +- .../com/bnorm/power/InfixFunctionTest.kt | 239 +++++++++++++++--- 7 files changed, 270 insertions(+), 52 deletions(-) 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 e4027a1004c..7a18ad0f8cf 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 @@ -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, + dispatchRoot: Node? = null, extensionRoot: Node? = null ): IrExpression { fun recursive( index: Int, - extension: IrExpression? = null, + dispatch: IrExpression?, + extension: IrExpression?, arguments: List, variables: List ): 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) { - buildDiagramNesting(extensionRoot) { extension, newVariables -> - recursive(0, extension, emptyList(), newVariables) + return buildDiagramNestingNullable(dispatchRoot) { dispatch, newVariables -> + buildDiagramNestingNullable(extensionRoot, newVariables) { extension, 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 } 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 index ec686f158b1..8a806a5d629 100644 --- 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 @@ -33,6 +33,7 @@ interface FunctionDelegate { fun buildCall( builder: IrBuilderWithScope, original: IrCall, + dispatchReceiver: IrExpression?, extensionReceiver: IrExpression?, valueArguments: List, messageArgument: IrExpression @@ -41,12 +42,13 @@ interface FunctionDelegate { fun IrBuilderWithScope.irCallCopy( overload: IrSimpleFunctionSymbol, original: IrCall, + dispatchReceiver: IrExpression?, extensionReceiver: IrExpression?, valueArguments: List, 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)) 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 index f099d325095..dccc644291f 100644 --- 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 @@ -33,6 +33,7 @@ class LambdaFunctionDelegate( override fun buildCall( builder: IrBuilderWithScope, original: IrCall, + dispatchReceiver: IrExpression?, extensionReceiver: IrExpression?, valueArguments: List, 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 + ) } } diff --git a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/SamConversionLambdaFunctionDelegate.kt b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/SamConversionLambdaFunctionDelegate.kt index 6bd0ae36f4e..902fedc9d8e 100644 --- a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/SamConversionLambdaFunctionDelegate.kt +++ b/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, 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 + ) } } 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 index e576e5022d7..b1b9221c5cd 100644 --- 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 @@ -31,8 +31,16 @@ class SimpleFunctionDelegate( override fun buildCall( builder: IrBuilderWithScope, original: IrCall, + dispatchReceiver: IrExpression?, extensionReceiver: IrExpression?, valueArguments: List, 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 + ) } diff --git a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/DiagramBuilder.kt b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/DiagramBuilder.kt index 9b6f0c757e0..70f38933972 100644 --- a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/DiagramBuilder.kt +++ b/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 = emptyList(), call: IrBuilderWithScope.(IrExpression, List) -> IrExpression ): IrExpression { - return buildExpression(root, listOf()) { argument, subStack -> + return buildExpression(root, variables) { argument, subStack -> call(argument, subStack) } } +fun IrBuilderWithScope.buildDiagramNestingNullable( + root: Node?, + variables: List = emptyList(), + call: IrBuilderWithScope.(IrExpression?, List) -> IrExpression +): IrExpression { + return if (root != null) buildDiagramNesting(root, variables, call) else call(null, variables) +} + private fun IrBuilderWithScope.buildExpression( node: Node, variables: List, diff --git a/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/InfixFunctionTest.kt b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/InfixFunctionTest.kt index 88cd3dfa82a..2adc69020c1 100644 --- a/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/InfixFunctionTest.kt +++ b/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`() { - val actual = execute( + fun `extension infix function call includes receiver`() { + val actual = runExtensionInfix( """ (1 + 1) mustEqual (2 + 4) """.trimIndent() @@ -44,8 +44,8 @@ class InfixFunctionTest { } @Test - fun `infix function call with constant receiver`() { - val actual = execute( + fun `extension infix function call with constant receiver`() { + val actual = runExtensionInfix( """ 1 mustEqual (2 + 4) """.trimIndent() @@ -61,8 +61,8 @@ class InfixFunctionTest { } @Test - fun `infix function call with constant parameter`() { - val actual = execute( + fun `extension infix function call with constant parameter`() { + val actual = runExtensionInfix( """ (1 + 1) mustEqual 6 """.trimIndent() @@ -78,8 +78,8 @@ class InfixFunctionTest { } @Test - fun `infix function call with only constants`() { - val actual = execute( + fun `extension infix function call with only constants`() { + val actual = runExtensionInfix( """ 2 mustEqual 6 """.trimIndent() @@ -93,8 +93,8 @@ class InfixFunctionTest { } @Test - fun `non-infix function call includes receiver`() { - val actual = execute( + fun `extension non-infix function call includes receiver`() { + val actual = runExtensionInfix( """ (1 + 1).mustEqual(2 + 4) """.trimIndent() @@ -111,8 +111,8 @@ class InfixFunctionTest { } @Test - fun `non-infix function call with constant receiver`() { - val actual = execute( + fun `extension non-infix function call with constant receiver`() { + val actual = runExtensionInfix( """ 1.mustEqual(2 + 4) """.trimIndent() @@ -128,8 +128,8 @@ class InfixFunctionTest { } @Test - fun `non-infix function call with constant parameter`() { - val actual = execute( + fun `extension non-infix function call with constant parameter`() { + val actual = runExtensionInfix( """ (1 + 1).mustEqual(6) """.trimIndent() @@ -145,8 +145,8 @@ class InfixFunctionTest { } @Test - fun `non-infix function call with only constants`() { - val actual = execute( + fun `extension non-infix function call with only constants`() { + val actual = runExtensionInfix( """ 2.mustEqual(6) """.trimIndent() @@ -159,24 +159,199 @@ class InfixFunctionTest { ) } - private fun execute(mainBody: String): String { - val file = SourceFile.kotlin( - name = "main.kt", - contents = """ -infix fun V.mustEqual(expected: V): Unit = assert(this == expected) - -fun V.mustEqual(expected: V, message: () -> String): Unit = - assert(this == expected, message) - -fun main() { - $mainBody -} -""", - trimIndent = false + @Test + fun `dispatch 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() + ) + } - val result = compile(listOf(file), PowerAssertComponentRegistrar(setOf(FqName("mustEqual")))) - assertEquals(KotlinCompilation.ExitCode.OK, result.exitCode) + @Test + 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.mustEqual(expected: V): Unit = assert(this == expected) + + fun 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( + 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): 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 }