Add renderAnnotations to ExtractableCodeDescriptor

The existing `ExtractFunctionDescriptorModifier::modifyDescriptor(..)`
that allows us to update the extraction descriptor returns
`ExtractableCodeDescriptor`. We have some cases like compose editor
of Android Studio that needs to change the behavior of the extraction
by updating the annotation rendering. The current annotation rendering
has a limitation that the rendered annotations must be from the
annotations of the container class. However, in the following example
the lambda argument must have `@Composable` annotation after extraction
even though its container function does not have it:
```
@Composable
fun myWidget(context: @Composable () -> Unit) {}

fun containerFunction() {
    myWidget { // This lambda argument is composable!
        <selection>print(true)</selection>
    }
}
```

Since we cannot handle the above case i.e., adding `@Composable` to the
extracted function for the lambda argument, this commit allows us to
use the rendered annotations instead of class Ids for
`ExtractableCodeDescriptor`.

^KTIJ-30695

Signed-off-by: Anna Kozlova <anna.kozlova@jetbrains.com>

GitOrigin-RevId: e4f2039882092544b67896f4f06a91ffee05d6d8

plugins/kotlin/idea/tests/testData/refactoring/extractFunctionModifier/composableFunction.kt

@@ -0,0 +1,10 @@
+package com.example
+
+@Target(AnnotationTarget.FUNCTION, AnnotationTarget.TYPE_PARAMETER,
+        AnnotationTarget.VALUE_PARAMETER)
+annotation class MyComposable
+
+@MyComposable
+fun sourceFunction() {
+    <selection>print(true)</selection>
+}

plugins/kotlin/idea/tests/testData/refactoring/extractFunctionModifier/composableFunction.kt.after

@@ -0,0 +1,15 @@
+package com.example
+
+@Target(AnnotationTarget.FUNCTION, AnnotationTarget.TYPE_PARAMETER,
+        AnnotationTarget.VALUE_PARAMETER)
+annotation class MyComposable
+
+@MyComposable
+fun sourceFunction() {
+    __dummyTestFun__()
+}
+
+@MyComposable
+private fun __dummyTestFun__() {
+    print(true)
+}

plugins/kotlin/idea/tests/testData/refactoring/extractFunctionModifier/composableLambda.kt

@@ -0,0 +1,14 @@
+package com.example
+
+@Target(AnnotationTarget.FUNCTION, AnnotationTarget.TYPE_PARAMETER,
+        AnnotationTarget.VALUE_PARAMETER)
+annotation class MyComposable
+
+@MyComposable
+fun myWidget(context: @MyComposable () -> Unit) {}
+
+fun setContent() {
+    myWidget {
+        <selection>print(true)</selection>
+    }
+}

plugins/kotlin/idea/tests/testData/refactoring/extractFunctionModifier/composableLambda.kt.after

@@ -0,0 +1,19 @@
+package com.example
+
+@Target(AnnotationTarget.FUNCTION, AnnotationTarget.TYPE_PARAMETER,
+        AnnotationTarget.VALUE_PARAMETER)
+annotation class MyComposable
+
+@MyComposable
+fun myWidget(context: @MyComposable () -> Unit) {}
+
+fun setContent() {
+    myWidget {
+        __dummyTestFun__()
+    }
+}
+
+@MyComposable
+private fun __dummyTestFun__() {
+    print(true)
+}

plugins/kotlin/refactorings/kotlin.refactorings.k2/src/org/jetbrains/kotlin/idea/k2/refactoring/extractFunction/ExtractableCodeDescriptor.kt

@@ -3,30 +3,14 @@ package org.jetbrains.kotlin.idea.k2.refactoring.extractFunction
 
 import com.intellij.psi.PsiElement
 import com.intellij.util.containers.MultiMap
-import org.jetbrains.kotlin.analysis.api.KaExperimentalApi
-import org.jetbrains.kotlin.analysis.api.KaSession
 import org.jetbrains.kotlin.analysis.api.analyze
