[Compose] KMT-546: Refactor composable control flow detection

Extended isInsideComposableControlFlow to handle missing lambda cases. Previously only explicitly annotated lambdas were detected, causing tools to miss composable calls. Added tests for folding and extracting.

GitOrigin-RevId: 9a39439105001b7b3ff942f5cad6546f13cf3361

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

@@ -2,10 +2,12 @@
 package com.intellij.compose.ide.plugin.shared
 
 import com.intellij.psi.PsiElement
+import com.intellij.psi.util.parentOfType
 import org.jetbrains.kotlin.analysis.api.KaContextParameterApi
 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.components.expectedType
 import org.jetbrains.kotlin.analysis.api.components.resolveToCall
 import org.jetbrains.kotlin.analysis.api.resolution.KaCallInfo
 import org.jetbrains.kotlin.analysis.api.resolution.singleConstructorCallOrNull
@@ -20,6 +22,7 @@ import org.jetbrains.kotlin.psi.KtLambdaArgument
 import org.jetbrains.kotlin.psi.KtLambdaExpression
 import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.KtPropertyAccessor
+import org.jetbrains.kotlin.psi.KtValueArgument
 import org.jetbrains.kotlin.psi.psiUtil.getAnnotationEntries
 
 /**
@@ -44,22 +47,22 @@ internal fun PsiElement?.isInsideComposableControlFlow(): Boolean = when (this)
   else -> parent.isInsideComposableControlFlow()
 }
 
+@OptIn(KaContextParameterApi::class)
 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)
+  getAnnotationEntries().containsAnnotationNamed(COMPOSABLE_ANNOTATION_FQ_NAME)
+  || expectedType?.isAnnotatedWith(COMPOSABLE_ANNOTATION_FQ_NAME) == true
+  || (parent as? KtValueArgument)?.ownsComposableControlFlow() == true
 
 @OptIn(KaContextParameterApi::class)
 context(_: KaSession)
-private fun KtLambdaArgument.ownsComposableControlFlow(): Boolean =
-  (parent as? KtCallExpression)
+private fun KtValueArgument.ownsComposableControlFlow(): Boolean =
+  parentOfType<KtCallExpression>()
     ?.resolveToCall()
     .let { functionCall -> functionCall != null && this.ownsComposableControlFlowWhenArgumentOf(functionCall) }
 
 context(_: KaSession)
-private fun KtLambdaArgument.ownsComposableControlFlowWhenArgumentOf(functionCall: KaCallInfo): Boolean =
+private fun KtValueArgument.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,
@@ -67,27 +70,27 @@ private fun KtLambdaArgument.ownsComposableControlFlowWhenArgumentOf(functionCal
     && parent.isInsideComposableControlFlow() // then check whether the parent is in Composable flow
   )
 
-private fun KtLambdaArgument.hasAnnotationIn(function: KaCallInfo, fqName: FqName): Boolean = analyze(this) {
+private fun KtValueArgument.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
+  ?: false
 }
 
 private fun KaAnnotated.isAnnotatedWith(fqName: FqName): Boolean =
   annotations.any { it.classId?.asSingleFqName() == fqName }
 
 context(_: KaSession)
-private fun KtLambdaArgument.isInlinedInside(function: KaCallInfo): Boolean =
+private fun KtValueArgument.isInlinedInside(function: KaCallInfo): Boolean =
   function.isInline() &&
-    function.parameterSymbolOf(this@isInlinedInside)
-      ?.let { !it.isNoinline && !it.isCrossinline }
-      ?: true
+  function.parameterSymbolOf(this@isInlinedInside)
+    ?.let { !it.isNoinline && !it.isCrossinline }
+  ?: true
 
-private fun KaCallInfo.parameterSymbolOf(argument: KtLambdaArgument): KaValueParameterSymbol? =
+private fun KaCallInfo.parameterSymbolOf(argument: KtValueArgument): KaValueParameterSymbol? =
   argument.getArgumentExpression()?.let { expression ->
     successfulFunctionCallOrNull()
-      ?.argumentMapping
+      ?.valueArgumentMapping
       ?.get(expression)
       ?.symbol
   }
@@ -97,7 +100,7 @@ private fun KaCallInfo.isInline(): Boolean =
     ?.symbol
     ?.let { it as? KaNamedFunctionSymbol }
     ?.isInline
-    ?: false
+  ?: false
 
 /** Not cached, but doesn't matter for this infrequently triggered extension */
 @OptIn(KaContextParameterApi::class)

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

@@ -26,7 +26,6 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.util.CachedValue
 import com.intellij.psi.util.CachedValueProvider
 import com.intellij.psi.util.CachedValuesManager
-import com.intellij.psi.util.parentOfType
 import com.intellij.util.concurrency.annotations.RequiresReadLock
 import org.jetbrains.annotations.ApiStatus
 import org.jetbrains.kotlin.analysis.api.KaSession
@@ -34,20 +33,15 @@ import org.jetbrains.kotlin.analysis.api.analyze
 import org.jetbrains.kotlin.analysis.api.resolution.KaCallableMemberCall
 import org.jetbrains.kotlin.analysis.api.resolution.calls
 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.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtAnnotated
-import org.jetbrains.kotlin.psi.KtAnnotatedExpression
 import org.jetbrains.kotlin.psi.KtAnnotationEntry
-import org.jetbrains.kotlin.psi.KtCallExpression
 import org.jetbrains.kotlin.psi.KtCallableDeclaration
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtLambdaExpression
 import org.jetbrains.kotlin.psi.KtNamedFunction
