From 10a422180a7c6495b5f7662c85ecac27a5aae76e Mon Sep 17 00:00:00 2001 From: Jaebaek Seo Date: Thu, 8 Feb 2024 16:44:09 -0800 Subject: [PATCH] K2: Propagate annotation of function type kind in compose plugin We want to test whether the propagation of function type kind annotation correctly works or not. This commit updates ComposableFunctionsTransformer that mimics the real compose compiler plugin to propagate the composable annotation i.e., @MyComposable from MyComposableFunction kind. Finally, the generated IR function for lambda expression must have the composable annotation, when the lambda expression does not have the explicit composable annotation, but we infer the composable annotation based on the annotation of function type for the value parameter. ^KT-64994 --- .../plugin/ComposableFunctionsTransformer.kt | 55 +++++++++++++++++-- .../box/composableFunction.fir.ir.txt | 8 +++ 2 files changed, 58 insertions(+), 5 deletions(-) diff --git a/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/ir/plugin/ComposableFunctionsTransformer.kt b/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/ir/plugin/ComposableFunctionsTransformer.kt index a5dec140c87..3f2834eca00 100644 --- a/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/ir/plugin/ComposableFunctionsTransformer.kt +++ b/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/ir/plugin/ComposableFunctionsTransformer.kt @@ -6,22 +6,28 @@ package org.jetbrains.kotlin.ir.plugin import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext -import org.jetbrains.kotlin.builtins.StandardNames import org.jetbrains.kotlin.builtins.functions.FunctionTypeKind import org.jetbrains.kotlin.fir.plugin.types.ComposableNames.FULL_COMPOSABLE_NAME_PREFIX import org.jetbrains.kotlin.ir.IrElement -import org.jetbrains.kotlin.ir.declarations.* -import org.jetbrains.kotlin.ir.expressions.* +import org.jetbrains.kotlin.ir.declarations.IrClass +import org.jetbrains.kotlin.ir.declarations.IrFunction +import org.jetbrains.kotlin.ir.declarations.IrValueParameter +import org.jetbrains.kotlin.ir.declarations.IrVariable +import org.jetbrains.kotlin.ir.expressions.IrCall +import org.jetbrains.kotlin.ir.expressions.IrExpression +import org.jetbrains.kotlin.ir.expressions.IrFunctionExpression +import org.jetbrains.kotlin.ir.expressions.IrFunctionReference +import org.jetbrains.kotlin.ir.expressions.impl.IrConstructorCallImpl import org.jetbrains.kotlin.ir.symbols.IrClassifierSymbol import org.jetbrains.kotlin.ir.types.* import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl import org.jetbrains.kotlin.ir.types.impl.IrTypeBase import org.jetbrains.kotlin.ir.types.impl.IrTypeProjectionImpl -import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable -import org.jetbrains.kotlin.ir.util.functions +import org.jetbrains.kotlin.ir.util.* import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid import org.jetbrains.kotlin.name.ClassId +import org.jetbrains.kotlin.name.FqName import org.jetbrains.kotlin.name.Name class ComposableFunctionsTransformer(val pluginContext: IrPluginContext) : IrElementVisitorVoid { @@ -66,6 +72,45 @@ class ComposableFunctionsTransformer(val pluginContext: IrPluginContext) : IrEle visitElement(expression) } + /** + * This function propagates the special function type kind for composable to function expressions like lambda expression. + */ + override fun visitFunctionExpression(expression: IrFunctionExpression) { + if (expression.type.isSyntheticComposableFunction()) { + expression.function.mark() + } + super.visitFunctionExpression(expression) + } + + /** + * This function propagates the special function type kind for composable to function references. + */ + override fun visitFunctionReference(expression: IrFunctionReference) { + if (expression.type.isSyntheticComposableFunction()) { + expression.symbol.owner.mark() + } + super.visitFunctionReference(expression) + } + + private fun IrType.isSyntheticComposableFunction() = + classOrNull?.owner?.let { + it.name.asString().startsWith("MyComposableFunction") && + it.packageFqName?.asString() == "some" + } ?: false + + private val composableClassId = ClassId(FqName("org.jetbrains.kotlin.fir.plugin"), FqName("MyComposable"), false) + + private val composableSymbol = pluginContext.referenceClass(composableClassId)!! + + private fun IrFunction.mark() { + if (!hasAnnotation(composableClassId)) { + annotations = annotations + IrConstructorCallImpl.fromSymbolOwner( + composableSymbol.owner.defaultType, + composableSymbol.constructors.single(), + ) + } + } + private fun updateReferenceInCallIfNeeded(call: IrCall) { val function = call.symbol.owner if (function.name != INVOKE) return diff --git a/plugins/fir-plugin-prototype/testData/box/composableFunction.fir.ir.txt b/plugins/fir-plugin-prototype/testData/box/composableFunction.fir.ir.txt index e83d5079f88..5e0c8bc0445 100644 --- a/plugins/fir-plugin-prototype/testData/box/composableFunction.fir.ir.txt +++ b/plugins/fir-plugin-prototype/testData/box/composableFunction.fir.ir.txt @@ -64,6 +64,8 @@ FILE fqName: fileName:/composableFunction.kt VAR name:l1 type:kotlin.Function0 [val] FUN_EXPR type=kotlin.Function0 origin=LAMBDA FUN LOCAL_FUNCTION_FOR_LAMBDA name: visibility:local modality:FINAL <> () returnType:kotlin.Unit + annotations: + MyComposable BLOCK_BODY TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit BLOCK type=kotlin.Int origin=POSTFIX_INCR @@ -77,6 +79,8 @@ FILE fqName: fileName:/composableFunction.kt VAR name:l2 type:@[MyComposable] kotlin.Function0 [val] FUN_EXPR type=kotlin.Function0 origin=LAMBDA FUN LOCAL_FUNCTION_FOR_LAMBDA name: visibility:local modality:FINAL <> () returnType:kotlin.Unit + annotations: + MyComposable BLOCK_BODY TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit BLOCK type=kotlin.Int origin=POSTFIX_INCR @@ -126,6 +130,8 @@ FILE fqName: fileName:/composableFunction.kt CALL 'public final fun runComposable (block: @[MyComposable] kotlin.Function0): kotlin.Unit declared in ' type=kotlin.Unit origin=null block: FUN_EXPR type=kotlin.Function0 origin=LAMBDA FUN LOCAL_FUNCTION_FOR_LAMBDA name: visibility:local modality:FINAL <> () returnType:kotlin.Unit + annotations: + MyComposable BLOCK_BODY TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit BLOCK type=kotlin.Int origin=POSTFIX_INCR @@ -157,6 +163,8 @@ FILE fqName: fileName:/composableFunction.kt CALL 'public final fun runComposable2 (block: kotlin.Function1): kotlin.String declared in ' type=kotlin.String origin=null block: FUN_EXPR type=kotlin.Function1 origin=LAMBDA FUN LOCAL_FUNCTION_FOR_LAMBDA name: visibility:local modality:FINAL <> (it:kotlin.String) returnType:kotlin.String + annotations: + MyComposable VALUE_PARAMETER name:it index:0 type:kotlin.String BLOCK_BODY RETURN type=kotlin.Nothing from='local final fun (it: kotlin.String): kotlin.String declared in .test_2'