Stop transforming non-composable fun interfaces in Compose compiler

The original bug (b/178663739) is already fixed upstream, so we can remove this transform, as it causes problem with Java sam conversions in K2.
I left the conversion for composable functions, as SAM conversion does not support restartable lambdas otherwise.

The side effect of this change is that SAM-converted lambdas are now memoized by Compose compiler. This CL also adjusts this behavior by ensuring remember happens outside of conversion. Without this fix, the converted class was delegating to remembered lambda, essentially un-doing the optimization.

This commit also improves the lowering for cases when SAM_CONVERSION was followed by IMPLICIT_CAST (see updates in `ComposerParamSignatureTest` as sample).

Test: Updated compiler tests + runtime tests
Bug: 178663739
Fixes: 288403625
Change-Id: Ieabb4a9df24f77e9b5e7e41adf360c50cede604f
diff --git a/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/ComposerParamSignatureTests.kt b/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/ComposerParamSignatureTests.kt
index 5d649af..f476366 100644
--- a/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/ComposerParamSignatureTests.kt
+++ b/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/ComposerParamSignatureTests.kt
@@ -1471,15 +1471,7 @@
             public final class TestKt {
               public final static Example(LA;)V
               public final static Usage()V
-              final static INNERCLASS TestKt%Usage%1 null null
-            }
-            final class TestKt%Usage%1 implements A {
-              <init>()V
-              public final compute(I)V
-              static <clinit>()V
-              public final static LTestKt%Usage%1; INSTANCE
-              OUTERCLASS TestKt Usage ()V
-              final static INNERCLASS TestKt%Usage%1 null null
+              private final static Usage%lambda%0(I)V
             }
         """
     )
@@ -1524,6 +1516,93 @@
         """
     )
 