-import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.psi.KtValueArgument
 import org.jetbrains.kotlin.psi.KtValueArgumentList
 
@@ -82,34 +76,6 @@ fun isComposeEnabledForElementModule(element: PsiElement): Boolean {
 internal fun PsiElement.isComposableFunction(): Boolean =
   this is KtNamedFunction && this.hasComposableAnnotation()
 
-internal fun KtProperty.isComposableGetter(): Boolean =
-  getter?.hasComposableAnnotation() == true
-
-internal fun KtCallableDeclaration.isComposableType(): Boolean =
-  typeReference?.hasComposableAnnotation() == true
-
-internal fun KtLambdaExpression.isComposableLambda(): Boolean {
-  val parent = this.parent
-  if (parent is KtAnnotatedExpression && parent.hasComposableAnnotation()) return true
-
-  if (parent !is KtValueArgument) return false
-
-  val callExpression = parent.parentOfType<KtCallExpression>() ?: return false
-  return analyze(callExpression) {
-    val functionCall = callExpression.resolveToCall()?.successfulFunctionCallOrNull() ?: return@analyze false
-    val parameterSymbol = functionCall.valueArgumentMapping[this@isComposableLambda]?.symbol ?: return@analyze false
-    parameterSymbol.returnType.annotations.any { it.classId == COMPOSABLE_ANNOTATION_CLASS_ID }
-  }
-}
-
-internal fun KtLambdaExpression.isStandardLambda(): Boolean {
-  if (parent is KtProperty) {
-    val property = parent as KtProperty
-    if (!property.isComposableType()) return true
-  }
-  return false
-}
-
 internal fun KtAnnotated.hasComposableAnnotation(): Boolean =
   this.getAnnotationWithCaching(COMPOSABLE_FUNCTION_KEY) { it.isComposableAnnotation() } != null
 

plugins/compose/intellij.compose.ide.plugin.shared/src/com/intellij/compose/ide/plugin/shared/folding/ComposeFoldingBuilder.kt

@@ -4,12 +4,8 @@ package com.intellij.compose.ide.plugin.shared.folding
 import com.intellij.compose.ide.plugin.shared.COMPOSE_MODIFIER_FQN
 import com.intellij.compose.ide.plugin.shared.callReturnTypeFqName
 import com.intellij.compose.ide.plugin.shared.isAndroidFile
-import com.intellij.compose.ide.plugin.shared.isComposableFunction
-import com.intellij.compose.ide.plugin.shared.isComposableGetter
-import com.intellij.compose.ide.plugin.shared.isComposableLambda
-import com.intellij.compose.ide.plugin.shared.isComposableType
 import com.intellij.compose.ide.plugin.shared.isComposeEnabledForElementModule
-import com.intellij.compose.ide.plugin.shared.isStandardLambda
+import com.intellij.compose.ide.plugin.shared.isInsideComposableControlFlow
 import com.intellij.lang.ASTNode
 import com.intellij.lang.folding.CustomFoldingBuilder
 import com.intellij.lang.folding.FoldingDescriptor
@@ -19,12 +15,8 @@ import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
 import org.jetbrains.annotations.ApiStatus
 import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
-import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi.KtLambdaExpression
-import org.jetbrains.kotlin.psi.KtNamedFunction
-import org.jetbrains.kotlin.psi.KtProperty
-import org.jetbrains.kotlin.psi.KtTreeVisitor
+import org.jetbrains.kotlin.psi.KtTreeVisitorVoid
 
 /**
  * Adds a folding region for a Modifier chain longer or equal than two.
@@ -47,14 +39,12 @@ abstract class ComposeFoldingBuilder : CustomFoldingBuilder() {
     // Do not run on modules that do not have Compose enabled
     if (!isComposeEnabledForElementModule(root)) return
 
-    root.accept(ComposeFoldingVisitor(descriptors), false)
+    root.accept(ComposeFoldingVisitor(descriptors))
   }
 
-  private fun KtElement.isModifierChainLongerThanTwo(): Boolean {
-    if (this !is KtDotQualifiedExpression) return false
-    if (this.receiverExpression !is KtDotQualifiedExpression) return false
-
-    return this.callReturnTypeFqName() == COMPOSE_MODIFIER_FQN
+  private fun KtDotQualifiedExpression.isModifierChainLongerThanOne(): Boolean {
+    if (receiverExpression !is KtDotQualifiedExpression) return false
+    return callReturnTypeFqName() == COMPOSE_MODIFIER_FQN
   }
 
   /** For Modifier.adjust().adjust() -> Modifier.(...) */
@@ -66,32 +56,18 @@ abstract class ComposeFoldingBuilder : CustomFoldingBuilder() {
 
   private inner class ComposeFoldingVisitor(
     private val descriptors: MutableList<FoldingDescriptor>
-  ) : KtTreeVisitor<Boolean>() {
+  ) : KtTreeVisitorVoid() {
 
-    override fun visitNamedFunction(function: KtNamedFunction, insideComposable: Boolean): Void?{
-      val isComposable = function.isComposableFunction() || function.isComposableType()
-      return super.visitNamedFunction(function, isComposable)
-    }
-
-    override fun visitProperty(property: KtProperty, insideComposable: Boolean): Void? {
-      val isComposable = insideComposable || property.isComposableGetter() || property.isComposableType()
-      return super.visitProperty(property, isComposable)
-    }
-
-    override fun visitLambdaExpression(lambdaExpression: KtLambdaExpression, insideComposable: Boolean): Void? {
-      val isComposable = when {
-        lambdaExpression.isComposableLambda() -> true
-        lambdaExpression.isStandardLambda() -> false
-        else -> insideComposable
-      }
-      return super.visitLambdaExpression(lambdaExpression, isComposable)
-    }
-
-    override fun visitDotQualifiedExpression(expression: KtDotQualifiedExpression, insideComposable: Boolean): Void? {
-      if (insideComposable && expression.parent !is KtDotQualifiedExpression && expression.isModifierChainLongerThanTwo()) {
+    override fun visitDotQualifiedExpression(expression: KtDotQualifiedExpression) {
+      if (expression.isFoldableModifierChain()) {
         descriptors.add(FoldingDescriptor(expression.node, expression.node.textRange))
       }
-      return super.visitDotQualifiedExpression(expression, insideComposable)
+      super.visitDotQualifiedExpression(expression)
     }
+
+    fun KtDotQualifiedExpression.isFoldableModifierChain(): Boolean =
+      parent !is KtDotQualifiedExpression &&
+      isModifierChainLongerThanOne() &&
+      isInsideComposableControlFlow()
   }
 }

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/BasicComposableArguments.kt

@@ -0,0 +1,51 @@
+package com.example
+
+import <fold text='...'>androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier</fold>
+
+suspend fun awaitApplication(content: @Composable () -> Unit) {}
+
+suspend fun multipleComposable(
+  header: @Composable () -> Unit,
+  footer: @Composable () -> Unit
+) {}
+
+suspend fun mixedNamed(
+  composable: @Composable () -> Unit,
+  suspendBlock: suspend () -> Unit
+) {}
+
+@Composable
+suspend fun composableFunction() <fold text='{...}'>{
+  awaitApplication(content = <fold text='{...}'>{
+    val folded = <fold text='Modifier.(...)'>Modifier
+      .adjust()
+      .adjust()</fold>
+  }</fold>)
+
+  multipleComposable<fold text='(...)'>(
+    header = <fold text='{...}'>{
+      val folded1 = <fold text='Modifier.(...)'>Modifier
+        .adjust()
+        .adjust()</fold>
+    }</fold>,
+    footer = <fold text='{...}'>{
+      val folded2 = <fold text='Modifier.(...)'>Modifier
+        .adjust()
+        .adjust()</fold>
+    }</fold>
+  )</fold>
+
+  mixedNamed<fold text='(...)'>(
+    composable = <fold text='{...}'>{
+      val folded = <fold text='Modifier.(...)'>Modifier
+        .adjust()
+        .adjust()</fold>
+    }</fold>,
+    suspendBlock = <fold text='{...}'>{
+      val notFolded = Modifier
+        .adjust()
+        .adjust()
+    }</fold>
+  )</fold>
+}</fold>

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/BasicComposableReturnType.kt

@@ -0,0 +1,38 @@
+package com.example
+
+import <fold text='...'>androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier</fold>
+
+fun returnsComposable(): @Composable () -> Unit <fold text='{...}'>{
+  return <fold text='{...}'>{
+    <fold text='Modifier.(...)'>Modifier
+      .adjust()
+      .adjust()</fold>
+  }</fold>
+}</fold>
+
+fun returnsComposableWithParam(): @Composable (Int) -> Unit <fold text='{...}'>{
+  return <fold text='{...}'>{
+    <fold text='Modifier.(...)'>Modifier
+      .adjust()
+      .adjust()
+      .adjust()</fold>
+  }</fold>
+}</fold>
+
+fun directReturn(): @Composable () -> Unit <fold text='{...}'>{
+  return <fold text='{...}'>{
+    <fold text='Modifier.(...)'>Modifier
+      .adjust()
+      .adjust()</fold>
+  }</fold>
+}</fold>
+
+fun foo(): Sequence<@Composable () -> Unit> = <fold text='{...}'>(0..10)
+  .asSequence()
+  .map <fold text='{...}'>{ return@map <fold text='{...}'>{
+    <fold text='Modifier.(...)'>Modifier
+      .adjust()
+      .adjust()</fold>
+  }</fold>
+}</fold></fold>

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/BasicInlineFunctions.kt

@@ -0,0 +1,26 @@
+package com.example
+
+import <fold text='...'>androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier</fold>
+
+inline fun <T> myRun(block: () -> T): T = block()
+
+@Composable
+fun InlineRunFolds() <fold text='{...}'>{
+  myRun <fold text='{...}'>{
+    <fold text='Modifier.(...)'>Modifier
+      .adjust()
+      .adjust()</fold>
+    }</fold>
+}</fold>
+
+@Composable
+fun NestedInlineFolds() <fold text='{...}'>{
+  myRun <fold text='{...}'>{
+    myRun <fold text='{...}'>{
+      <fold text='Modifier.(...)'>Modifier
+        .adjust()
+        .adjust()</fold>
+      }</fold>
+    }</fold>
+}</fold>

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/BasicSuspendFunction.kt

@@ -0,0 +1,33 @@
+package com.example
+
+import <fold text='...'>androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier</fold>
+
+suspend fun awaitApplication(content: @Composable () -> Unit) {}
+suspend fun regularSuspendLambda(block: suspend () -> Unit) {}
+
+suspend fun main() = <fold text='{...}'>awaitApplication <fold text='{...}'>{
+  val m = <fold text='Modifier.(...)'>Modifier
+    .adjust()
+    .adjust()</fold>
+}</fold></fold>
+
+suspend fun noFoldOutsideComposable() <fold text='{...}'>{
+  val notFolded = Modifier
+    .adjust()
+    .adjust()
+
+  awaitApplication <fold text='{...}'>{
+    val folded = <fold text='Modifier.(...)'>Modifier
+      .adjust()
+      .adjust()</fold>
+  }</fold>
+}</fold>
+
+suspend fun noFoldInSuspendLambda() <fold text='{...}'>{
+  regularSuspendLambda <fold text='{...}'>{
+    val notFolded = Modifier
+      .adjust()
+      .adjust()
+  }</fold>
+}</fold>

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/InlineCrossinlineLambda.kt

@@ -0,0 +1,22 @@
+package com.example
+
+import <fold text='...'>androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier</fold>
+
+inline fun fold(crossinline a: @Composable () -> Unit) {}
+inline fun noFold(crossinline block: () -> Unit) { block() }
+
+@Composable
+fun ComposableFunction() <fold text='{...}'>{
+  noFold <fold text='{...}'>{
+    Modifier
+      .adjust()
+      .adjust()
+  }</fold>
+
+  fold <fold text='{...}'>{
+     <fold text='Modifier.(...)'>Modifier
+      .adjust()
+      .adjust()</fold>
+  }</fold>
+}</fold>

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/InlineNoinlineLambda.kt

@@ -0,0 +1,24 @@
+package com.example
+
+import <fold text='...'>androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier</fold>
+
+inline fun fold(noinline a: @Composable () -> Unit) {}
+inline fun noFold(noinline block: () -> Unit) { block() }
+
+@Composable
+fun ComposableFunction() <fold text='{...}'>{
+
+  fold <fold text='{...}'>{
+    <fold text='Modifier.(...)'>Modifier
+      .adjust()
+      .adjust()</fold>
+  }</fold>
+
+  noFold <fold text='{...}'>{
+    Modifier
+      .adjust()
+      .adjust()
+  }</fold>
+
+}</fold>

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedAnnotatedLambda.kt

@@ -20,12 +20,6 @@ val annotatedLambda = @Composable <fold text='{...}'>{
       .adjust()</fold>
   }</fold>
 
-  val myModifier = remember <fold text='{...}'>{
-    <fold text='Modifier.(...)'>Modifier
-      .adjust()
-      .adjust()</fold>
-  }</fold>
-
   @Composable
   fun ComposableFunction() <fold text='{...}'>{
     <fold text='Modifier.(...)'>Modifier

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedComposableReturnType.kt

@@ -3,27 +3,68 @@ package com.example
 import <fold text='...'>androidx.compose.runtime.Composable
 import androidx.compose.ui.Modifier</fold>
 
-fun returnsComposable(): @Composable () -> Unit <fold text='{...}'>{
-  return <fold text='{...}'>{
-    <fold text='Modifier.(...)'>Modifier
-      .adjust()
-      .adjust()</fold>
+fun ifElseReturn(condition: Boolean): @Composable () -> Unit <fold text='{...}'>{
+  return if (condition) <fold text='{...}'>{
+    <fold text='{...}'>{
+      <fold text='Modifier.(...)'>Modifier
+        .adjust()
+        .adjust()</fold>
+    }</fold>
+  }</fold> else <fold text='{...}'>{
+    <fold text='{...}'>{
+      <fold text='Modifier.(...)'>Modifier
+        .adjust()
+        .adjust()</fold>
+    }</fold>
   }</fold>
 }</fold>
 
-fun returnsComposableWithParam(): @Composable (Int) -> Unit <fold text='{...}'>{
-  return <fold text='{...}'>{
-    <fold text='Modifier.(...)'>Modifier
-      .adjust()
-      .adjust()
-      .adjust()</fold>
-  }</fold>
+fun whenReturn(x: Int): @Composable () -> Unit <fold text='{...}'>{
+  return when (x) <fold text='{...}'>{
+    1 -> <fold text='{...}'>{
+      <fold text='{...}'>{
+        <fold text='Modifier.(...)'>Modifier
+          .adjust()
+          .adjust()</fold>
+       }</fold>
+     }</fold>
+    2 -> <fold text='{...}'>{
+      <fold text='{...}'>{
+        <fold text='Modifier.(...)'>Modifier
+          .adjust()
+          .adjust()</fold>
+      }</fold>
+    }</fold>
+    else -> <fold text='{...}'>{
+      <fold text='{...}'>{
+        <fold text='Modifier.(...)'>Modifier
+          .adjust()
+          .adjust()</fold>
+        }</fold>
+      }</fold>
+    }</fold>
 }</fold>
 
-fun notReturningComposable(): () -> Unit <fold text='{...}'>{
-  return <fold text='{...}'>{
-    Modifier
-      .adjust()
-      .adjust()
+fun nestedIfReturn(a: Boolean, b: Boolean): @Composable () -> Unit <fold text='{...}'>{
+  return if (a) <fold text='{...}'>{
+    if (b) <fold text='{...}'>{
+      <fold text='{...}'>{
+        <fold text='Modifier.(...)'>Modifier
+          .adjust()
+          .adjust()</fold>
+      }</fold>
+    }</fold> else <fold text='{...}'>{
+      <fold text='{...}'>{
+        <fold text='Modifier.(...)'>Modifier
+          .adjust()
+          .adjust()</fold>
+      }</fold>
+    }</fold>
+  }</fold> else <fold text='{...}'>{
+    <fold text='{...}'>{
+      <fold text='Modifier.(...)'>Modifier
+        .adjust()
+        .adjust()</fold>
+    }</fold>
   }</fold>
 }</fold>

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedExplicitTypeLambda.kt

@@ -20,13 +20,6 @@ val explicitUnit: @Composable () -> Unit = <fold text='{...}'>{
       .adjust()</fold>
   }</fold>
 
-  val blockGetter: Modifier
-  @Composable get() <fold text='{...}'>{
-    return <fold text='Modifier.(...)'>Modifier
-      .adjust()
-      .adjust()</fold>
-  }</fold>
-
   @Composable
   fun ComposableFunction() <fold text='{...}'>{
     <fold text='Modifier.(...)'>Modifier

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedInheritedComposable.kt

@@ -19,12 +19,6 @@ fun MyScreen() <fold text='{...}'>{
         .adjust()
         .adjust()
     }</fold>
-
-    val myModifier = remember <fold text='{...}'>{
-      <fold text='Modifier.(...)'>Modifier
-        .adjust()
-        .adjust()</fold>
-     }</fold>
   }</fold>
 }</fold>
 
@@ -34,13 +28,13 @@ fun Row(content: @Composable () -> Unit) {}
 @Composable
 fun NestedContainerCalls() <fold text='{...}'>{
   Container <fold text='{...}'>{
-  Row <fold text='{...}'>{
-    Container <fold text='{...}'>{
-      <fold text='Modifier.(...)'>Modifier
-        .adjust()
-        .adjust()</fold>
+    Row <fold text='{...}'>{
+      Container <fold text='{...}'>{
+        <fold text='Modifier.(...)'>Modifier
+          .adjust()
+          .adjust()</fold>
+        }</fold>
       }</fold>
-    }</fold>
   }</fold>
 }</fold>
 

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedNoFoldComposableReturnType.kt

@@ -0,0 +1,68 @@
+package com.example
+
+import <fold text='...'>androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier</fold>
+
+fun nestedReturnsComposable(condition: Boolean): @Composable () -> Unit <fold text='{...}'>{
+  return if (condition) <fold text='{...}'>{
+    <fold text='{...}'>{
+      val m = <fold text='Modifier.(...)'>Modifier
+        .adjust()
+        .adjust()</fold>
+    }</fold>
+  }</fold> else <fold text='{...}'>{
+    <fold text='{...}'>{
+      val m = <fold text='Modifier.(...)'>Modifier
+        .adjust()
+        .adjust()</fold>
+    }</fold>
+  }</fold>
+}</fold>
+
+fun notReturnedIfElse(): @Composable () -> Unit <fold text='{...}'>{
+  val notReturned = if (true) <fold text='{...}'>{
+    <fold text='{...}'>{
+      Modifier
+        .adjust()
+        .adjust()
+    }</fold>
+  }</fold> else <fold text='{...}'>{
+    <fold text='{...}'>{
+      Modifier
+        .adjust()
+        .adjust()
+    }</fold>
+  }</fold>
+  return {}
+}</fold>
+
+fun notReturnedWhen(): @Composable () -> Unit <fold text='{...}'>{
+  val notReturned = when (1) <fold text='{...}'>{
+    1 -> <fold text='{...}'>{
+      <fold text='{...}'>{
+        Modifier
+          .adjust()
+          .adjust()
+      }</fold>
+    }</fold>
+    else -> <fold text='{...}'>{
+      <fold text='{...}'>{
+        Modifier
+          .adjust()
+          .adjust()
+      }</fold>
+    }</fold>
+  }</fold>
+  return {}
+}</fold>
+
+fun lambdaInsideCallInsideIf(): @Composable () -> Unit <fold text='{...}'>{
+  if (true) <fold text='{...}'>{
+    listOf(1).map <fold text='{...}'>{
+      Modifier
+        .adjust()
+        .adjust()
+    }</fold>
+  }</fold>
+  return {}
+}</fold>

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedNoFoldInlineFunction.kt

@@ -0,0 +1,44 @@
+package com.example
+
+import <fold text='...'>androidx.compose.runtime.Composable
+import androidx.compose.runtime.DisallowComposableCalls
+import androidx.compose.ui.Modifier</fold>
+
+inline fun <T> myRun(block: () -> T): T = block()
+fun <T> myRunNonInline(block: () -> T): T = block()
+
+inline fun f(block: @DisallowComposableCalls () -> Unit) = Unit
+
+@Composable
+fun composableFunction() <fold text='{...}'>{
+  // Non inline
+  myRunNonInline <fold text='{...}'>{
+    val notFolded = Modifier
+      .adjust()
+      .adjust()
+  }</fold>
+
+  // Inline inside non inline
+  myRunNonInline <fold text='{...}'>{
+    myRun <fold text='{...}'>{
+       val notFolded = Modifier
+         .adjust()
+         .adjust()
+     }</fold>
+  }</fold>
+
+  // DisallowComposableCalls block
+  f <fold text='{...}'>{
+    val notFolded = Modifier
+      .adjust()
+      .adjust()
+  }</fold>
+}</fold>
+
+fun NonComposableWithInline() <fold text='{...}'>{
+  myRun <fold text='{...}'>{
+    val notFolded = Modifier
+      .adjust()
+      .adjust()
+  }</fold>
+}</fold>

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedNoFoldLambda.kt

@@ -3,17 +3,35 @@ package com.example
 import <fold text='...'>androidx.compose.runtime.Composable
 import androidx.compose.ui.Modifier</fold>
 
+@Composable
+fun SideEffect(effect: () -> Unit) {}
+
+@Composable
+fun LaunchedEffect(key: Any?, block: suspend () -> Unit) {}
+
 @Composable
 fun ComposableFunction() <fold text='{...}'>{
   val annotatedLambda = <fold text='{...}'>{
-     Modifier
+    val notFolded = Modifier
        .adjust()
        .adjust()
    }</fold>
 
   val explicitLambda: () -> Unit = <fold text='{...}'>{
-    val m = Modifier
-      .padding(8.dp)
-      .fillMaxWidth()
+    val notFolded = Modifier
+      .adjust()
+      .adjust()
+  }</fold>
+
+  SideEffect <fold text='{...}'>{
+    val notFolded = Modifier
+      .adjust()
+      .adjust()
+  }</fold>
+
+  LaunchedEffect(Unit) <fold text='{...}'>{
+    val notFolded = Modifier
+      .adjust()
+      .adjust()
   }</fold>
 }</fold>

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NoFoldComposableReturnType.kt

@@ -0,0 +1,28 @@
+package com.example
+
+import <fold text='...'>androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier</fold>
+
+fun notReturningComposable(): () -> Unit <fold text='{...}'>{
+  return <fold text='{...}'>{
+    Modifier
+      .adjust()
+      .adjust()
+  }</fold>
+}</fold>
+
+fun nonComposableReturn(condition: Boolean): () -> Unit <fold text='{...}'>{
+  return if (condition) <fold text='{...}'>{
+    <fold text='{...}'>{
+      Modifier
+        .adjust()
+        .adjust()
+    }</fold>
+  }</fold> else <fold text='{...}'>{
+    <fold text='{...}'>{
+      Modifier
+        .adjust()
+        .adjust()
+    }</fold>
+  }</fold>
+}</fold>

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NoFoldPropertyLambda.kt

@@ -28,4 +28,18 @@ val standardLambda = <fold text='{...}'>{
     .adjust()
     .adjust()
     .adjust()
+}</fold>
+
+val explicitLambdaIfElse: () -> Unit = if (true) <fold text='{...}'>{
+  <fold text='{...}'>{
+    Modifier
+      .adjust()
+      .adjust()
+  }</fold>
+}</fold> else <fold text='{...}'>{
+  <fold text='{...}'>{
+    Modifier
+      .adjust()
+      .adjust()
+  }</fold>
 }</fold>

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/PropertyLambda.kt

@@ -25,4 +25,35 @@ val modifierType: @Composable Modifier.() -> Modifier = <fold text='{...}'>{
   <fold text='this.(...)'>this
     .adjust()
     .adjust()</fold>
+}</fold>
+
+val ifElseLambda: @Composable () -> Unit = if (true) <fold text='{...}'>{
+  <fold text='{...}'>{
+    <fold text='Modifier.(...)'>Modifier
+      .adjust()
+      .adjust()</fold>
+  }</fold>
+}</fold> else <fold text='{...}'>{
+  <fold text='{...}'>{
+    <fold text='Modifier.(...)'>Modifier
+      .adjust()
+      .adjust()</fold>
+  }</fold>
+}</fold>
+
+val whenLambda: @Composable () -> Unit = when (1) <fold text='{...}'>{
+  1 -> <fold text='{...}'>{
+    <fold text='{...}'>{
+      <fold text='Modifier.(...)'>Modifier
+        .adjust()
+        .adjust()</fold>
+    }</fold>
+  }</fold>
+  else -> <fold text='{...}'>{
+    <fold text='{...}'>{
+      <fold text='Modifier.(...)'>Modifier
+        .adjust()
+        .adjust()</fold>
+    }</fold>
+  }</fold>
 }</fold>

plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/SuspendFunction.kt

@@ -1,76 +0,0 @@
-package com.example
-
-import <fold text='...'>androidx.compose.runtime.Composable
-import androidx.compose.ui.Modifier</fold>
-
-suspend fun awaitApplication(content: @Composable () -> Unit) {}
-
-suspend fun main() = <fold text='{...}'>awaitApplication <fold text='{...}'>{
-  val m = <fold text='Modifier.(...)'>Modifier
-    .adjust()
-    .adjust()</fold>
-}</fold></fold>
-
-suspend fun anotherSuspendMain() = <fold text='{...}'>awaitApplication <fold text='{...}'>{
-  @Composable
-  fun NestedComposable() <fold text='{...}'>{
-    val inner = <fold text='Modifier.(...)'>Modifier
-      .adjust()
-      .adjust()
-      .adjust()</fold>
-  }</fold>
-
-  val outer = <fold text='Modifier.(...)'>Modifier
-    .adjust()
-    .adjust()</fold>
-}</fold></fold>
-
-suspend fun withNestedAwait() <fold text='{...}'>{
-  awaitApplication <fold text='{...}'>{
-    val m = <fold text='Modifier.(...)'>Modifier
-      .adjust()
-      .adjust()</fold>
-  }</fold>
-}</fold>
-
-suspend fun noFoldOutsideComposable() <fold text='{...}'>{
-  val notFolded = Modifier
-    .adjust()
-    .adjust()
-
-  awaitApplication <fold text='{...}'>{
-    val folded = <fold text='Modifier.(...)'>Modifier
-      .adjust()
-      .adjust()</fold>
-  }</fold>
-}</fold>
-
-suspend fun regularSuspendLambda(block: suspend () -> Unit) {}
-
-suspend fun noFoldInSuspendLambda() <fold text='{...}'>{
-  regularSuspendLambda <fold text='{...}'>{
-    val notFolded = Modifier
-      .adjust()
-      .adjust()
-  }</fold>
-}</fold>
-
-suspend fun mixedParams(
-  composableContent: @Composable () -> Unit,
-  suspendBlock: suspend () -> Unit
-) {}
-
-suspend fun mixedParamsTest() <fold text='{...}'>{
-  mixedParams<fold text='(...)'>(
-    composableContent = <fold text='{...}'>{
-      val folded = <fold text='Modifier.(...)'>Modifier
-        .adjust()
-        .adjust()</fold>
-    }</fold>,
-    suspendBlock = <fold text='{...}'>{
-      val notFolded = Modifier
-        .adjust()
-        .adjust()
-    }</fold>
-  )</fold>
-}</fold>

plugins/compose/intellij.compose.ide.plugin.shared/tests/com/intellij/compose/ide/plugin/shared/ComposableAnnotationToExtractedFunctionAdderTest.kt

@@ -221,6 +221,31 @@ abstract class ComposableAnnotationToExtractedFunctionAdderTest : KotlinLightCod
       """.trimIndent())
   }
 
+  fun `test inside property with explicit Composable type`() {
+    assertThatGivenComposeProjectWithState("""
+      import androidx.compose.material.Text
+      import androidx.compose.runtime.Composable
+  
+      val myComposable: @Composable () -> Unit = {
+          <selection>Text("Hello")</selection>
+      }
+    """.trimIndent())
+      .invokingExtractFunction()
+      .resultsIn("""
+        import androidx.compose.material.Text
+        import androidx.compose.runtime.Composable
+  
+        val myComposable: @Composable () -> Unit = {
+            newFunction()
+        }
+  
+        @Composable
+        private fun newFunction() {
+            Text("Hello")
+        }
+      """.trimIndent())
+  }
+
   fun `test inside non-inline Composable call with Composable lambda`() {
     assertThatGivenComposeProjectWithState("""
       import androidx.compose.foundation.layout.Column
@@ -908,6 +933,61 @@ abstract class ComposableAnnotationToExtractedFunctionAdderTest : KotlinLightCod
       """.trimIndent())
   }
 
+  fun `test inside returned lambda from function with Composable return type`() {
+    assertThatGivenComposeProjectWithState("""
+      import androidx.compose.material.Text
+      import androidx.compose.runtime.Composable
+  
+      fun returnsComposable(): @Composable () -> Unit {
+          return {
+              <selection>Text("Hello")</selection>
+          }
+      }
+    """.trimIndent())
+      .invokingExtractFunction()
+      .resultsIn("""
+        import androidx.compose.material.Text
+        import androidx.compose.runtime.Composable
+  
+        fun returnsComposable(): @Composable () -> Unit {
+            return {
+                newFunction()
+            }
+        }
+  
+        @Composable
+        private fun newFunction() {
+            Text("Hello")
+        }
+      """.trimIndent())
+  }
+
+  fun `test inside returned lambda from function with non-Composable return type`() {
+    assertThatGivenComposeProjectWithState("""
+        import androidx.compose.runtime.Composable
+    
+        fun returnsLambda(): () -> Unit {
+            return {
+                <selection>println("Hello")</selection>
+            }
+        }
+      """.trimIndent())
+      .invokingExtractFunction()
+      .resultsIn("""
+        import androidx.compose.runtime.Composable
+  
+        fun returnsLambda(): () -> Unit {
+            return {
+                newFunction()
+            }
+        }
+  
+        private fun newFunction() {
+            println("Hello")
+        }
+      """.trimIndent())
+  }
+
   fun `test Composable is not imported`() {
     assertThatGivenComposeProjectWithState("""
       @Composable
@@ -928,6 +1008,132 @@ abstract class ComposableAnnotationToExtractedFunctionAdderTest : KotlinLightCod
       """.trimIndent())
   }
 
+  fun `test inside named composable argument`() {
+    assertThatGivenComposeProjectWithState("""
+      import androidx.compose.material.Text
+      import androidx.compose.runtime.Composable
+  
+      fun withComposableParam(content: @Composable () -> Unit) {}
+  
+      fun f() {
+          withComposableParam(content = {
+              <selection>Text("Hello")</selection>
+          })
+      }
+    """.trimIndent())
+      .invokingExtractFunction()
+      .resultsIn("""
+        import androidx.compose.material.Text
+        import androidx.compose.runtime.Composable
+  
+        fun withComposableParam(content: @Composable () -> Unit) {}
+  
+        fun f() {
+            withComposableParam(content = {
+                newFunction()
+            })
+        }
+  
+        @Composable
+        private fun newFunction() {
+            Text("Hello")
+        }
+      """.trimIndent())
+  }
+
+  fun `test inside mixed named composable argument`() {
+    assertThatGivenComposeProjectWithState("""
+      import androidx.compose.runtime.Composable
+  
+      suspend fun mixedNamed(
+          composable: @Composable () -> Unit,
+          suspendBlock: suspend () -> Unit
+      ) {}
+  
+      fun f() {
+          mixedNamed(
+              composable = {
+                  <selection>Text("Hello")</selection>
+              },
+              suspendBlock = {
+                  println("World")
+              }
+          )
+      }
+    """.trimIndent())
+      .invokingExtractFunction()
+      .resultsIn("""
+        import androidx.compose.runtime.Composable
+  
+        suspend fun mixedNamed(
+            composable: @Composable () -> Unit,
+            suspendBlock: suspend () -> Unit
+        ) {}
+    
+        fun f() {
+            mixedNamed(
+                composable = {
+                    newFunction()
+                },
+                suspendBlock = {
+                    println("World")
+                }
+            )
+        }
+        
+        @Composable
+        private fun newFunction() {
+            Text("Hello")
+        }
+      """.trimIndent())
+  }
+
+  fun `test inside mixed named non-composable argument`() {
+    assertThatGivenComposeProjectWithState("""
+      import androidx.compose.runtime.Composable
+  
+      suspend fun mixedNamed(
+          composable: @Composable () -> Unit,
+          suspendBlock: suspend () -> Unit
+      ) {}
+  
+      fun f() {
+          mixedNamed(
+              composable = {
+                  Text("Hello")
+              },
+              suspendBlock = {
+                  <selection>println("World")</selection>
+              }
+          )
+      }
+    """.trimIndent())
+      .invokingExtractFunction()
+      .resultsIn("""
+        import androidx.compose.runtime.Composable
+  
+        suspend fun mixedNamed(
+            composable: @Composable () -> Unit,
+            suspendBlock: suspend () -> Unit
+        ) {}
+    
+        fun f() {
+            mixedNamed(
+                composable = {
+                    Text("Hello")
+                },
+                suspendBlock = {
+                    newFunction()
+                }
+            )
+        }
+  
+        private fun newFunction() {
+            println("World")
+        }
+      """.trimIndent())
+  }
+
   fun `test Composable is not on classpath`() {
     myFixture.configureByText(
       "test.kt",

plugins/compose/intellij.compose.ide.plugin.shared/tests/com/intellij/compose/ide/plugin/shared/ComposeFoldingBuilderTest.kt

@@ -43,11 +43,25 @@ abstract class ComposeFoldingBuilderTest : KotlinLightCodeInsightFixtureTestCase
       }
       """
     )
+
   }
+
+  fun `test basic composable arguments`() {
+    myFixture.testFolding("$testPath/BasicComposableArguments.kt")
+  }
+
+  fun `test basic composable return type`() {
+    myFixture.testFolding("$testPath/BasicComposableReturnType.kt")
+  }
+
   fun `test basic conditional expressions`() {
     myFixture.testFolding("$testPath/BasicConditionalExpressions.kt")
   }
 
+  fun `test basic inline function`() {
+    myFixture.testFolding("$testPath/BasicInlineFunctions.kt")
+  }
+
   fun `test basic modifier chain`() {
     myFixture.testFolding("$testPath/BasicModifierChain.kt")
   }
@@ -60,6 +74,18 @@ abstract class ComposeFoldingBuilderTest : KotlinLightCodeInsightFixtureTestCase
     myFixture.testFolding("$testPath/BasicScopeFunctions.kt")
   }
 
+  fun `test basic suspend function`() {
+    myFixture.testFolding("$testPath/BasicSuspendFunction.kt")
+  }
+
+  fun `test inline crossinline lambda`() {
+    myFixture.testFolding("$testPath/InlineCrossinlineLambda.kt")
+  }
+
+  fun `test inline noinline lambda`() {
+    myFixture.testFolding("$testPath/InlineNoinlineLambda.kt")
+  }
+
   fun `test nested annotated lambdas`() {
     myFixture.testFolding("$testPath/NestedAnnotatedLambda.kt")
   }
@@ -88,6 +114,14 @@ abstract class ComposeFoldingBuilderTest : KotlinLightCodeInsightFixtureTestCase
     myFixture.testFolding("$testPath/NestedInheritedComposable.kt")
   }
 
+  fun `test nested no fold composable return type`() {
+    myFixture.testFolding("$testPath/NestedNoFoldComposableReturnType.kt")
+  }
+
+  fun `test nested no fold inline function`() {
+    myFixture.testFolding("$testPath/NestedNoFoldInlineFunction.kt")
+  }
+
   fun `test nested no fold inner function`() {
     myFixture.testFolding("$testPath/NestedNoFoldInnerFunction.kt")
   }
@@ -100,8 +134,12 @@ abstract class ComposeFoldingBuilderTest : KotlinLightCodeInsightFixtureTestCase
     myFixture.testFolding("$testPath/NoFoldBadEnding.kt")
   }
 
+  fun `test no fold composable return type`() {
+    myFixture.testFolding("$testPath/NoFoldComposableReturnType.kt")
+  }
+
   fun `test no fold lambdas`() {
-    myFixture.testFolding("$testPath/NoFoldLambdas.kt")
+    myFixture.testFolding("$testPath/NoFoldPropertyLambda.kt")
   }
 
   fun `test no fold non modifier`() {
@@ -113,7 +151,7 @@ abstract class ComposeFoldingBuilderTest : KotlinLightCodeInsightFixtureTestCase
   }
 
   fun `test no fold scope standard`() {
-    myFixture.testFolding("$testPath/NoFoldScopeStandard.kt")
+    myFixture.testFolding("$testPath/NoFoldNonComposableFunction.kt")
   }
 
   fun `test property getter`() {
@@ -123,8 +161,4 @@ abstract class ComposeFoldingBuilderTest : KotlinLightCodeInsightFixtureTestCase
   fun `test property lambda`() {
     myFixture.testFolding("$testPath/PropertyLambda.kt")
   }
-
-  fun `test suspend function`() {
-    myFixture.testFolding("$testPath/SuspendFunction.kt")
-  }
 }

plugins/compose/intellij.compose.ide.plugin.shared/tests/com/intellij/compose/ide/plugin/shared/util/ComposeTestUtils.kt

@@ -56,6 +56,7 @@ private fun JavaCodeInsightTestFixture.createComposableAnnotationInFile(fileRela
                   package androidx.compose.runtime
   
                   annotation class Composable
+                  annotation class DisallowComposableCalls
               """.trimIndent()
   this.addFileToProject(fileRelativePath, composableAnnotationDefinitionFileText)
 }