-import org.jetbrains.kotlin.analysis.api.annotations.KaAnnotated
-import org.jetbrains.kotlin.analysis.api.annotations.KaAnnotation
-import org.jetbrains.kotlin.analysis.api.renderer.base.annotations.KaRendererAnnotationsFilter
-import org.jetbrains.kotlin.analysis.api.renderer.declarations.impl.KaDeclarationRendererForSource
 import org.jetbrains.kotlin.analysis.api.types.KaType
-import org.jetbrains.kotlin.analysis.utils.printer.PrettyPrinter
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.ControlFlow
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.DuplicateInfo
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.ExtractableCodeDescriptorWithConflictsResult
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.IExtractableCodeDescriptor
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.IExtractableCodeDescriptorWithConflicts
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.IReplacement
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.TypeParameter
+import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.*
 import org.jetbrains.kotlin.lexer.KtKeywordToken
 import org.jetbrains.kotlin.lexer.KtModifierKeywordToken
-import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
-import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 
 data class ExtractableCodeDescriptor(
     val context: KtElement,
@@ -41,7 +25,7 @@ data class ExtractableCodeDescriptor(
     override val returnType: KaType,
     override val modifiers: List<KtKeywordToken> = emptyList(),
     override val optInMarkers: List<FqName> = emptyList(),
-    val annotationClassIds: Set<ClassId> = emptySet()
+    val renderedAnnotations: List<String> = emptyList(),
 ) : IExtractableCodeDescriptor<KaType> {
     override val name: String get() = suggestedNames.firstOrNull() ?: ""
 
@@ -51,28 +35,11 @@ data class ExtractableCodeDescriptor(
 
     override fun isUnitReturnType(): Boolean = isUnitReturn
 
-    @OptIn(KaExperimentalApi::class)
     override val annotationsText: String
-        get() {
-            if (annotationClassIds.isEmpty()) return ""
-            val container = extractionData.commonParent.getStrictParentOfType<KtNamedFunction>() ?: return ""
-            return analyze(container) {
-                val filteredRenderer = KaDeclarationRendererForSource.WITH_QUALIFIED_NAMES.annotationRenderer.with {
-                    annotationFilter = annotationFilter.and(object : KaRendererAnnotationsFilter {
-                        override fun filter(
-                            analysisSession: KaSession,
-                            annotation: KaAnnotation,
-                            owner: KaAnnotated
-                        ): Boolean {
-                            return annotation.classId in annotationClassIds
-                        }
-                    })
-
-                }
-                val printer = PrettyPrinter()
-                filteredRenderer.renderAnnotations(useSiteSession, container.symbol, printer)
-                printer.toString() + "\n"
-            }
+        get() = if (renderedAnnotations.isEmpty()) {
+            ""
+        } else {
+            renderedAnnotations.joinToString(separator = "\n", postfix = "\n")
         }
 }
 

plugins/kotlin/refactorings/kotlin.refactorings.k2/src/org/jetbrains/kotlin/idea/k2/refactoring/extractFunction/ui/KotlinFirExtractFunctionDialog.java

@@ -303,6 +303,6 @@ public class KotlinFirExtractFunctionDialog extends DialogWrapper {
                                              returnType,
                                              originalDescriptor.getModifiers(),
                                              originalDescriptor.getOptInMarkers(),
-                                             originalDescriptor.getAnnotationClassIds());
+                                             originalDescriptor.getRenderedAnnotations());
     }
 }

plugins/kotlin/refactorings/kotlin.refactorings.k2/src/org/jetbrains/kotlin/idea/k2/refactoring/introduce/extractionEngine/ExtractionDataAnalyzer.kt

@@ -9,63 +9,30 @@ import com.intellij.psi.util.descendantsOfType
 import com.intellij.refactoring.util.RefactoringUIUtil
 import com.intellij.util.containers.MultiMap
 import org.jetbrains.kotlin.analysis.api.*