+    @Test
+    fun testFunInterfacesInComposableCall() = checkApi(
+        """
+            fun interface MeasurePolicy {
+                fun compute(value: Int): Unit
+            }
+
+            @NonRestartableComposable
+            @Composable fun Text() {
+                Layout { value ->
+                    println(value)
+                }
+            }
+
+            @Composable inline fun Layout(policy: MeasurePolicy) {
+                policy.compute(0)
+            }
+        """,
+        """
+            public abstract interface MeasurePolicy {
+              public abstract compute(I)V
+            }
+            public final class TestKt {
+              public final static Text(Landroidx/compose/runtime/Composer;I)V
+              public final static Layout(LMeasurePolicy;Landroidx/compose/runtime/Composer;I)V
+              private final static Text%lambda%0(I)V
+            }
+        """,
+    )
+
+    @Test
+        fun testComposableFunInterfacesInVariance() = checkApi(
+        """
+           import androidx.compose.runtime.*
+
+            fun interface Consumer<T> {
+                @Composable fun consume(t: T)
+            }
+
+            class Repro<T : Any>() {
+                fun test(consumer: Consumer<in T>) {}
+            }
+
+            fun test() {
+                Repro<String>().test { string ->
+                    println(string)
+                }
+            }
+        """,
+        """
+            public abstract interface Consumer {
+              public abstract consume(Ljava/lang/Object;Landroidx/compose/runtime/Composer;I)V
+            }
+            public final class Repro {
+              public <init>()V
+              public final test(LConsumer;)V
+              static <clinit>()V
+              public final static I %stable
+            }
+            public final class TestKt {
+              public final static test()V
+              final static INNERCLASS TestKt%test%1 null null
+            }
+            final class TestKt%test%1 implements Consumer {
+              <init>()V
+              public final consume(Ljava/lang/String;Landroidx/compose/runtime/Composer;I)V
+              public synthetic bridge consume(Ljava/lang/Object;Landroidx/compose/runtime/Composer;I)V
+              static <clinit>()V
+              public final static LTestKt%test%1; INSTANCE
+              OUTERCLASS TestKt test ()V
+              final static INNERCLASS TestKt%test%1 null null
+              final static INNERCLASS TestKt%test%1%consume%1 null null
+            }
+            final class TestKt%test%1%consume%1 extends kotlin/jvm/internal/Lambda implements kotlin/jvm/functions/Function2 {
+              <init>(LTestKt%test%1;Ljava/lang/String;I)V
+              public final invoke(Landroidx/compose/runtime/Composer;I)V
+              public synthetic bridge invoke(Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;
+              final synthetic LTestKt%test%1; %tmp0_rcvr
+              final synthetic Ljava/lang/String; %string
+              final synthetic I %%changed
+              OUTERCLASS TestKt%test%1 consume (Ljava/lang/String;Landroidx/compose/runtime/Composer;I)V
+              final static INNERCLASS TestKt%test%1 null null
+              final static INNERCLASS TestKt%test%1%consume%1 null null
+            }
+        """
+    )
+
     val hashCodeEqualsAndToString = if (useFir) {
         """
               public static equals-impl(ILjava/lang/Object;)Z
diff --git a/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/FunctionalInterfaceExtensionReceiverTransformTests.kt b/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/FunctionalInterfaceExtensionReceiverTransformTests.kt
deleted file mode 100644
index 29e745d..0000000
--- a/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/FunctionalInterfaceExtensionReceiverTransformTests.kt
+++ /dev/null
@@ -1,124 +0,0 @@
-/*
- * Copyright 2023 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.compiler.plugins.kotlin
-
-import org.junit.Test
-
-class FunctionalInterfaceExtensionReceiverTransformTests(
-    useFir: Boolean
-) : AbstractControlFlowTransformTests(useFir) {
-    @Test
-    fun testFunctionalInterfaceWithExtensionReceiverTransformation() {
-        verifyComposeIrTransform(
-            source = """
-                import androidx.compose.runtime.*
-                fun interface TestContent {
-                    @Composable
-                    fun String.Content()
-                }
-                @Composable
-                fun Test(content: TestContent) {
-                    with(content) {
-                        "".Content()
-                    }
-                }
-
-                @Composable
-                fun CallTest() {
-                    Test { this.length }
-                }
-            """.trimIndent(),
-            expectedTransformed = """
-            interface TestContent {
-              @Composable
-              abstract fun String.Content(%composer: Composer?, %changed: Int)
-            }
-            @Composable
-            @ComposableInferredTarget(scheme = "[0[0]]")
-            fun Test(content: TestContent, %composer: Composer?, %changed: Int) {
-              %composer = %composer.startRestartGroup(<>)
-              sourceInformation(%composer, "C(Test)*<Conten...>:Test.kt")
-              val %dirty = %changed
-              if (%changed and 0b1110 === 0) {
-                %dirty = %dirty or if (%composer.changed(content)) 0b0100 else 0b0010
-              }
-              if (%dirty and 0b1011 !== 0b0010 || !%composer.skipping) {
-                if (isTraceInProgress()) {
-                  traceEventStart(<>, %changed, -1, <>)
-                }
-                with(content) {
-                  %this%with.Content(%composer, 0b0110)
-                }
-                if (isTraceInProgress()) {
-                  traceEventEnd()
-                }
-              } else {
-                %composer.skipToGroupEnd()
-              }
-              %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
-                Test(content, %composer, updateChangedFlags(%changed or 0b0001))
-              }
-            }
-            @Composable
-            fun CallTest(%composer: Composer?, %changed: Int) {
-              %composer = %composer.startRestartGroup(<>)
-              sourceInformation(%composer, "C(CallTest)<Test>:Test.kt")
-              if (%changed !== 0 || !%composer.skipping) {
-                if (isTraceInProgress()) {
-                  traceEventStart(<>, %changed, -1, <>)
-                }
-                Test(class <no name provided> : TestContent {
-                  @Composable
-                  override fun Content(%this%Test: String, %composer: Composer?, %changed: Int) {
-                    %composer = %composer.startRestartGroup(<>)
-                    sourceInformation(%composer, "C(Content):Test.kt")
-                    val %dirty = %changed
-                    if (%changed and 0b1110 === 0) {
-                      %dirty = %dirty or if (%composer.changed(%this%Test)) 0b0100 else 0b0010
-                    }
-                    if (%dirty and 0b1011 !== 0b0010 || !%composer.skipping) {
-                      if (isTraceInProgress()) {
-                        traceEventStart(<>, %changed, -1, <>)
-                      }
-                      %this%Test.length
-                      if (isTraceInProgress()) {
-                        traceEventEnd()
-                      }
-                    } else {
-                      %composer.skipToGroupEnd()
-                    }
-                    val tmp0_rcvr = <this>
-                    %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
-                      tmp0_rcvr.Content(%this%Test, %composer, updateChangedFlags(%changed or 0b0001))
-                    }
-                  }
-                }
-                <no name provided>(), %composer, 0)
-                if (isTraceInProgress()) {
-                  traceEventEnd()
-                }
-              } else {
-                %composer.skipToGroupEnd()
-              }
-              %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
-                CallTest(%composer, updateChangedFlags(%changed or 0b0001))
-              }
-            }
-            """.trimIndent()
-        )
-    }
-}
diff --git a/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/FunctionalInterfaceTransformTests.kt b/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/FunctionalInterfaceTransformTests.kt
new file mode 100644
index 0000000..606cc70
--- /dev/null
+++ b/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/FunctionalInterfaceTransformTests.kt
@@ -0,0 +1,390 @@
+/*
+ * Copyright 2023 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.compiler.plugins.kotlin
+
+import org.junit.Test
+
+class FunctionalInterfaceTransformTests(
+    useFir: Boolean
+) : AbstractControlFlowTransformTests(useFir) {
+    @Test
+    fun testFunctionalInterfaceWithExtensionReceiverTransformation() {
+        verifyComposeIrTransform(
+            source = """
+                import androidx.compose.runtime.*
+                fun interface TestContent {
+                    @Composable
+                    fun String.Content()
+                }
+                @Composable
+                fun Test(content: TestContent) {
+                    with(content) {
+                        "".Content()
+                    }
+                }
+
+                @Composable
+                fun CallTest() {
+                    Test { this.length }
+                }
+            """.trimIndent(),
+            expectedTransformed = """
+            interface TestContent {
+              @Composable
+              abstract fun String.Content(%composer: Composer?, %changed: Int)
+            }
+            @Composable
+            @ComposableInferredTarget(scheme = "[0[0]]")
+            fun Test(content: TestContent, %composer: Composer?, %changed: Int) {
+              %composer = %composer.startRestartGroup(<>)
+              sourceInformation(%composer, "C(Test)*<Conten...>:Test.kt")
+              val %dirty = %changed
+              if (%changed and 0b1110 === 0) {
+                %dirty = %dirty or if (%composer.changed(content)) 0b0100 else 0b0010
+              }
+              if (%dirty and 0b1011 !== 0b0010 || !%composer.skipping) {
+                if (isTraceInProgress()) {
+                  traceEventStart(<>, %changed, -1, <>)
+                }
+                with(content) {
+                  %this%with.Content(%composer, 0b0110)
+                }
+                if (isTraceInProgress()) {
+                  traceEventEnd()
+                }
+              } else {
+                %composer.skipToGroupEnd()
+              }
+              %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+                Test(content, %composer, updateChangedFlags(%changed or 0b0001))
+              }
+            }
+            @Composable
+            fun CallTest(%composer: Composer?, %changed: Int) {
+              %composer = %composer.startRestartGroup(<>)
+              sourceInformation(%composer, "C(CallTest)<Test>:Test.kt")
+              if (%changed !== 0 || !%composer.skipping) {
+                if (isTraceInProgress()) {
+                  traceEventStart(<>, %changed, -1, <>)
+                }
+                Test(class <no name provided> : TestContent {
+                  @Composable
+                  override fun Content(%this%Test: String, %composer: Composer?, %changed: Int) {
+                    %composer = %composer.startRestartGroup(<>)
+                    sourceInformation(%composer, "C(Content):Test.kt")
+                    val %dirty = %changed
+                    if (%changed and 0b1110 === 0) {
+                      %dirty = %dirty or if (%composer.changed(%this%Test)) 0b0100 else 0b0010
+                    }
+                    if (%dirty and 0b1011 !== 0b0010 || !%composer.skipping) {
+                      if (isTraceInProgress()) {
+                        traceEventStart(<>, %changed, -1, <>)
+                      }
+                      %this%Test.length
+                      if (isTraceInProgress()) {
+                        traceEventEnd()
+                      }
+                    } else {
+                      %composer.skipToGroupEnd()
+                    }
+                    val tmp0_rcvr = <this>
+                    %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+                      tmp0_rcvr.Content(%this%Test, %composer, updateChangedFlags(%changed or 0b0001))
+                    }
+                  }
+                }
+                <no name provided>(), %composer, 0)
+                if (isTraceInProgress()) {
+                  traceEventEnd()
+                }
+              } else {
+                %composer.skipToGroupEnd()
+              }
+              %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+                CallTest(%composer, updateChangedFlags(%changed or 0b0001))
+              }
+            }
+            """.trimIndent()
+        )
+    }
+
+    @Test
+    fun testFunInterfaces() = verifyComposeIrTransform(
+        """
+            import androidx.compose.runtime.*
+
+            fun interface A {
+                fun compute(value: Int): Unit
+            }
+
+            @Composable
+            fun Example(a: A) {
+                Example { it -> a.compute(it) }
+            }
+        """,
+        """
+            interface A {
+              abstract fun compute(value: Int)
+            }
+            @Composable
+            fun Example(a: A, %composer: Composer?, %changed: Int) {
+              %composer = %composer.startRestartGroup(<>)
+              sourceInformation(%composer, "C(Example)<Exampl...>:Test.kt")
+              val %dirty = %changed
+              if (%changed and 0b1110 === 0) {
+                %dirty = %dirty or if (%composer.changed(a)) 0b0100 else 0b0010
+              }
+              if (%dirty and 0b1011 !== 0b0010 || !%composer.skipping) {
+                if (isTraceInProgress()) {
+                  traceEventStart(<>, %changed, -1, <>)
+                }
+                Example(A { it: Int ->
+                  a.compute(it)
+                }, %composer, 0)
+                if (isTraceInProgress()) {
+                  traceEventEnd()
+                }
+              } else {
+                %composer.skipToGroupEnd()
+              }
+              %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+                Example(a, %composer, updateChangedFlags(%changed or 0b0001))
+              }
+            }
+        """
+    )
+
+    @Test
+    fun testComposableFunInterfaces() = verifyComposeIrTransform(
+        """
+            import androidx.compose.runtime.*
+
+            fun interface A {
+                @Composable fun compute(value: Int): Unit
+            }
+            fun Example(a: A) {
+                Example { it -> a.compute(it) }
+            }
+        """,
+        """
+            interface A {
+              @Composable
+              abstract fun compute(value: Int, %composer: Composer?, %changed: Int)
+            }
+            fun Example(a: A) {
+              Example(class <no name provided> : A {
+                @Composable
+                override fun compute(it: Int, %composer: Composer?, %changed: Int) {
+                  %composer = %composer.startRestartGroup(<>)
+                  sourceInformation(%composer, "C(compute)<comput...>:Test.kt")
+                  val %dirty = %changed
+                  if (%changed and 0b1110 === 0) {
+                    %dirty = %dirty or if (%composer.changed(it)) 0b0100 else 0b0010
+                  }
+                  if (%dirty and 0b1011 !== 0b0010 || !%composer.skipping) {
+                    if (isTraceInProgress()) {
+                      traceEventStart(<>, %dirty, -1, <>)
+                    }
+                    a.compute(it, %composer, 0b1110 and %dirty)
+                    if (isTraceInProgress()) {
+                      traceEventEnd()
+                    }
+                  } else {
+                    %composer.skipToGroupEnd()
+                  }
+                  val tmp0_rcvr = <this>
+                  %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+                    tmp0_rcvr.compute(it, %composer, updateChangedFlags(%changed or 0b0001))
+                  }
+                }
+              }
+              <no name provided>())
+            }
+        """
+    )
+
+    @Test
+    fun testComposableFunInterfacesInVariance() = verifyComposeIrTransform(
+        """
+            import androidx.compose.runtime.*
+
+            fun interface Consumer<T> {
+                @Composable fun consume(t: T)
+            }
+
+            class Repro<T : Any> {
+                fun test(consumer: Consumer<in T>) {}
+            }
+
+            fun test() {
+                Repro<String>().test { string ->
+                    println(string)
+                }
+            }
+        """,
+        """
+            interface Consumer<T>  {
+              @Composable
+              abstract fun consume(t: T, %composer: Composer?, %changed: Int)
+            }
+            @StabilityInferred(parameters = 0)
+            class Repro<T: Any>  {
+              fun test(consumer: Consumer<in T>) { }
+              static val %stable: Int = 0
+            }
+            fun test() {
+              Repro().test(class <no name provided> : Consumer<Any?> {
+                @Composable
+                override fun consume(string: String, %composer: Composer?, %changed: Int) {
+                  %composer = %composer.startRestartGroup(<>)
+                  sourceInformation(%composer, "C(consume):Test.kt")
+                  val %dirty = %changed
+                  if (%changed and 0b1110 === 0) {
+                    %dirty = %dirty or if (%composer.changed(string)) 0b0100 else 0b0010
+                  }
+                  if (%dirty and 0b1011 !== 0b0010 || !%composer.skipping) {
+                    if (isTraceInProgress()) {
+                      traceEventStart(<>, %changed, -1, <>)
+                    }
+                    println(string)
+                    if (isTraceInProgress()) {
+                      traceEventEnd()
+                    }
+                  } else {
+                    %composer.skipToGroupEnd()
+                  }
+                  val tmp0_rcvr = <this>
+                  %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+                    tmp0_rcvr.consume(string, %composer, updateChangedFlags(%changed or 0b0001))
+                  }
+                }
+              }
+              <no name provided>())
+            }
+        """
+    )
+
+    @Test
+    fun testCaptureStableFunInterface() = verifyComposeIrTransform(
+        """
+            import androidx.compose.runtime.*
+
+            fun interface Consumer {
+                fun consume(t: Int)
+            }
+
+            @Composable fun Test(int: Int) {
+                Example {
+                    println(int)
+                }
+            }
+
+            @Composable inline fun Example(consumer: Consumer) {
+            }
+        """,
+        """
+            interface Consumer {
+              abstract fun consume(t: Int)
+            }
+            @Composable
+            fun Test(int: Int, %composer: Composer?, %changed: Int) {
+              %composer = %composer.startRestartGroup(<>)
+              sourceInformation(%composer, "C(Test)<{>,<Exampl...>:Test.kt")
+              val %dirty = %changed
+              if (%changed and 0b1110 === 0) {
+                %dirty = %dirty or if (%composer.changed(int)) 0b0100 else 0b0010
+              }
+              if (%dirty and 0b1011 !== 0b0010 || !%composer.skipping) {
+                if (isTraceInProgress()) {
+                  traceEventStart(<>, %dirty, -1, <>)
+                }
+                Example(remember(int, {
+                  Consumer { it: Int ->
+                    println(int)
+                  }
+                }, %composer, 0b1110 and %dirty), %composer, 0)
+                if (isTraceInProgress()) {
+                  traceEventEnd()
+                }
+              } else {
+                %composer.skipToGroupEnd()
+              }
+              %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+                Test(int, %composer, updateChangedFlags(%changed or 0b0001))
+              }
+            }
+            @Composable
+            fun Example(consumer: Consumer, %composer: Composer?, %changed: Int) {
+              %composer.startReplaceableGroup(<>)
+              sourceInformation(%composer, "CC(Example):Test.kt")
+              %composer.endReplaceableGroup()
+            }
+        """
+    )
+
+    @Test
+    fun testNoCaptureFunInterface() = verifyComposeIrTransform(
+        """
+            import androidx.compose.runtime.*
+
+            fun interface Consumer {
+                fun consume(t: Int)
+            }
+
+            @Composable fun Test(int: Int) {
+                Example {
+                    println(it)
+                }
+            }
+
+            @Composable inline fun Example(consumer: Consumer) {
+            }
+        """,
+        """
+            interface Consumer {
+              abstract fun consume(t: Int)
+            }
+            @Composable
+            fun Test(int: Int, %composer: Composer?, %changed: Int) {
+              %composer = %composer.startRestartGroup(<>)
+              sourceInformation(%composer, "C(Test)<Exampl...>:Test.kt")
+              if (%changed and 0b0001 !== 0 || !%composer.skipping) {
+                if (isTraceInProgress()) {
+                  traceEventStart(<>, %changed, -1, <>)
+                }
+                Example(Consumer { it: Int ->
+                  println(it)
+                }, %composer, 0b0110)
+                if (isTraceInProgress()) {
+                  traceEventEnd()
+                }
+              } else {
+                %composer.skipToGroupEnd()
+              }
+              %composer.endRestartGroup()?.updateScope { %composer: Composer?, %force: Int ->
+                Test(int, %composer, updateChangedFlags(%changed or 0b0001))
+              }
+            }
+            @Composable
+            fun Example(consumer: Consumer, %composer: Composer?, %changed: Int) {
+              %composer.startReplaceableGroup(<>)
+              sourceInformation(%composer, "CC(Example):Test.kt")
+              %composer.endReplaceableGroup()
+            }
+        """
+    )
+}
diff --git a/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/TargetAnnotationsTransformTests.kt b/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/TargetAnnotationsTransformTests.kt
index d3f96e3..3817942 100644
--- a/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/TargetAnnotationsTransformTests.kt
+++ b/compose/compiler/compiler-hosted/integration-tests/src/test/java/androidx/compose/compiler/plugins/kotlin/TargetAnnotationsTransformTests.kt
@@ -890,16 +890,13 @@
           }
           static val %stable: Int = 0
         }
-        val localBoxMeasurePolicy: MeasurePolicy = class <no name provided> : MeasurePolicy {
-          override fun measure(%this%MeasurePolicy: MeasureScope, <unused var>: List<Measurable>, constraints: Constraints): MeasureResult {
-            return %this%MeasurePolicy.layout(
-              width = constraints.minWidth,
-              height = constraints.minHeight
-            ) {
-            }
+        val localBoxMeasurePolicy: MeasurePolicy = MeasurePolicy { <unused var>: List<Measurable>, constraints: Constraints ->
+          %this%MeasurePolicy.layout(
+            width = constraints.minWidth,
+            height = constraints.minHeight
+          ) {
           }
         }
-        <no name provided>()
         @Composable
         @ComposableInferredTarget(scheme = "[androidx.compose.ui.UiComposable[androidx.compose.ui.UiComposable]]")
         fun LocalBox(modifier: Modifier?, content: @[ExtensionFunctionType] Function3<LocalBoxScope, Composer, Int, Unit>, %composer: Composer?, %changed: Int, %default: Int) {
@@ -983,12 +980,9 @@
               sourceInformationMarkerStart(%composer, <>, "C:Test.kt")
               Unit
               sourceInformationMarkerEnd(%composer)
-            }, null, class <no name provided> : MeasurePolicy {
-              override fun measure(%this%Layout: MeasureScope, <unused var>: List<Measurable>, <unused var>: Constraints): MeasureResult {
-                ${if (!useFir) "return " else ""}error("")
-              }
-            }
-            <no name provided>(), %composer, 0, 0b0010)
+            }, null, MeasurePolicy { <unused var>: List<Measurable>, <unused var>: Constraints ->
+              error("")
+            }, %composer, 0b000110000000, 0b0010)
             if (isTraceInProgress()) {
               traceEventEnd()
             }
@@ -1012,12 +1006,9 @@
             if (isTraceInProgress()) {
               traceEventStart(<>, %dirty, -1, <>)
             }
-            Layout(content, null, class <no name provided> : MeasurePolicy {
-              override fun measure(%this%Layout: MeasureScope, <unused var>: List<Measurable>, <unused var>: Constraints): MeasureResult {
-                ${if (!useFir) "return " else ""}error("")
-              }
-            }
-            <no name provided>(), %composer, 0b1110 and %dirty, 0b0010)
+            Layout(content, null, MeasurePolicy { <unused var>: List<Measurable>, <unused var>: Constraints ->
+              error("")
+            }, %composer, 0b000110000000 or 0b1110 and %dirty, 0b0010)
             if (isTraceInProgress()) {
               traceEventEnd()
             }
diff --git a/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/AbstractComposeLowering.kt b/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/AbstractComposeLowering.kt
index e1cd194b..066e0d3 100644
--- a/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/AbstractComposeLowering.kt
+++ b/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/AbstractComposeLowering.kt
@@ -70,6 +70,7 @@
 import org.jetbrains.kotlin.ir.expressions.IrMemberAccessExpression
 import org.jetbrains.kotlin.ir.expressions.IrReturn
 import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
 import org.jetbrains.kotlin.ir.expressions.IrVararg
 import org.jetbrains.kotlin.ir.expressions.impl.IrBlockImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrBranchImpl
@@ -886,7 +887,8 @@
                     else -> false
                 }
             }
-            is IrFunctionExpression ->
+            is IrFunctionExpression,
+            is IrTypeOperatorCall ->
                 context.irTrace[ComposeWritableSlices.IS_STATIC_FUNCTION_EXPRESSION, this] ?: false
             is IrGetField ->
                 // K2 sometimes produces `IrGetField` for reads from constant properties
diff --git a/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposableFunInterfaceLowering.kt b/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposableFunInterfaceLowering.kt
index cb8f689..9bacca8 100644
--- a/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposableFunInterfaceLowering.kt
+++ b/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposableFunInterfaceLowering.kt
@@ -16,15 +16,20 @@
 
 package androidx.compose.compiler.plugins.kotlin.lower
 
+import androidx.compose.compiler.plugins.kotlin.hasComposableAnnotation
 import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
+import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrFunctionExpression
-import org.jetbrains.kotlin.ir.expressions.IrTypeOperator
+import org.jetbrains.kotlin.ir.expressions.IrTypeOperator.IMPLICIT_CAST
+import org.jetbrains.kotlin.ir.expressions.IrTypeOperator.SAM_CONVERSION
 import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
+import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.IrTypeSystemContextImpl
 import org.jetbrains.kotlin.ir.types.classOrNull
+import org.jetbrains.kotlin.ir.util.functions
 import org.jetbrains.kotlin.ir.util.isLambda
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.platform.jvm.isJvm
@@ -40,32 +45,10 @@
         }
     }
 
-    private fun isFunInterfaceConversion(expression: IrTypeOperatorCall): Boolean {
-        val argument = expression.argument
-        val operator = expression.operator
-        val type = expression.typeOperand
-        val functionClass = type.classOrNull
-        return operator == IrTypeOperator.SAM_CONVERSION &&
-            argument is IrFunctionExpression &&
-            argument.origin.isLambda &&
-            functionClass != null &&
-            functionClass.owner.isFun
-        // IMPORTANT(b/178663739):
-        // We are transforming not just SAM conversions for composable fun interfaces, but ALL
-        // fun interfaces temporarily until KT-44622 gets fixed in the version of kotlin we
-        // are using, which should be in 1.4.30.
-        // Once it does, we should either add the below additional condition to this predicate,
-        // or, if possible, remove this lowering all together if kotlin's lowering works for
-        // composable fun interfaces as well.
-        //
-        // functionClass.functions.single {
-        //    it.owner.modality == Modality.ABSTRACT
-        // }.owner.annotations.hasAnnotation(ComposeFqNames.Composable)
-    }
-
     override fun visitTypeOperator(expression: IrTypeOperatorCall): IrExpression {
-        if (isFunInterfaceConversion(expression)) {
-            val argument = expression.argument.transform(this, null) as IrFunctionExpression
+        val functionExpr = expression.findSamFunctionExpr()
+        if (functionExpr != null && expression.typeOperand.isComposableFunInterface()) {
+            val argument = functionExpr.transform(this, null) as IrFunctionExpression
             val superType = expression.typeOperand
             val superClass = superType.classOrNull ?: error("Expected non-null class")
             return FunctionReferenceBuilder(
@@ -81,3 +64,38 @@
         return super.visitTypeOperator(expression)
     }
 }
+
+private fun IrType.isComposableFunInterface(): Boolean =
+    classOrNull
+        ?.functions
+        ?.single { it.owner.modality == Modality.ABSTRACT }
+        ?.owner
+        ?.hasComposableAnnotation() == true
+
+internal fun IrTypeOperatorCall.findSamFunctionExpr(): IrFunctionExpression? {
+    val argument = argument
+    val operator = operator
+    val type = typeOperand
+    val functionClass = type.classOrNull
+
+    val isFunInterfaceConversion = operator == SAM_CONVERSION &&
+        functionClass != null &&
+        functionClass.owner.isFun
+
+    return if (isFunInterfaceConversion) {
+        // if you modify this logic, make sure to update wrapping of type operators
+        // in ComposerLambdaMemoization.kt
+        when {
+            argument is IrFunctionExpression && argument.origin.isLambda -> argument
+            // some expressions are wrapped with additional implicit cast operator
+            // unwrapping that allows to avoid SAM conversion that capture FunctionN and box params.
+            argument is IrTypeOperatorCall && argument.operator == IMPLICIT_CAST -> {
+                val functionExpr = argument.argument
+                functionExpr as? IrFunctionExpression
+            }
+            else -> null
+        }
+    } else {
+        null
+    }
+}
diff --git a/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposerLambdaMemoization.kt b/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposerLambdaMemoization.kt
index 3a8ba5a..166307c 100644
--- a/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposerLambdaMemoization.kt
+++ b/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposerLambdaMemoization.kt
@@ -69,11 +69,14 @@
 import org.jetbrains.kotlin.ir.expressions.IrFunctionExpression
 import org.jetbrains.kotlin.ir.expressions.IrFunctionReference
 import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.expressions.IrTypeOperator
+import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
 import org.jetbrains.kotlin.ir.expressions.IrValueAccessExpression
 import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionReferenceImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetObjectValueImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrInstanceInitializerCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrTypeOperatorCallImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrVarargImpl
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
@@ -509,8 +512,70 @@
         return result
     }
 
+    override fun visitTypeOperator(expression: IrTypeOperatorCall): IrExpression {
+        // SAM conversions are handled by Kotlin compiler
+        // We only need to make sure that remember is handled correctly around type operator
+        if (
+            expression.operator != IrTypeOperator.SAM_CONVERSION ||
+                currentFunctionContext?.canRemember != true
+        ) {
+            return super.visitTypeOperator(expression)
+        }
+
+        // Unwrap function from type operator
+        val originalFunctionExpression =
+            expression.findSamFunctionExpr() ?: return super.visitTypeOperator(expression)
+
+        // Record capture variables for this scope
+        val collector = CaptureCollector()
+        startCollector(collector)
+        // Handle inside of the function expression
+        val result = super.visitFunctionExpression(originalFunctionExpression)
+        stopCollector(collector)
+
+        // If the ancestor converted this then return
+        val newFunctionExpression = result as? IrFunctionExpression ?: return result
+
+        // Construct new type operator call to wrap remember around.
+        val newArgument = when (val argument = expression.argument) {
+            is IrFunctionExpression -> newFunctionExpression
+            is IrTypeOperatorCall -> {
+                require(
+                    argument.operator == IrTypeOperator.IMPLICIT_CAST &&
+                        argument.argument == originalFunctionExpression
+                ) {
+                    "Only implicit cast is supported inside SAM conversion"
+                }
+                IrTypeOperatorCallImpl(
+                    argument.startOffset,
+                    argument.endOffset,
+                    argument.type,
+                    argument.operator,
+                    argument.typeOperand,
+                    newFunctionExpression
+                )
+            }
+            else -> error("Unknown ")
+        }
+
+        val expressionToRemember =
+            IrTypeOperatorCallImpl(
+                expression.startOffset,
+                expression.endOffset,
+                expression.type,
+                IrTypeOperator.SAM_CONVERSION,
+                expression.typeOperand,
+                newArgument
+            )
+        return rememberExpression(
+            currentFunctionContext!!,
+            expressionToRemember,
+            collector.captures.toList()
+        )
+    }
+
     private fun visitNonComposableFunctionExpression(
-        expression: IrFunctionExpression
+        expression: IrFunctionExpression,
     ): IrExpression {
         val functionContext = currentFunctionContext
             ?: return super.visitFunctionExpression(expression)
diff --git a/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/IrSourcePrinter.kt b/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/IrSourcePrinter.kt
index 6103318..04c7f2c 100644
--- a/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/IrSourcePrinter.kt
+++ b/compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/IrSourcePrinter.kt
@@ -616,10 +616,12 @@
             IrTypeOperator.NOT_INSTANCEOF -> {
                 expression.argument.print()
             }
-            IrTypeOperator.CAST, IrTypeOperator.IMPLICIT_CAST, IrTypeOperator.SAFE_CAST -> {
+            IrTypeOperator.CAST, IrTypeOperator.SAFE_CAST, IrTypeOperator.IMPLICIT_CAST -> {
                 expression.argument.print()
             }
             IrTypeOperator.SAM_CONVERSION -> {
+                print(expression.type.renderSrc())
+                print(" ")
                 expression.argument.print()
             }
             IrTypeOperator.IMPLICIT_NOTNULL -> {
diff --git a/compose/runtime/runtime/src/nonEmulatorCommonTest/kotlin/androidx/compose/runtime/CompositionTests.kt b/compose/runtime/runtime/src/nonEmulatorCommonTest/kotlin/androidx/compose/runtime/CompositionTests.kt
index 62a25cc..8f8615b 100644
--- a/compose/runtime/runtime/src/nonEmulatorCommonTest/kotlin/androidx/compose/runtime/CompositionTests.kt
+++ b/compose/runtime/runtime/src/nonEmulatorCommonTest/kotlin/androidx/compose/runtime/CompositionTests.kt
@@ -3895,6 +3895,35 @@
         assertEquals(1, consumer.invokeCount)
     }
 
+    fun interface TestFunInterface {
+        fun compute(value: Int)
+    }
+
+    @Composable fun TestMemoizedFun(compute: TestFunInterface) {
+        val oldCompute = remember { compute }
+        assertEquals(oldCompute, compute)
+    }
+
+    @Test
+    fun funInterface_isMemoized() = compositionTest {
+        var recomposeTrigger by mutableStateOf(0)
+        val capture = 0
+        compose {
+            use(recomposeTrigger)
+            TestMemoizedFun {
+                // no captures
+                use(it)
+            }
+            TestMemoizedFun {
+                // stable captures
+                use(capture)
+            }
+        }
+
+        recomposeTrigger++
+        advance()
+    }
+
     private inline fun CoroutineScope.withGlobalSnapshotManager(block: CoroutineScope.() -> Unit) {
         val channel = Channel<Unit>(Channel.CONFLATED)
         val job = launch {
diff --git a/compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/AndroidLayoutDrawTest.kt b/compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/AndroidLayoutDrawTest.kt
index 3a416f3..93604d6 100644
--- a/compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/AndroidLayoutDrawTest.kt
+++ b/compose/ui/ui/src/androidAndroidTest/kotlin/androidx/compose/ui/AndroidLayoutDrawTest.kt
@@ -2952,6 +2952,7 @@
     @Test
     fun instancesKeepDelegates() {
         var color by mutableStateOf(Color.Red)
+        var size by mutableStateOf(30)
         var m: Measurable? = null
         val layoutCaptureModifier = object : LayoutModifier {
             override fun MeasureScope.measure(
@@ -2960,15 +2961,22 @@
             ): MeasureResult {
                 m = measurable
                 val p = measurable.measure(constraints)
-                drawLatch.countDown()
                 return layout(p.width, p.height) {
                     p.place(0, 0)
                 }
             }
         }
+        val drawCaptureModifier = object : DrawModifier {
+            override fun ContentDrawScope.draw() {
+                drawLatch.countDown()
+            }
+        }
         activityTestRule.runOnUiThread {
             activity.setContent {
-                FixedSize(30, layoutCaptureModifier.background(color)) {}
+                FixedSize(
+                    size = size,
+                    modifier = layoutCaptureModifier.background(color).then(drawCaptureModifier)
+                ) {}
             }
         }
         assertTrue(drawLatch.await(1, TimeUnit.SECONDS))
@@ -2977,6 +2985,7 @@
 
         activityTestRule.runOnUiThread {
             m = null
+            size = 40
             color = Color.Blue
         }