[KMT-523][Compose] Add and improve support for maintaining @Composable when extracting a function

Adopt the feature from Android Studio, port it to K2, and add support for a bunch of cases where it produced red code previously.

https://youtrack.jetbrains.com/issue/KMT-523

GitOrigin-RevId: b78afe1af1558f1a738852c0580e674104fab5f7

plugins/compose/intellij.compose.ide.plugin.k2/BUILD.bazel

@@ -29,6 +29,8 @@ jvm_library(
     "//plugins/kotlin/highlighting/highlighting-k2:kotlin-highlighting-k2",
     "//platform/util",
     "//platform/analysis-impl",
+    "//plugins/kotlin/refactorings/kotlin.refactorings.common:kotlin-refactorings-common",
+    "//plugins/kotlin/refactorings/kotlin.refactorings.k2:kotlin-refactorings-k2",
     "@lib//:kotlinc-kotlin-compiler-common-provided",
     "@lib//:kotlinc-analysis-api-provided",
     "//plugins/kotlin/base/analysis-api/analysis-api-utils:kotlin-base-analysis-api-utils",
@@ -58,7 +60,12 @@ jvm_library(
   deps = [
     "//plugins/compose/intellij.compose.ide.plugin.shared:ide-plugin-shared",
     "//plugins/compose/intellij.compose.ide.plugin.shared:ide-plugin-shared-tests_test_lib",
+    "//java/testFramework",
+    "//plugins/kotlin/refactorings/kotlin.refactorings.k2:kotlin-refactorings-k2",
+    "//plugins/kotlin/refactorings/kotlin.refactorings.common:kotlin-refactorings-common",
+    "//platform/refactoring",
     "//platform/core-api:core",
+    "//platform/core-ui",
     "//platform/editor-ui-api:editor-ui",
     "//plugins/kotlin/base/test:test_test_lib",
     "//plugins/kotlin/base/plugin",

plugins/compose/intellij.compose.ide.plugin.k2/intellij.compose.ide.plugin.k2.iml

@@ -33,6 +33,8 @@
     <orderEntry type="module" module-name="kotlin.highlighting.k2" />
     <orderEntry type="module" module-name="intellij.platform.util" />
     <orderEntry type="module" module-name="intellij.platform.analysis.impl" />
+    <orderEntry type="module" module-name="kotlin.refactorings.common" />
+    <orderEntry type="module" module-name="kotlin.refactorings.k2" />
     <orderEntry type="library" scope="PROVIDED" name="kotlinc.kotlin-compiler-common" level="project" />
     <orderEntry type="library" scope="PROVIDED" name="kotlinc.analysis-api" level="project" />
     <orderEntry type="module" module-name="kotlin.base.analysis-api.utils" />

plugins/compose/intellij.compose.ide.plugin.k2/intellij.compose.ide.plugin.k2.tests.iml

@@ -26,7 +26,12 @@
     <orderEntry type="module" module-name="intellij.compose.ide.plugin.k2" scope="TEST" />
     <orderEntry type="module" module-name="intellij.compose.ide.plugin.shared" scope="TEST" />
     <orderEntry type="module" module-name="intellij.compose.ide.plugin.shared.tests" scope="TEST" />
+    <orderEntry type="module" module-name="intellij.java.testFramework" scope="TEST" />
+    <orderEntry type="module" module-name="kotlin.refactorings.k2" scope="TEST" />
+    <orderEntry type="module" module-name="kotlin.refactorings.common" scope="TEST" />
+    <orderEntry type="module" module-name="intellij.platform.refactoring" scope="TEST" />
     <orderEntry type="module" module-name="intellij.platform.core" scope="TEST" />
+    <orderEntry type="module" module-name="intellij.platform.core.ui" scope="TEST" />
     <orderEntry type="module" module-name="intellij.platform.editor.ui" scope="TEST" />
     <orderEntry type="module" module-name="kotlin.base.test" scope="TEST" />
     <orderEntry type="module" module-name="kotlin.base.plugin" scope="TEST" />

plugins/compose/intellij.compose.ide.plugin.k2/resources/intellij.compose.ide.plugin.k2.xml

@@ -1,12 +1,15 @@
 <idea-plugin package="com.intellij.compose.ide.plugin.k2">
   <dependencies>
     <plugin id="com.intellij.modules.kotlin.k2"/>
+    <module name="kotlin.refactorings.k2"/>
     <module name="kotlin.highlighting.k2"/>
     <module name="intellij.compose.ide.plugin.shared" />
   </dependencies>
 
   <extensions defaultExtensionNs="org.jetbrains.kotlin">
     <callHighlighterExtension implementation="com.intellij.compose.ide.plugin.k2.highlighting.ComposableFunctionCallHighlighterExtension"/>
+    <extractFunctionDescriptorModifier implementation="com.intellij.compose.ide.plugin.k2.K2ComposableAnnotationToExtractedFunctionAdder"
+                                       order="last"/>
   </extensions>
 
   <extensions defaultExtensionNs="com.intellij">

plugins/compose/intellij.compose.ide.plugin.k2/src/com/intellij/compose/ide/plugin/k2/K2ComposableAnnotationToExtractedFunctionAdder.kt

@@ -0,0 +1,20 @@
+package com.intellij.compose.ide.plugin.k2
+
+import com.intellij.compose.ide.plugin.shared.COMPOSABLE_ANNOTATION_FQ_NAME
+import com.intellij.compose.ide.plugin.shared.isFunctionExtractionInComposableControlFlow
+import org.jetbrains.kotlin.idea.k2.refactoring.extractFunction.ExtractFunctionDescriptorModifier
+import org.jetbrains.kotlin.idea.k2.refactoring.extractFunction.ExtractableCodeDescriptor
+
+/**
+ * Responsible for adding `@Composable` annotation to functions extracted from inside Composable control flow.
+ *
+ * For details see: [com.intellij.compose.ide.plugin.shared.isFunctionExtractionInComposableControlFlow]
+ */
+internal class K2ComposableAnnotationToExtractedFunctionAdder : ExtractFunctionDescriptorModifier {
+  override fun modifyDescriptor(descriptor: ExtractableCodeDescriptor): ExtractableCodeDescriptor =
+    if (composableAnnotationText !in descriptor.annotationsText && descriptor.extractionData.isFunctionExtractionInComposableControlFlow())
+      descriptor.copy(renderedAnnotations = descriptor.renderedAnnotations + "$composableAnnotationText ")
+    else descriptor
+}
+
+private val composableAnnotationText = "@${COMPOSABLE_ANNOTATION_FQ_NAME.asString()}"

plugins/compose/intellij.compose.ide.plugin.k2/tests/com/intellij/compose/ide/plugin/k2/K2ComposableAnnotationToExtractedFunctionAdderTest.kt

@@ -0,0 +1,111 @@
+package com.intellij.compose.ide.plugin.k2
+
+import com.intellij.compose.ide.plugin.shared.ComposableAnnotationToExtractedFunctionAdderTest
+import com.intellij.openapi.application.Application
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.project.Project
+import com.intellij.testFramework.fixtures.JavaCodeInsightTestFixture
+import org.jetbrains.kotlin.analysis.api.permissions.KaAllowAnalysisOnEdt
+import org.jetbrains.kotlin.analysis.api.permissions.allowAnalysisOnEdt
+import org.jetbrains.kotlin.idea.base.plugin.KotlinPluginMode
+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.ExtractionGeneratorConfiguration
+import org.jetbrains.kotlin.idea.k2.refactoring.extractFunction.ExtractionResult
+import org.jetbrains.kotlin.idea.k2.refactoring.extractFunction.KotlinFirExtractFunctionHandler
+import org.jetbrains.kotlin.idea.k2.refactoring.introduce.extractionEngine.ExtractionEngineHelper
+import org.jetbrains.kotlin.idea.k2.refactoring.introduceConstant.INTRODUCE_CONSTANT
+import org.jetbrains.kotlin.idea.k2.refactoring.introduceConstant.KotlinIntroduceConstantHandler
+import org.jetbrains.kotlin.idea.refactoring.introduce.extractFunction.EXTRACT_FUNCTION
+import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.ExtractionGeneratorOptions
+import org.jetbrains.kotlin.name.FqName
+import java.util.concurrent.Semaphore
+
+class K2ComposableAnnotationToExtractedFunctionAdderTest : ComposableAnnotationToExtractedFunctionAdderTest() {
+  override val pluginMode: KotlinPluginMode get() = KotlinPluginMode.K2
+
+  override fun JavaCodeInsightTestFixture.invokeExtractFunctionIn(application: Application, existingAnnotationFqNames: List<FqName>) {
+    val helper = ExtractionHelper(existingAnnotationFqNames)
+    application.invokeAndWait {
+      KotlinFirExtractFunctionHandler(helper = helper)
+        .invoke(this.project, this.editor, this.file!!, null)
+    }
+    helper.waitUntilFinished()
+  }
+
+  override fun JavaCodeInsightTestFixture.invokeExtractConstantIn(application: Application) {
+    val helper = InteractiveExtractionHelper()
+    application.invokeAndWait {
+      KotlinIntroduceConstantHandler(helper = helper)
+        .invoke(this.project, this.editor, this.file!!, null)
+    }
+    helper.waitUntilFinished()
+  }
+
+}
+
+// Following 2 are heavily inspired/copied from AOSP
+private class ExtractionHelper(private val existingAnnotationFqNames: List<FqName>) : ExtractionEngineHelper(EXTRACT_FUNCTION) {
+  private val finishedSemaphore = Semaphore(0)
+
+  fun waitUntilFinished() {
+    finishedSemaphore.acquire()
+  }
+
+  @OptIn(KaAllowAnalysisOnEdt::class)
+  override fun configureAndRun(
+    project: Project,
+    editor: Editor,
+    descriptorWithConflicts: ExtractableCodeDescriptorWithConflicts,
+    onFinish: (ExtractionResult) -> Unit,
+  ) {
+    // `descriptorWithConflicts.descriptor.copy(..)` runs the constructor of
+    // ExtractableCodeDescriptor
+    // that calls AA for return type check. Since the copied one has the exactly same type as the
+    // initial descriptor, if we only copy the boolean value, we do not need `allowAnalysisOnEdt`.
+    // To do so, we have to update `ExtractableCodeDescriptor`.
+    allowAnalysisOnEdt {
+      val newDescriptor =
+        descriptorWithConflicts.descriptor.copy(
+          suggestedNames = listOf("newFunction"),
+          renderedAnnotations = existingAnnotationFqNames.map { "@${it.asString()}" } +
+                                descriptorWithConflicts.descriptor.renderedAnnotations
+        )
+      doRefactor(
+        ExtractionGeneratorConfiguration(newDescriptor, ExtractionGeneratorOptions.DEFAULT)
+      ) { er: ExtractionResult ->
+        onFinish(er)
+        finishedSemaphore.release()
+      }
+    }
+  }
+}
+
+private class InteractiveExtractionHelper : ExtractionEngineHelper(INTRODUCE_CONSTANT) {
+  private val finishedSemaphore = Semaphore(0)
+
+  fun waitUntilFinished() {
+    finishedSemaphore.acquire()
+  }
+
+  override fun validate(
+    descriptor: ExtractableCodeDescriptor
+  ): ExtractableCodeDescriptorWithConflicts =
+    KotlinIntroduceConstantHandler.InteractiveExtractionHelper.validate(descriptor)
+
+  override fun configureAndRun(
+    project: Project,
+    editor: Editor,
+    descriptorWithConflicts: ExtractableCodeDescriptorWithConflicts,
+    onFinish: (ExtractionResult) -> Unit,
+  ) {
+    KotlinIntroduceConstantHandler.InteractiveExtractionHelper.configureAndRun(
+      project,
+      editor,
+      descriptorWithConflicts,
+    ) { er: ExtractionResult ->
+      onFinish(er)
+      finishedSemaphore.release()
+    }
+  }
+}

plugins/compose/intellij.compose.ide.plugin.shared/BUILD.bazel

@@ -1,6 +1,12 @@
 ### auto-generated section `build intellij.compose.ide.plugin.shared` start
+load("//build:compiler-options.bzl", "create_kotlinc_options")
 load("@rules_jvm//:jvm.bzl", "jvm_library", "jvm_resources")
 
+create_kotlinc_options(
+  name = "custom",
+  context_receivers = True
+)
+
 jvm_resources(
   name = "ide-plugin-shared_resources",
   files = glob(["resources/**/*"]),
@@ -12,6 +18,7 @@ jvm_library(
   module_name = "intellij.compose.ide.plugin.shared",
   visibility = ["//visibility:public"],
   srcs = glob(["src/**/*.kt", "src/**/*.java"], allow_empty = True),
+  kotlinc_opts = ":custom",
   deps = [
     "//plugins/compose:compose-ide-plugin",
     "//platform/analysis-api:analysis",
@@ -28,6 +35,7 @@ jvm_library(
     "//uast/uast-common:uast",
     "//platform/util",
     "//plugins/kotlin/base/indices",
+    "//plugins/kotlin/refactorings/kotlin.refactorings.common:kotlin-refactorings-common",
     "@lib//:kotlinc-kotlin-compiler-common-provided",
     "@lib//:kotlinc-analysis-api-provided",
     "//java/java-psi-api:psi",

plugins/compose/intellij.compose.ide.plugin.shared/intellij.compose.ide.plugin.shared.iml

@@ -1,5 +1,21 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <module type="JAVA_MODULE" version="4">
+  <component name="FacetManager">
+    <facet type="kotlin-language" name="Kotlin">
+      <configuration version="5" platform="JVM 17" allPlatforms="JVM [17]" useProjectSettings="false">
+        <compilerSettings>
+          <option name="additionalArguments" value="-Xjvm-default=all -Xcontext-receivers" />
+        </compilerSettings>
+        <compilerArguments>
+          <stringArguments>
+            <stringArg name="jvmTarget" arg="17" />
+            <stringArg name="apiVersion" arg="2.2" />
+            <stringArg name="languageVersion" arg="2.2" />
+          </stringArguments>
+        </compilerArguments>
+      </configuration>
+    </facet>
+  </component>
   <component name="NewModuleRootManager" inherit-compiler-output="true">
     <exclude-output />
     <content url="file://$MODULE_DIR$">
@@ -23,6 +39,7 @@
     <orderEntry type="module" module-name="intellij.platform.uast" />
     <orderEntry type="module" module-name="intellij.platform.util" />
     <orderEntry type="module" module-name="kotlin.base.indices" />
+    <orderEntry type="module" module-name="kotlin.refactorings.common" />
     <orderEntry type="library" scope="PROVIDED" name="kotlinc.kotlin-compiler-common" level="project" />
     <orderEntry type="library" scope="PROVIDED" name="kotlinc.analysis-api" level="project" />
     <orderEntry type="module" module-name="intellij.java.psi" />

plugins/compose/intellij.compose.ide.plugin.shared/src/com/intellij/compose/ide/plugin/shared/ComposableAnnotationToExtractedFunctionAdderUtil.kt

@@ -0,0 +1,15 @@
+package com.intellij.compose.ide.plugin.shared
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.annotations.ApiStatus
+import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.IExtractionData
+
+/** Checks whether [this] represents the intention to extract a function inside a Composable control flow */
+@ApiStatus.Internal
+fun IExtractionData.isFunctionExtractionInComposableControlFlow(): Boolean =
+  !options.extractAsProperty // We only care for function extractions!
+  && composeIsEnabledInModuleOf(commonParent) // Cheap check for performance
+  && commonParent.isInsideComposableControlFlow()
+
+private fun composeIsEnabledInModuleOf(element: PsiElement): Boolean =
+  element.module.let { it != null && isComposeEnabledInModule(it) }

plugins/compose/intellij.compose.ide.plugin.shared/src/com/intellij/compose/ide/plugin/shared/ComposeLibraryNamespace.kt

@@ -11,6 +11,7 @@ val COMPOSABLE_ANNOTATION_CLASS_ID: ClassId = ClassId.topLevel(COMPOSABLE_ANNOTA
 val COMPOSE_MODIFIER_NAME: Name = Name.identifier("Modifier")
 val COMPOSE_MODIFIER_FQN: FqName = FqName("androidx.compose.ui.$COMPOSE_MODIFIER_NAME")
 val COMPOSE_MODIFIER_CLASS_ID: ClassId = ClassId.topLevel(COMPOSE_MODIFIER_FQN)
+val DISALLOW_COMPOSABLE_CALLS_FQ_NAME: FqName = FqName("androidx.compose.runtime.DisallowComposableCalls")
 
 val PREVIEW_CLASS_NAME: Name = Name.identifier("Preview")
 val PREVIEW_PARAMETER_CLASS_NAME: Name = Name.identifier("PreviewParameter")

plugins/compose/intellij.compose.ide.plugin.shared/src/com/intellij/compose/ide/plugin/shared/IsComposableControlFlowUtil.kt

@@ -0,0 +1,96 @@
+package com.intellij.compose.ide.plugin.shared
+
+import com.intellij.psi.PsiElement
+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.resolution.KaCallInfo
+import org.jetbrains.kotlin.analysis.api.resolution.singleConstructorCallOrNull
+import org.jetbrains.kotlin.analysis.api.resolution.successfulFunctionCallOrNull
+import org.jetbrains.kotlin.analysis.api.resolution.symbol
+import org.jetbrains.kotlin.analysis.api.symbols.KaNamedFunctionSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KaValueParameterSymbol
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getAnnotationEntries
+
+/**
+ * Checks whether [this] is inside a Composable control flow.
+ * Examples:
+ * ```kt
+ * fun bar() {
+ *     somePsiElement // <-- Will be false for this element. Because it would be a compiler error to call @Composable functions here.
+ * }
+ *
+ * @Composable fun foo() {
+ *     somePsiElement // <-- Will be true for this element. Because it's legal to call @Composable functions here.
+ * }
+ * ```
+ * For detailed behavior, see: [com.intellij.compose.ide.plugin.shared.ComposableAnnotationToExtractedFunctionAddingAnalyserTest]
+ */
+internal fun PsiElement?.isInsideComposableControlFlow(): Boolean = when (this) {
+  null -> false // Reached root parent
+  is KtPropertyAccessor, is KtNamedFunction -> hasComposableAnnotation()
+  is KtLambdaExpression -> analyze(this) { ownsComposableControlFlow() }
+  is KtLambdaArgument -> analyze(this) { ownsComposableControlFlow() }
+  else -> parent.isInsideComposableControlFlow()
+}
+
+context(KaSession)
+private fun KtLambdaExpression.ownsComposableControlFlow(): Boolean =
+  (parent as? KtLambdaArgument)?.ownsComposableControlFlow() // Only check lambda that is NOT an argument
+  // One would simply expect us to infer the type and then check whether the type has a `@Composable` annotation.
+  // But currently, annotations are not inferred. See https://jetbrains.slack.com/archives/C061DS4G41J/p1738180610032509
+  ?: getAnnotationEntries().containsAnnotationNamed(COMPOSABLE_ANNOTATION_FQ_NAME)
+
+context(KaSession)
+private fun KtLambdaArgument.ownsComposableControlFlow(): Boolean =
+  (parent as? KtCallExpression)
+    ?.resolveToCall()
+    .let { functionCall -> functionCall != null && this.ownsComposableControlFlowWhenArgumentOf(functionCall) }
+
+context(KaSession)
+private fun KtLambdaArgument.ownsComposableControlFlowWhenArgumentOf(functionCall: KaCallInfo): Boolean =
+  this.hasAnnotationIn(functionCall, COMPOSABLE_ANNOTATION_FQ_NAME) || // We are explicitly in Composable flow!
+  (
+    !this.hasAnnotationIn(functionCall, DISALLOW_COMPOSABLE_CALLS_FQ_NAME) // If our flow is not explicitly forbidding Composable calls,
+    && this.isInlinedInside(functionCall) // and we are in an inlined lambda,
+    && parent.isInsideComposableControlFlow() // then check whether the parent is in Composable flow
+  )
+
+private fun KtLambdaArgument.hasAnnotationIn(function: KaCallInfo, fqName: FqName): Boolean = analyze(this) {
+  function.parameterSymbolOf(this@hasAnnotationIn)
+    ?.returnType
+    ?.isAnnotatedWith(fqName) // Not cached, but doesn't matter for this infrequently triggered extension
+    ?: false
+}
+
+private fun KaAnnotated.isAnnotatedWith(fqName: FqName): Boolean =
+  annotations.any { it.classId?.asSingleFqName() == fqName }
+
+context(KaSession)
+private fun KtLambdaArgument.isInlinedInside(function: KaCallInfo): Boolean =
+  function.isInline() &&
+    function.parameterSymbolOf(this@isInlinedInside)
+      ?.let { !it.isNoinline && !it.isCrossinline }
+      ?: true
+
+private fun KaCallInfo.parameterSymbolOf(argument: KtLambdaArgument): KaValueParameterSymbol? =
+  argument.getArgumentExpression()?.let { expression ->
+    successfulFunctionCallOrNull()
+      ?.argumentMapping
+      ?.get(expression)
+      ?.symbol
+  }
+
+private fun KaCallInfo.isInline(): Boolean =
+  successfulFunctionCallOrNull()
+    ?.symbol
+    ?.let { it as? KaNamedFunctionSymbol }
+    ?.isInline
+    ?: false
+
+/** Not cached, but doesn't matter for this infrequently triggered extension */
+context(KaSession)
+private fun Iterable<KtAnnotationEntry>.containsAnnotationNamed(fqName: FqName): Boolean =
+  any { it.resolveToCall()?.singleConstructorCallOrNull()?.symbol?.containingClassId?.asSingleFqName() == fqName }

plugins/compose/intellij.compose.ide.plugin.shared/src/com/intellij/compose/ide/plugin/shared/PsiUtils.kt

@@ -1,6 +1,7 @@
 // Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.compose.ide.plugin.shared
 
+import com.intellij.openapi.module.Module
 import com.intellij.openapi.module.ModuleUtilCore
 import com.intellij.openapi.roots.ProjectRootModificationTracker
 import com.intellij.openapi.roots.impl.ProjectFileIndexFacade
@@ -57,7 +58,13 @@ fun isComposeEnabledForElementModule(element: PsiElement): Boolean {
 }
 
 internal fun PsiElement.isComposableFunction(): Boolean =
-  (this as? KtNamedFunction)?.getAnnotationWithCaching(COMPOSABLE_FUNCTION_KEY) { it.isComposableAnnotation() } != null
+  this is KtNamedFunction && this.hasComposableAnnotation()
+
+internal fun KtAnnotated.hasComposableAnnotation(): Boolean =
+  this.getAnnotationWithCaching(COMPOSABLE_FUNCTION_KEY) { it.isComposableAnnotation() } != null
+
+internal val PsiElement.module: Module?
+  get() = ModuleUtilCore.findModuleForPsiElement(this)
 
 private val COMPOSABLE_FUNCTION_KEY: Key<CachedValue<KtAnnotationEntry?>> =
   Key.create("com.intellij.compose.ide.plugin.shared.isComposableFunction")