-import org.jetbrains.kotlin.analysis.api.annotations.*
+import org.jetbrains.kotlin.analysis.api.annotations.KaAnnotation
+import org.jetbrains.kotlin.analysis.api.annotations.KaAnnotationValue
 import org.jetbrains.kotlin.analysis.api.components.KaDataFlowExitPointSnapshot
 import org.jetbrains.kotlin.analysis.api.components.KaDiagnosticCheckerFilter
 import org.jetbrains.kotlin.analysis.api.fir.diagnostics.KaFirDiagnostic
 import org.jetbrains.kotlin.analysis.api.projectStructure.KaDanglingFileResolutionMode
+import org.jetbrains.kotlin.analysis.api.renderer.declarations.impl.KaDeclarationRendererForSource
 import org.jetbrains.kotlin.analysis.api.symbols.KaCallableSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KaAnnotatedSymbol
 import org.jetbrains.kotlin.analysis.api.types.KaClassType
 import org.jetbrains.kotlin.analysis.api.types.KaType
+import org.jetbrains.kotlin.analysis.utils.printer.PrettyPrinter
 import org.jetbrains.kotlin.idea.base.resources.KotlinBundle
 import org.jetbrains.kotlin.idea.base.util.names.FqNames
 import org.jetbrains.kotlin.idea.base.util.names.FqNames.OptInFqNames.isRequiresOptInFqName
-import org.jetbrains.kotlin.idea.k2.refactoring.extractFunction.ExtractFunctionDescriptorModifier
-import org.jetbrains.kotlin.idea.k2.refactoring.extractFunction.ExtractableCodeDescriptor
-import org.jetbrains.kotlin.idea.k2.refactoring.extractFunction.ExtractableCodeDescriptorWithConflicts
-import org.jetbrains.kotlin.idea.k2.refactoring.extractFunction.ExtractionData
-import org.jetbrains.kotlin.idea.k2.refactoring.extractFunction.ExtractionGeneratorConfiguration
-import org.jetbrains.kotlin.idea.k2.refactoring.extractFunction.ExtractionResult
-import org.jetbrains.kotlin.idea.k2.refactoring.extractFunction.MutableParameter
-import org.jetbrains.kotlin.idea.k2.refactoring.extractFunction.inferParametersInfo
-import org.jetbrains.kotlin.idea.k2.refactoring.extractFunction.targetKey
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.AbstractExtractionDataAnalyzer
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.ControlFlow
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.ExtractionGeneratorOptions
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.ExtractionTarget
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.IExtractableCodeDescriptor
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.IExtractionData
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.IReplacement
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.OutputDescriptor
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.ParametersInfo
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.ResolveResult
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.TypeDescriptor
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.TypeParameter
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.getDefaultVisibility
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.resolveResult
+import org.jetbrains.kotlin.idea.k2.refactoring.extractFunction.*
+import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.*
 import org.jetbrains.kotlin.idea.references.ReadWriteAccessChecker
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.lexer.KtModifierKeywordToken
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.psi.KtBlockExpression
-import org.jetbrains.kotlin.psi.KtBreakExpression
-import org.jetbrains.kotlin.psi.KtContinueExpression
-import org.jetbrains.kotlin.psi.KtDeclaration
-import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.KtNamedDeclaration
-import org.jetbrains.kotlin.psi.KtNamedFunction
-import org.jetbrains.kotlin.psi.KtParameter
-import org.jetbrains.kotlin.psi.KtProperty
-import org.jetbrains.kotlin.psi.KtReferenceExpression
-import org.jetbrains.kotlin.psi.KtReturnExpression
-import org.jetbrains.kotlin.psi.KtSimpleNameExpression
-import org.jetbrains.kotlin.psi.KtThisExpression
-import org.jetbrains.kotlin.psi.KtTreeVisitorVoid
-import org.jetbrains.kotlin.psi.KtTypeParameter
-import org.jetbrains.kotlin.psi.KtUserType
+import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 
 internal class ExtractionDataAnalyzer(private val extractionData: ExtractionData) :
