Support dispatch infix functions

This commit is contained in:
Brian Norman
2022-08-01 22:28:08 -05:00
parent b5d8143f4c
commit 3d0919b3f3
7 changed files with 270 additions and 52 deletions
@@ -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) {
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
}
@@ -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))
@@ -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
)
}
}
@@ -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
)
}
}
@@ -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
)
}
@@ -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>,
@@ -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> 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 = 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> 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 }