diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/src/com/intellij/compose/ide/plugin/shared/IsComposableControlFlowUtil.kt b/plugins/compose/intellij.compose.ide.plugin.shared/src/com/intellij/compose/ide/plugin/shared/IsComposableControlFlowUtil.kt
index 7b7f46ef2109..503839caf55d 100644
--- a/plugins/compose/intellij.compose.ide.plugin.shared/src/com/intellij/compose/ide/plugin/shared/IsComposableControlFlowUtil.kt
+++ b/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)
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/src/com/intellij/compose/ide/plugin/shared/PsiUtils.kt b/plugins/compose/intellij.compose.ide.plugin.shared/src/com/intellij/compose/ide/plugin/shared/PsiUtils.kt
index 77acb2296157..ac6300d3a6e1 100644
--- a/plugins/compose/intellij.compose.ide.plugin.shared/src/com/intellij/compose/ide/plugin/shared/PsiUtils.kt
+++ b/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
 
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/src/com/intellij/compose/ide/plugin/shared/folding/ComposeFoldingBuilder.kt b/plugins/compose/intellij.compose.ide.plugin.shared/src/com/intellij/compose/ide/plugin/shared/folding/ComposeFoldingBuilder.kt
index 3ff243544a1e..3dad8bfeb8bc 100644
--- a/plugins/compose/intellij.compose.ide.plugin.shared/src/com/intellij/compose/ide/plugin/shared/folding/ComposeFoldingBuilder.kt
+++ b/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()
   }
 }
\ No newline at end of file
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/BasicComposableArguments.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/BasicComposableArguments.kt
new file mode 100644
index 000000000000..cd4fcb7ad185
--- /dev/null
+++ b/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>
\ No newline at end of file
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/BasicComposableReturnType.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/BasicComposableReturnType.kt
new file mode 100644
index 000000000000..1167ed43a2f3
--- /dev/null
+++ b/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>
\ No newline at end of file
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/BasicInlineFunctions.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/BasicInlineFunctions.kt
new file mode 100644
index 000000000000..cba49c4895cd
--- /dev/null
+++ b/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>
\ No newline at end of file
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/BasicSuspendFunction.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/BasicSuspendFunction.kt
new file mode 100644
index 000000000000..f510c8b88176
--- /dev/null
+++ b/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>
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/InlineCrossinlineLambda.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/InlineCrossinlineLambda.kt
new file mode 100644
index 000000000000..478e8cffed71
--- /dev/null
+++ b/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>
\ No newline at end of file
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/InlineNoinlineLambda.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/InlineNoinlineLambda.kt
new file mode 100644
index 000000000000..087dee38220b
--- /dev/null
+++ b/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>
\ No newline at end of file
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedAnnotatedLambda.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedAnnotatedLambda.kt
index eed496a4b2dd..0f02e35d38b8 100644
--- a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedAnnotatedLambda.kt
+++ b/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
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedComposableReturnType.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedComposableReturnType.kt
index e37362f1e4ca..98437e5fdd40 100644
--- a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedComposableReturnType.kt
+++ b/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>
\ No newline at end of file
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedExplicitTypeLambda.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedExplicitTypeLambda.kt
index cb845714b38c..74537ee572f8 100644
--- a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedExplicitTypeLambda.kt
+++ b/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
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedInheritedComposable.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedInheritedComposable.kt
index 4792bfd53a6b..a633d34f5539 100644
--- a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedInheritedComposable.kt
+++ b/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>
 
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedNoFoldComposableReturnType.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedNoFoldComposableReturnType.kt
new file mode 100644
index 000000000000..b4267798267d
--- /dev/null
+++ b/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>
\ No newline at end of file
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedNoFoldInlineFunction.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedNoFoldInlineFunction.kt
new file mode 100644
index 000000000000..5272f245cd46
--- /dev/null
+++ b/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>
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedNoFoldLambda.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedNoFoldLambda.kt
index c70093c601ed..7bf8fd384437 100644
--- a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NestedNoFoldLambda.kt
+++ b/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>
\ No newline at end of file
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NoFoldComposableReturnType.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NoFoldComposableReturnType.kt
new file mode 100644
index 000000000000..ba8e44cac053
--- /dev/null
+++ b/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>
\ No newline at end of file
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NoFoldScopeStandard.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NoFoldNonComposableFunction.kt
similarity index 100%
rename from plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NoFoldScopeStandard.kt
rename to plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NoFoldNonComposableFunction.kt
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NoFoldLambdas.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NoFoldPropertyLambda.kt
similarity index 67%
rename from plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NoFoldLambdas.kt
rename to plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NoFoldPropertyLambda.kt
index 88135036332e..fef239a5bf62 100644
--- a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/NoFoldLambdas.kt
+++ b/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>
\ No newline at end of file
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/PropertyLambda.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/PropertyLambda.kt
index d9357296f152..c82b5a7a2ddb 100644
--- a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/PropertyLambda.kt
+++ b/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>
\ No newline at end of file
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/SuspendFunction.kt b/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/SuspendFunction.kt
deleted file mode 100644
index 77df24868b86..000000000000
--- a/plugins/compose/intellij.compose.ide.plugin.shared/testData/folding/SuspendFunction.kt
+++ /dev/null
@@ -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>
\ No newline at end of file
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/tests/com/intellij/compose/ide/plugin/shared/ComposableAnnotationToExtractedFunctionAdderTest.kt b/plugins/compose/intellij.compose.ide.plugin.shared/tests/com/intellij/compose/ide/plugin/shared/ComposableAnnotationToExtractedFunctionAdderTest.kt
index 6ac53b0f6936..59782aaae4ad 100644
--- a/plugins/compose/intellij.compose.ide.plugin.shared/tests/com/intellij/compose/ide/plugin/shared/ComposableAnnotationToExtractedFunctionAdderTest.kt
+++ b/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",
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/tests/com/intellij/compose/ide/plugin/shared/ComposeFoldingBuilderTest.kt b/plugins/compose/intellij.compose.ide.plugin.shared/tests/com/intellij/compose/ide/plugin/shared/ComposeFoldingBuilderTest.kt
index 01105c23446c..3db1ef33c5dd 100644
--- a/plugins/compose/intellij.compose.ide.plugin.shared/tests/com/intellij/compose/ide/plugin/shared/ComposeFoldingBuilderTest.kt
+++ b/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")
-  }
 }
diff --git a/plugins/compose/intellij.compose.ide.plugin.shared/tests/com/intellij/compose/ide/plugin/shared/util/ComposeTestUtils.kt b/plugins/compose/intellij.compose.ide.plugin.shared/tests/com/intellij/compose/ide/plugin/shared/util/ComposeTestUtils.kt
index c210b4b37739..3c30b5603e9c 100644
--- a/plugins/compose/intellij.compose.ide.plugin.shared/tests/com/intellij/compose/ide/plugin/shared/util/ComposeTestUtils.kt
+++ b/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)
 }