From a8249d60c7cd14a459469fe0ef2099721d3dd699 Mon Sep 17 00:00:00 2001
From: Andrei Shikov <ashikov@google.com>
Date: Wed, 15 May 2024 19:00:18 +0100
Subject: [PATCH] Allow memoizing lambdas in composable inline functions

Memoization was disabled in inline functions for legacy reasons while technically it can be allowed.

Test: compiler tests
Fixes: 340606661 ( https://issuetracker.google.com/issues/340606661 )
Change-Id: Iee8edda58f1e3a2b7e7dcc5ef1abcc663fab7fda ( https://android-review.googlesource.com/q/Iee8edda58f1e3a2b7e7dcc5ef1abcc663fab7fda )

Moved from: https://github.com/androidx/androidx/commit/dc92291c662e0c843fb6888e160da19e0d5bb86d
---
 .../kotlin/LambdaMemoizationRegressionTests.kt  | 17 +++++++++++++++++
 .../kotlin/lower/ComposerLambdaMemoization.kt   |  8 ++------
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/LambdaMemoizationRegressionTests.kt b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/LambdaMemoizationRegressionTests.kt
index 847e4c90bb34c..e718729142710 100644
--- a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/LambdaMemoizationRegressionTests.kt
+++ b/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/LambdaMemoizationRegressionTests.kt
@@ -60,4 +60,21 @@ class LambdaMemoizationRegressionTests(useFir: Boolean) : AbstractIrTransformTes
             val x = @Composable {}
         """
     )
+
+    // regression test for b/340606661
+    @Test
+    fun testMemoizationInInlineFunction() = verifyGoldenComposeIrTransform(
+        """
+            import androidx.compose.runtime.*
+
+            @Composable
+            inline fun Test(
+                someBool: Boolean,
+            ) {
+                val someInt = remember { 1 }
+                val lambda = { someInt }
+                println(lambda.hashCode())
+            }
+        """
+    )
 }
diff --git a/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposerLambdaMemoization.kt b/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposerLambdaMemoization.kt
index 8968c2a9bd92b..d7666cbf023e8 100644
--- a/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposerLambdaMemoization.kt
+++ b/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposerLambdaMemoization.kt
@@ -202,13 +202,13 @@ private class FunctionLocalSymbol(
 private class FunctionContext(
     override val declaration: IrFunction,
     override val composable: Boolean,
-    val canRemember: Boolean
 ) : DeclarationContext() {
     override val symbol get() = declaration.symbol
     override val functionContext: FunctionContext get() = this
     val locals = mutableSetOf<IrValueDeclaration>()
     override val captures: MutableSet<IrValueDeclaration> = mutableSetOf()
     var collectors = mutableListOf<CaptureCollector>()
+    val canRemember: Boolean get() = composable
 
     init {
         declaration.valueParameters.forEach {
@@ -459,11 +459,7 @@ class ComposerLambdaMemoization(
 
     override fun visitFunction(declaration: IrFunction): IrStatement {
         val composable = declaration.allowsComposableCalls
-        val canRemember = composable &&
-            // Don't use remember in an inline function
-            !declaration.isInline
-
-        val context = FunctionContext(declaration, composable, canRemember)
+        val context = FunctionContext(declaration, composable)
         if (declaration.isLocal) {
             declarationContextStack.recordLocalDeclaration(context)
         }