@@ -222,7 +189,7 @@ internal class ExtractionDataAnalyzer(private val extractionData: ExtractionData
             returnType = returnType,
             modifiers = emptyList(),
             optInMarkers = experimentalMarkers.optInMarkers,
-            annotationClassIds = experimentalMarkers.propagatingMarkerClassIds
+            renderedAnnotations = getRenderedAnnotations(experimentalMarkers.propagatingMarkerClassIds)
         )
         val config = ExtractionGeneratorConfiguration(
             descriptor,
@@ -248,6 +215,24 @@ internal class ExtractionDataAnalyzer(private val extractionData: ExtractionData
         }
         return descriptor
     }
+
+    @OptIn(KaExperimentalApi::class)
+    private fun getRenderedAnnotations(annotationClassIds: Set<ClassId>): List<String> {
+        if (annotationClassIds.isEmpty()) return emptyList()
+        val container = extractionData.commonParent.getStrictParentOfType<KtNamedFunction>() ?: return emptyList()
+        return analyze(container) {
+            val owner = container.symbol
+            val renderer = KaDeclarationRendererForSource.WITH_QUALIFIED_NAMES.annotationRenderer
+            owner.annotations.filter { it.classId in annotationClassIds }.map { annotation ->
+                val printer = PrettyPrinter()
+                printer.append('@')
+                renderer.annotationUseSiteTargetRenderer.renderUseSiteTarget(useSiteSession, annotation, owner, renderer, printer)
+                renderer.annotationsQualifiedNameRenderer.renderQualifier(useSiteSession, annotation, owner, renderer, printer)
+                printer.append('\n')
+                printer.toString()
+            }
+        }
+    }
 }
 
 private data class ExperimentalMarkers(

plugins/kotlin/refactorings/kotlin.refactorings.tests.k2/test/org/jetbrains/kotlin/idea/k2/refactoring/introduce/AbstractK2IntroduceFunctionWithExtractFunctionModifierTest.kt

@@ -0,0 +1,54 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.kotlin.idea.k2.refactoring.introduce
+
+import com.intellij.psi.util.parentOfType
+import org.jetbrains.kotlin.analysis.api.analyze
+import org.jetbrains.kotlin.analysis.api.resolution.singleFunctionCallOrNull
+import org.jetbrains.kotlin.idea.k2.refactoring.extractFunction.ExtractFunctionDescriptorModifier
+import org.jetbrains.kotlin.idea.k2.refactoring.extractFunction.ExtractableCodeDescriptor
+import org.jetbrains.kotlin.idea.util.findAnnotation
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.KtAnnotated
+import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.KtLambdaArgument
+
+private val MY_COMPOSABLE_CLASS_ID = ClassId(FqName("com.example"), FqName("MyComposable"), false)
+
+abstract class AbstractK2IntroduceFunctionWithExtractFunctionModifierTest: AbstractK2IntroduceFunctionTest() {
+    override fun setUp() {
+        super.setUp()
+        ExtractFunctionDescriptorModifier.EP_NAME.point.registerExtension(
+            MyComposeExtractFunctionDescriptorModifier(), testRootDisposable
+        )
+    }
+}
+
+private class MyComposeExtractFunctionDescriptorModifier : ExtractFunctionDescriptorModifier {
+    private fun KtLambdaArgument.isComposable(): Boolean {
+        val callExpression = parent as KtCallExpression
+        val lambdaExpression = getLambdaExpression() ?: return false
+        analyze(callExpression) {
+            val call = callExpression.resolveToCall()?.singleFunctionCallOrNull() ?: return false
+            val parameterTypeForLambda = call.argumentMapping[lambdaExpression]?.returnType ?: return false
+            return parameterTypeForLambda.annotations.classIds.any { it == MY_COMPOSABLE_CLASS_ID }
+        }
+    }
+
+    override fun modifyDescriptor(descriptor: ExtractableCodeDescriptor): ExtractableCodeDescriptor {
+        val sourceFunction = descriptor.extractionData.targetSibling
+        if (sourceFunction is KtAnnotated) {
+            sourceFunction.findAnnotation(MY_COMPOSABLE_CLASS_ID)?.let {
+                return descriptor.copy(
+                    renderedAnnotations = descriptor.renderedAnnotations + "@${MY_COMPOSABLE_CLASS_ID.asFqNameString()}"
+                )
+            }
+        }
+        val outsideLambda = descriptor.extractionData.commonParent.parentOfType<KtLambdaArgument>(true) ?: return descriptor
+        return if (outsideLambda.isComposable()) {
+            descriptor.copy(renderedAnnotations = descriptor.renderedAnnotations + "@${MY_COMPOSABLE_CLASS_ID.asFqNameString()}")
+        } else {
+            descriptor
+        }
+    }
+}

plugins/kotlin/refactorings/kotlin.refactorings.tests.k2/test/org/jetbrains/kotlin/idea/k2/refactoring/introduce/K2IntroduceFunctionWithExtractFunctionModifierTestGenerated.java

@@ -0,0 +1,42 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+
+package org.jetbrains.kotlin.idea.k2.refactoring.introduce;
+
+import com.intellij.testFramework.TestDataPath;
+import org.jetbrains.kotlin.idea.base.plugin.KotlinPluginMode;
+import org.jetbrains.kotlin.idea.base.test.TestRoot;
+import org.jetbrains.kotlin.idea.test.JUnit3RunnerWithInners;
+import org.jetbrains.kotlin.idea.test.KotlinTestUtils;
+import org.jetbrains.kotlin.test.TestMetadata;
+import org.junit.runner.RunWith;
+
+/**
+ * This class is generated by {@link org.jetbrains.kotlin.testGenerator.generator.TestGenerator}.
+ * DO NOT MODIFY MANUALLY.
+ */
+@SuppressWarnings("all")
+@TestRoot("refactorings/kotlin.refactorings.tests.k2")
+@TestDataPath("$CONTENT_ROOT")
+@RunWith(JUnit3RunnerWithInners.class)
+@TestMetadata("../../idea/tests/testData/refactoring/extractFunctionModifier")
+public class K2IntroduceFunctionWithExtractFunctionModifierTestGenerated extends AbstractK2IntroduceFunctionWithExtractFunctionModifierTest {
+    @java.lang.Override
+    @org.jetbrains.annotations.NotNull
+    public final KotlinPluginMode getPluginMode() {
+        return KotlinPluginMode.K2;
+    }
+
+    private void runTest(String testDataFilePath) throws Exception {
+        KotlinTestUtils.runTest(this::doExtractFunctionTest, this, testDataFilePath);
+    }
+
+    @TestMetadata("composableFunction.kt")
+    public void testComposableFunction() throws Exception {
+        runTest("../../idea/tests/testData/refactoring/extractFunctionModifier/composableFunction.kt");
+    }
+
+    @TestMetadata("composableLambda.kt")
+    public void testComposableLambda() throws Exception {
+        runTest("../../idea/tests/testData/refactoring/extractFunctionModifier/composableLambda.kt");
+    }
+}

plugins/kotlin/util/test-generator-fir/test/org/jetbrains/kotlin/fir/testGenerator/GenerateK2RefactoringsTests.kt

@@ -82,6 +82,10 @@ internal fun MutableTWorkspace.generateK2RefactoringsTests() {
             model("refactoring/extractFunction", pattern = Patterns.KT_OR_KTS, testMethodName = "doExtractFunctionTest")
         }
 
+        testClass<AbstractK2IntroduceFunctionWithExtractFunctionModifierTest> {
+            model("refactoring/extractFunctionModifier", pattern = Patterns.KT_OR_KTS, testMethodName = "doExtractFunctionTest")
+        }
+
         testClass<AbstractK2InplaceIntroduceFunctionTest> {
             model("refactoring/extractFunctionInplace")
         }