Merge "Re-enable startRestartGroup/endRestartGroup for functions" into androidx-master-dev

6 files changed, 285 insertions, 228 deletions

diff --git a/compose/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ComposeCallResolutionInterceptorExtension.kt b/compose/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ComposeCallResolutionInterceptorExtension.kt
index 96811f0b386..fa66b06a294 100644
--- a/compose/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ComposeCallResolutionInterceptorExtension.kt
+++ b/compose/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ComposeCallResolutionInterceptorExtension.kt
@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.ConstructorDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.extensions.internal.CallResolutionInterceptorExtension
 import org.jetbrains.kotlin.incremental.components.LookupLocation
@@ -112,10 +113,10 @@ open class ComposeCallResolutionInterceptorExtension : CallResolutionInterceptor
         // to do any work at all, since it will never be anything we intercept
         if (!needToLookupComposer) return candidates
 
-        val temporaryTraceForKtxCall =
+        val temporaryTraceForComposeableCall =
             TemporaryTraceAndCache.create(
                 resolutionContext,
-                "trace to resolve ktx call", element
+                "trace to resolve composable call", element
             )
 
         val composableAnnotationChecker =
@@ -125,7 +126,7 @@ open class ComposeCallResolutionInterceptorExtension : CallResolutionInterceptor
         // TODO(lmr): there ought to be a better way to do this
         var walker: PsiElement? = call.callElement
         var isComposableContext = false
-//        var composableDescriptor: SimpleFunctionDescriptor? = null
+        var composableDescriptor: SimpleFunctionDescriptor? = null
         while (walker != null) {
             val descriptor = try {
                 resolutionContext.trace[BindingContext.FUNCTION, walker]
@@ -134,12 +135,12 @@ open class ComposeCallResolutionInterceptorExtension : CallResolutionInterceptor
             }
             if (descriptor != null) {
                 val composability = composableAnnotationChecker.analyze(
-                    temporaryTraceForKtxCall.trace,
+                    temporaryTraceForComposeableCall.trace,
                     descriptor
                 )
                 isComposableContext =
                     composability != ComposableAnnotationChecker.Composability.NOT_COMPOSABLE
-//                if (isComposableContext) composableDescriptor = descriptor
+                if (isComposableContext) composableDescriptor = descriptor
 
                 // If the descriptor is for an inlined lambda, infer composability from the
                 // outer scope
@@ -197,6 +198,42 @@ open class ComposeCallResolutionInterceptorExtension : CallResolutionInterceptor
             return nonComposablesNonConstructors + constructors
         }
 
+        val context = ExpressionTypingContext.newContext(
+            resolutionContext.trace,
+            resolutionContext.scope,
+            resolutionContext.dataFlowInfo,
+            resolutionContext.expectedType,
+            resolutionContext.languageVersionSettings,
+            resolutionContext.dataFlowValueFactory
+        )
+
+        val temporaryForComposableCall = context.replaceTraceAndCache(
+            temporaryTraceForComposeableCall
+        )
+
+        // Once we know we have a valid composer record the composer on the function descriptor
+        // for the restart transform
+        val trace = temporaryTraceForComposeableCall.trace
+        if (composableDescriptor != null && trace.get(
+                ComposeWritableSlices.RESTART_COMPOSER_NEEDED,
+                composableDescriptor
+            ) != false) {
+            val recordingContext = TemporaryTraceAndCache.create(
+                context,
+                "trace to resolve composable restart group methods",
+                element
+            )
+            val recordingTrace = recordingContext.trace
+            recordingTrace.record(
+                ComposeWritableSlices.RESTART_COMPOSER_NEEDED,
+                composableDescriptor,
+                false
+            )
+            recordingTrace.record(
+                ComposeWritableSlices.RESTART_COMPOSER, composableDescriptor, composerCall)
+            recordingTrace.commit()
+        }
+
         // If there are no constructors, then all of the candidates are either composables or
         // non-composable functions, and we follow normal resolution rules.
         if (constructors.isEmpty()) {
@@ -227,63 +264,14 @@ open class ComposeCallResolutionInterceptorExtension : CallResolutionInterceptor
             composerMetadata
         )
 
-        val context = ExpressionTypingContext.newContext(
-            resolutionContext.trace,
-            resolutionContext.scope,
-            resolutionContext.dataFlowInfo,
-            resolutionContext.expectedType,
-            resolutionContext.languageVersionSettings,
-            resolutionContext.dataFlowValueFactory
-        )
-
-        val temporaryForKtxCall = context.replaceTraceAndCache(temporaryTraceForKtxCall)
-
         val emitCandidates = emitResolver.resolveCandidates(
             call,
             emittables,
             composerCall,
             name,
-            temporaryForKtxCall
+            temporaryForComposableCall
         )
 
-        // TODO(lmr): deal with the RESTART_CALLS_NEEDED stuff for restartGroups
-        // Once we know we have a valid binding to a composable function call see if the scope need
-        // the startRestartGroup and endRestartGroup information
-//        val trace = temporaryTraceForKtxCall.trace
-//        if (trace.get(
-//                ComposeWritableSlices.RESTART_CALLS_NEEDED,
-//                composableDescriptor!!
-//            ) != false) {
-//            val recordingContext = TemporaryTraceAndCache.create(
-//                context,
-//                "trace to resolve ktx call", element
-//            )
-//            val recordingTrace = recordingContext.trace
-//            recordingTrace.record(
-//                ComposeWritableSlices.RESTART_CALLS_NEEDED,
-//                composableDescriptor,
-//                false
-//            )
-//
-//            val startRestartGroup = ktxCallResolver.resolveStartRestartGroup(
-//                call,
-//                temporaryForKtxCall
-//            )
-//            val endRestartGroup = ktxCallResolver.resolveEndRestartGroup(call, temporaryForKtxCall)
-//            val composerCall = ktxCallResolver.resolveComposerCall()
-//            if (startRestartGroup != null && endRestartGroup != null && composerCall != null) {
-//                recordingTrace.record(
-//                    ComposeWritableSlices.RESTART_CALLS, composableDescriptor,
-//                    ResolvedRestartCalls(
-//                        startRestartGroup = startRestartGroup,
-//                        endRestartGroup = endRestartGroup,
-//                        composer = composerCall
-//                    )
-//                )
-//            }
-//            recordingTrace.commit()
-//        }
-
         return nonComposablesNonConstructors +
                 composables.map {
                     ComposableFunctionDescriptor(it, composerCall, composerMetadata)
diff --git a/compose/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ResolvedRestartCalls.kt b/compose/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ResolvedRestartCalls.kt
deleted file mode 100644
index 9bf1eaacd60..00000000000
--- a/compose/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/ResolvedRestartCalls.kt
+++ /dev/null
@@ -1,25 +0,0 @@
-/*
- * Copyright 2019 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package androidx.compose.plugins.kotlin
-
-import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
-
-class ResolvedRestartCalls(
-    val composer: ResolvedCall<*>,
-    val startRestartGroup: ResolvedCall<*>,
-    val endRestartGroup: ResolvedCall<*>
-)
\ No newline at end of file
diff --git a/compose/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/analysis/ComposeWritableSlices.kt b/compose/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/analysis/ComposeWritableSlices.kt
index b2df5aaf290..dc98a4cf8bb 100644
--- a/compose/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/analysis/ComposeWritableSlices.kt
+++ b/compose/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/analysis/ComposeWritableSlices.kt
@@ -2,15 +2,12 @@ package androidx.compose.plugins.kotlin.analysis
 
 import androidx.compose.plugins.kotlin.ComposableAnnotationChecker
 import androidx.compose.plugins.kotlin.ComposerMetadata
-import androidx.compose.plugins.kotlin.ResolvedRestartCalls
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.ParameterDescriptor
 import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.psi.Call
 import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtReferenceExpression
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.util.slicedMap.BasicWritableSlice
 import org.jetbrains.kotlin.util.slicedMap.RewritePolicy
@@ -28,9 +25,9 @@ object ComposeWritableSlices {
         BasicWritableSlice(RewritePolicy.DO_NOTHING)
     val STABLE_TYPE: WritableSlice<KotlinType, Boolean?> =
         BasicWritableSlice(RewritePolicy.DO_NOTHING)
-    val RESTART_CALLS_NEEDED: WritableSlice<SimpleFunctionDescriptor, Boolean> =
+    val RESTART_COMPOSER_NEEDED: WritableSlice<SimpleFunctionDescriptor, Boolean> =
         BasicWritableSlice(RewritePolicy.DO_NOTHING)
-    val RESTART_CALLS: WritableSlice<SimpleFunctionDescriptor, ResolvedRestartCalls> =
+    val RESTART_COMPOSER: WritableSlice<SimpleFunctionDescriptor, ResolvedCall<*>> =
         BasicWritableSlice(RewritePolicy.DO_NOTHING)
     val COMPOSER_METADATA: WritableSlice<VariableDescriptor, ComposerMetadata> =
         BasicWritableSlice(RewritePolicy.DO_NOTHING)
diff --git a/compose/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/compiler/lower/ComposeObservePatcher.kt b/compose/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/compiler/lower/ComposeObservePatcher.kt
index 16979001663..73aa339020f 100644
--- a/compose/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/compiler/lower/ComposeObservePatcher.kt
+++ b/compose/compose-compiler-hosted/src/main/java/androidx/compose/plugins/kotlin/compiler/lower/ComposeObservePatcher.kt
@@ -48,6 +48,7 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtFunctionLiteral
 import androidx.compose.plugins.kotlin.ComposableAnnotationChecker
 import androidx.compose.plugins.kotlin.ComposeUtils.generateComposePackageName
+import androidx.compose.plugins.kotlin.KtxNameConventions
 import androidx.compose.plugins.kotlin.KtxNameConventions.UPDATE_SCOPE
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrReturn
@@ -60,6 +61,7 @@ import org.jetbrains.kotlin.backend.common.lower.irNot
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.PropertyDescriptor
 import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
+import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.builders.irBlock
 import org.jetbrains.kotlin.ir.builders.irEqeqeq
 import org.jetbrains.kotlin.ir.builders.irGet
@@ -68,17 +70,21 @@ import org.jetbrains.kotlin.ir.builders.irReturn
 import org.jetbrains.kotlin.ir.builders.irTemporary
 import org.jetbrains.kotlin.ir.declarations.IrDeclaration
 import org.jetbrains.kotlin.ir.expressions.IrBlockBody
+import org.jetbrains.kotlin.ir.expressions.IrBreak
 import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrContinue
 import org.jetbrains.kotlin.ir.expressions.impl.IrBlockImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrTryImpl
 import org.jetbrains.kotlin.ir.util.referenceFunction
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
 import org.jetbrains.kotlin.resolve.DelegatingBindingTrace
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.resolve.inline.InlineUtil
+import org.jetbrains.kotlin.resolve.scopes.MemberScope
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.typeUtil.isUnit
 
@@ -137,30 +143,41 @@ class ComposeObservePatcher(val context: JvmBackendContext) :
             }
         }
 
-        // Do not insert an observe scope if the funciton has a return result
+        // Do not insert an observe scope if the function has a return result
         if (descriptor.returnType.let { it == null || !it.isUnit() })
             return declaration
 
         // Check if the descriptor has restart scope calls resolved
         val bindingContext = context.state.bindingContext
-        if (descriptor is SimpleFunctionDescriptor) {
-            val restartCalls = bindingContext.get(ComposeWritableSlices.RESTART_CALLS, descriptor)
+        if (descriptor is SimpleFunctionDescriptor &&
+            // TODO(chuckj): Support lambdas
+            declaration.origin != IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA &&
+            declaration.origin != IrDeclarationOrigin.LOCAL_FUNCTION_NO_CLOSURE) {
+            val composerResolvedCall =
+                bindingContext.get(ComposeWritableSlices.RESTART_COMPOSER, descriptor)
 
             val oldBody = declaration.body
-            // TODO(b/142286425): Re-enable restart groups
-            if (restartCalls != null && oldBody != null /* disabled */ && descriptor.isInline) {
-                val composerResolvedCall = restartCalls.composer
-                val symbols = context.ir.symbols
-                val symbolTable = symbols.externalSymbolTable
-
-                // Create call to get the composer
-                val unitType = context.irBuiltIns.unitType
-                val composer = irGetProperty(composerResolvedCall)
-
-                // Create call to startRestartGroup
-                val startRestartGroup = irMethodCall(composer, restartCalls.startRestartGroup)
-                    .apply {
-                        putValueArgument(0,
+            if (composerResolvedCall != null && oldBody != null) {
+                val composerDescriptor =
+                    composerResolvedCall.resultingDescriptor as PropertyDescriptor
+                val startRestartGroupDescriptor = getStartRestartGroup(composerDescriptor)
+                val endRestartGroupDescriptor = getEndRestartGroup(composerDescriptor)
+
+                if (startRestartGroupDescriptor != null && endRestartGroupDescriptor != null) {
+                    val symbols = context.ir.symbols
+                    val symbolTable = symbols.externalSymbolTable
+
+                    // Create call to get the composer
+                    val unitType = context.irBuiltIns.unitType
+                    fun composer() = irGetProperty(composerResolvedCall)
+
+                    // Create call to startRestartGroup
+                    val startRestartGroup = irMethodCall(
+                        composer(),
+                        startRestartGroupDescriptor
+                    ).apply {
+                        putValueArgument(
+                            0,
                             keyExpression(
                                 descriptor,
                                 declaration.startOffset,
@@ -169,139 +186,148 @@ class ComposeObservePatcher(val context: JvmBackendContext) :
                         )
                     }
 
-                // Create call to endRestartGroup
-                val endRestartGroup = irMethodCall(composer, restartCalls.endRestartGroup)
-
-                // Create self-invoke lambda
-                val endRestartGroupResolvedCall = restartCalls.endRestartGroup
-                val endRestartGroupDescriptor = (
-                        endRestartGroupResolvedCall.resultingDescriptor as? FunctionDescriptor
-                    ) ?: error("Expected function descriptor")
-                val updateScopeDescriptor =
-                    endRestartGroupDescriptor.returnType?.memberScope?.getContributedFunctions(
-                        UPDATE_SCOPE,
-                        NoLookupLocation.FROM_BACKEND
-                    )?.singleOrNull()
-                        ?: error("updateScope not found in result type of endRestartGroup")
-                val blockParameterDescriptor = updateScopeDescriptor.valueParameters.singleOrNull()
-                    ?: error("expected a single block parameter for updateScope")
-                val blockParameterType = blockParameterDescriptor.type
-                val selfSymbol = declaration.symbol
-
-                val lambdaDescriptor = AnonymousFunctionDescriptor(
-                    declaration.descriptor,
-                    Annotations.EMPTY,
-                    CallableMemberDescriptor.Kind.DECLARATION,
-                    SourceElement.NO_SOURCE,
-                    false
-                ).apply {
-                    initialize(
-                        null,
-                        null,
-                        emptyList(),
-                        emptyList(),
-                        blockParameterType,
-                        Modality.FINAL,
-                        Visibilities.LOCAL
-                    )
-                }
+                    // Create call to endRestartGroup
+                    val endRestartGroup = irMethodCall(composer(), endRestartGroupDescriptor)
+
+                    // Create self-invoke lambda
+                    val updateScopeDescriptor =
+                        endRestartGroupDescriptor.returnType?.memberScope?.getContributedFunctions(
+                            UPDATE_SCOPE,
+                            NoLookupLocation.FROM_BACKEND
+                        )?.singleOrNull()
+                            ?: error("updateScope not found in result type of endRestartGroup")
+                    val blockParameterDescriptor =
+                        updateScopeDescriptor.valueParameters.singleOrNull()
+                            ?: error("expected a single block parameter for updateScope")
+                    val blockParameterType = blockParameterDescriptor.type
+                    val selfSymbol = declaration.symbol
+
+                    val lambdaDescriptor = AnonymousFunctionDescriptor(
+                        declaration.descriptor,
+                        Annotations.EMPTY,
+                        CallableMemberDescriptor.Kind.DECLARATION,
+                        SourceElement.NO_SOURCE,
+                        false
+                    ).apply {
+                        initialize(
+                            null,
+                            null,
+                            emptyList(),
+                            emptyList(),
+                            blockParameterType,
+                            Modality.FINAL,
+                            Visibilities.LOCAL
+                        )
+                    }
 
-                val fn = IrFunctionImpl(
-                    UNDEFINED_OFFSET, UNDEFINED_OFFSET,
-                    IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA,
-                    IrSimpleFunctionSymbolImpl(lambdaDescriptor),
-                    context.irBuiltIns.unitType
-                ).also {
-                    val irBuilder = context.createIrBuilder(it.symbol)
-                    it.body = irBuilder.irBlockBody {
-                        // Call the function again with the same parameters
-                        +irReturn(irCall(selfSymbol).apply {
-                            descriptor.valueParameters.forEachIndexed { index, valueParameter ->
-                                val value = declaration.valueParameters[index].symbol
-                                putValueArgument(
-                                    index, IrGetValueImpl(
-                                        UNDEFINED_OFFSET,
-                                        UNDEFINED_OFFSET,
-                                        valueParameter.type.toIrType(),
-                                        value
+                    val fn = IrFunctionImpl(
+                        UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+                        IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA,
+                        IrSimpleFunctionSymbolImpl(lambdaDescriptor),
+                        context.irBuiltIns.unitType
+                    ).also {
+                        it.parent = declaration
+                        val irBuilder = context.createIrBuilder(it.symbol)
+                        it.body = irBuilder.irBlockBody {
+                            // Call the function again with the same parameters
+                            +irReturn(irCall(selfSymbol).apply {
+                                descriptor.valueParameters.forEachIndexed { index, valueParameter ->
+                                    val value = declaration.valueParameters[index].symbol
+                                    putValueArgument(
+                                        index, IrGetValueImpl(
+                                            UNDEFINED_OFFSET,
+                                            UNDEFINED_OFFSET,
+                                            valueParameter.type.toIrType(),
+                                            value
+                                        )
                                     )
+                                }
+                                descriptor.dispatchReceiverParameter?.let { receiverDescriptor ->
+                                    dispatchReceiver = irGet(
+                                        receiverDescriptor.type.toIrType(),
+                                        declaration.dispatchReceiverParameter?.symbol
+                                            ?: error("Expected dispatch receiver on declaration")
+                                    )
+                                }
+                                descriptor.extensionReceiverParameter?.let { receiverDescriptor ->
+                                    extensionReceiver = irGet(
+                                        receiverDescriptor.type.toIrType(),
+                                        declaration.extensionReceiverParameter?.symbol
+                                            ?: error("Expected extension receiver on declaration")
+                                    )
+                                }
+                                descriptor.typeParameters.forEachIndexed { index, descriptor ->
+                                    putTypeArgument(index, descriptor.defaultType.toIrType())
+                                }
+                            })
+                        }
+                    }
+
+                    val irBuilder = context.createIrBuilder(declaration.symbol)
+                    val endRestartGroupCallBlock = irBuilder.irBlock(
+                        UNDEFINED_OFFSET,
+                        UNDEFINED_OFFSET
+                    ) {
+                        val result = irTemporary(endRestartGroup)
+                        val updateScopeSymbol =
+                            symbolTable.referenceSimpleFunction(updateScopeDescriptor)
+                        +irIfThen(irNot(irEqeqeq(irGet(result.type, result.symbol), irNull())),
+                            irCall(updateScopeSymbol).apply {
+                                dispatchReceiver = irGet(result.type, result.symbol)
+                                putValueArgument(0,
+                                    irBlock(origin = IrStatementOrigin.LAMBDA) {
+                                        +fn
+                                        +IrFunctionReferenceImpl(
+                                            UNDEFINED_OFFSET,
+                                            UNDEFINED_OFFSET,
+                                            blockParameterType.toIrType(),
+                                            fn.symbol,
+                                            fn.descriptor,
+                                            0,
+                                            IrStatementOrigin.LAMBDA
+                                        )
+                                    }
                                 )
                             }
-                            descriptor.dispatchReceiverParameter?.let { receiverDescriptor ->
-                                dispatchReceiver = irGet(receiverDescriptor.type.toIrType(),
-                                    declaration.dispatchReceiverParameter?.symbol
-                                        ?: error("Expected dispatch receiver on declaration")
-                                )
-                            }
-                            descriptor.extensionReceiverParameter?.let { receiverDescriptor ->
-                                extensionReceiver = irGet(receiverDescriptor.type.toIrType(),
-                                    declaration.extensionReceiverParameter?.symbol
-                                        ?: error("Expected extension receiver on declaration")
-                                )
-                            }
-                            descriptor.typeParameters.forEachIndexed { index, descriptor ->
-                                putTypeArgument(index, descriptor.defaultType.toIrType())
-                            }
-                        })
+                        )
                     }
-                }
 
-                val irBuilder = context.createIrBuilder(declaration.symbol)
-                val endRestartGroupCallBlock = irBuilder.irBlock(
-                    UNDEFINED_OFFSET,
-                    UNDEFINED_OFFSET
-                ) {
-                    val result = irTemporary(endRestartGroup)
-                    val updateScopeSymbol =
-                        symbolTable.referenceSimpleFunction(updateScopeDescriptor)
-                    +fn
-                    +irIfThen(irNot(irEqeqeq(irGet(result.type, result.symbol), irNull())),
-                        irCall(updateScopeSymbol).apply {
-                            dispatchReceiver = irGet(result.type, result.symbol)
-                            putValueArgument(0,
-                                IrFunctionReferenceImpl(
-                                    UNDEFINED_OFFSET,
-                                    UNDEFINED_OFFSET,
-                                    blockParameterType.toIrType(),
-                                    fn.symbol,
-                                    fn.descriptor,
-                                    0,
-                                    IrStatementOrigin.LAMBDA
-                                )
-                            )
+                    when (oldBody) {
+                        is IrBlockBody -> {
+                            if (containsPotentialEarly(oldBody)) {
+                                // Transform the block into
+                                // composer.startRestartGroup()
+                                // try {
+                                //   ... old statements ...
+                                // } finally {
+                                //    composer.endRestartGroup()
+                                // }
+                                declaration.body = irBuilder.irBlockBody {
+                                    +IrTryImpl(
+                                        oldBody.startOffset, oldBody.endOffset, unitType,
+                                        IrBlockImpl(
+                                            UNDEFINED_OFFSET,
+                                            UNDEFINED_OFFSET,
+                                            unitType
+                                        ).apply {
+                                            statements.add(startRestartGroup)
+                                            statements.addAll(oldBody.statements)
+                                        },
+                                        catches = emptyList(),
+                                        finallyExpression = endRestartGroupCallBlock
+                                    )
+                                }
+                            } else {
+                                // Insert the start and end calls into the block
+                                oldBody.statements.add(0, startRestartGroup)
+                                oldBody.statements.add(endRestartGroupCallBlock)
+                            }
+                            return declaration
                         }
-                    )
-                }
-
-                when (oldBody) {
-                    is IrBlockBody -> {
-                        // Transform the block into
-                        // composer.startRestartGroup()
-                        // try {
-                        //   ... old statements ...
-                        // } finally {
-                        //    composer.endRestartGroup()
-                        // }
-                        declaration.body = irBuilder.irBlockBody {
-                            +IrTryImpl(
-                                oldBody.startOffset, oldBody.endOffset, unitType,
-                                IrBlockImpl(
-                                    UNDEFINED_OFFSET,
-                                    UNDEFINED_OFFSET,
-                                    unitType
-                                ).apply {
-                                    statements.add(startRestartGroup)
-                                    statements.addAll(oldBody.statements)
-                                },
-                                catches = emptyList(),
-                                finallyExpression = endRestartGroupCallBlock
-                            )
+                        else -> {
+                            // Composable function do not use IrExpressionBody as they are converted
+                            // by the call lowering to IrBlockBody to introduce the call temporaries.
                         }
-                        return declaration
-                    }
-                    else -> {
-                        // Composable function do not use IrExpressionBody as they are converted
-                        // by the call lowering to IrBlockBody to introduce the call temporaries.
                     }
                 }
             }
@@ -427,11 +453,30 @@ class ComposeObservePatcher(val context: JvmBackendContext) :
     fun irMethodCall(target: IrExpression, resolvedCall: ResolvedCall<*>): IrCall {
         val descriptor = (resolvedCall.resultingDescriptor as?
                 FunctionDescriptor) ?: error("Expected function descriptor")
+        return irMethodCall(target, descriptor)
+    }
+
+    fun irMethodCall(target: IrExpression, descriptor: FunctionDescriptor): IrCall {
         return irCall(descriptor).apply {
             dispatchReceiver = target
         }
     }
 
+    fun MemberScope.firstFunctionOrNull(name: Name) =
+        getContributedFunctions(name, NoLookupLocation.FROM_BACKEND).firstOrNull()
+
+    fun getStartRestartGroup(descriptor: PropertyDescriptor): FunctionDescriptor? =
+        descriptor
+            .type
+            .memberScope
+            .firstFunctionOrNull(KtxNameConventions.STARTRESTARTGROUP)
+
+    fun getEndRestartGroup(descriptor: PropertyDescriptor): FunctionDescriptor? =
+        descriptor
+            .type
+            .memberScope
+            .firstFunctionOrNull(KtxNameConventions.ENDRESTARTGROUP)
+
     private fun isComposable(declaration: IrFunction): Boolean {
         val tmpTrace =
             DelegatingBindingTrace(
@@ -464,3 +509,30 @@ class ComposeObservePatcher(val context: JvmBackendContext) :
         )
     }
 }
+
+private fun containsPotentialEarly(block: IrBlockBody): Boolean {
+    var result = false
+    block.accept(object : IrElementVisitor<Unit, Unit> {
+        override fun visitElement(element: IrElement, data: Unit) {
+            if (!result)
+                element.acceptChildren(this, Unit)
+        }
+
+        override fun visitBreak(jump: IrBreak, data: Unit) {
+            result = true
+        }
+
+        override fun visitContinue(jump: IrContinue, data: Unit) {
+            result = true
+        }
+
+        override fun visitReturn(expression: IrReturn, data: Unit) {
+            result = true
+        }
+
+        override fun visitDeclaration(declaration: IrDeclaration, data: Unit) {
+            // Skip bodies of declarations
+        }
+    }, Unit)
+    return result
+}
\ No newline at end of file
diff --git a/compose/compose-runtime/src/androidAndroidTest/kotlin/androidx/compose/ViewComposerTests.kt b/compose/compose-runtime/src/androidAndroidTest/kotlin/androidx/compose/ViewComposerTests.kt
index eb08c5f3d36..faa8f26369f 100644
--- a/compose/compose-runtime/src/androidAndroidTest/kotlin/androidx/compose/ViewComposerTests.kt
+++ b/compose/compose-runtime/src/androidAndroidTest/kotlin/androidx/compose/ViewComposerTests.kt
@@ -32,6 +32,8 @@ import junit.framework.TestCase.assertEquals
 import org.junit.Rule
 import org.junit.runner.RunWith
 import org.junit.Test
+import java.lang.AssertionError
+import java.lang.IllegalStateException
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.TimeUnit
 import kotlin.test.assertTrue
@@ -1025,11 +1027,32 @@ class DisposeTests {
 internal val Activity.root get() = findViewById(ComposerComposeTestCase.ROOT_ID) as ViewGroup
 
 internal fun Activity.uiThread(block: () -> Unit) {
+    val latch = CountDownLatch(1)
+    var throwable: Throwable? = null
     runOnUiThread(object : Runnable {
         override fun run() {
-            block()
+            try {
+                block()
+            } catch (e: Throwable) {
+                throwable = e
+            } finally {
+                latch.countDown()
+            }
         }
     })
+
+    val completed = latch.await(5, TimeUnit.SECONDS)
+    if (!completed) error("UI thread work did not complete within 5 seconds")
+    throwable?.let {
+        throw when (it) {
+            is AssertionError -> AssertionError(it.localizedMessage, it)
+            else ->
+                IllegalStateException(
+                    "UI thread threw an exception: ${it.localizedMessage}",
+                    it
+                )
+        }
+    }
 }
 
 internal fun Activity.disposeTestComposition() {
diff --git a/compose/compose-runtime/src/commonMain/kotlin/androidx/compose/SlotTable.kt b/compose/compose-runtime/src/commonMain/kotlin/androidx/compose/SlotTable.kt
index 4e75e4c95c8..092ed335b3e 100644
--- a/compose/compose-runtime/src/commonMain/kotlin/androidx/compose/SlotTable.kt
+++ b/compose/compose-runtime/src/commonMain/kotlin/androidx/compose/SlotTable.kt
@@ -671,6 +671,8 @@ class SlotTable(internal var slots: Array<Any?> = arrayOf()) {
         return SlotWriter(this)
     }
 
+    val size: Int get() = slots.size - gapLen
+
     internal fun close(reader: SlotReader) {
         require(reader.table === this && readers > 0) { "Unexpected reader close()" }
         readers--