Bound function reference to stable class's method with arguments breaks skipping [302680514]

Fixed

Feature Request

Status Update

No update yet.

Description

el...@reaktor.com

created issue #1

Sep 30, 2023 02:18PM

Jetpack Compose version: 1.6.0-alpha06
Jetpack Compose component(s) used: @Composable
Kotlin version: 1.9.10
Compose compiler version: 1.5.3

Steps to Reproduce or Code Sample to Reproduce:

class MainActivity2 : ComponentActivity() {
    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)

        val cl = StableClass()
        setContent {
            Root(cl)
        }
    }
}

@Composable
fun Root(stableClass: StableClass) {
    Log.w("TEST", "Recomposing root")
    var internalCounter by remember { mutableIntStateOf(0) }

    Button(onClick = { internalCounter += 1 }) {
        Text("Increment ($internalCounter)")
    }

    internalCounter // recompose this scope when changed
    Skippable(stableClass::method)
}

var previous: Any? = null

class StableClass {
    fun method(arg: Int) {}
}

@Composable
fun Skippable(functionRef: (Int) -> Unit) {
    Log.w("TEST", "Recomposing while equals=${functionRef == previous}")
    previous = functionRef
}

Results in log output:

 W  Recomposing root
 W  Recomposing while equals=false
 W  Recomposing root
 W  Recomposing while equals=true
 W  Recomposing root
 W  Recomposing while equals=true
 W  Recomposing root
 W  Recomposing while equals=true
[...]

The Android Studio debugger claims that the parameter has changed even though equals returns true.

Composable fun Skippable(): Arguments: Different: ["functionRef"]

If the method signature is changed to not have arguments, skipping occurs normally:
fun method() {}
If a new lambda is created instead of passing the function reference, skipping occurs normally:
Skippable { stableClass.method(it) }

Comments

as...@google.com <as...@google.com> #2Oct 3, 2023 12:59PM

Status: Won't Fix (Intended Behavior)

This is expected behavior, we are checking for referential equality on function references, as their equals implementation doesn't take the receiver (stableClass) into account.

js...@google.com <js...@google.com> #3Oct 3, 2023 03:11PM

I don't feel strongly either way, but if we cared we could probably rewrite function references to behave more like lambdas which preserve stability. It does seem like maybe a valid request/optimization even if a low-priority one?

el...@reaktor.com <el...@reaktor.com> #4Oct 4, 2023 06:55AM

Thanks for the explanation. I didn't realize that using a bound function reference actually creates a new reference on every access, since I assumed .equals() would be using referential equality!

However, there still seems to be a bug.

To my understanding the compiler should generate a remember expression around parameter-less function reference, but it doesn't seem to generate one for function references with more than 0 arguments. I would imagine this breaking it:

        if (expression.valueArgumentsCount != 0) {
            // If this syntax is as a curry syntax in the future, don't memoize.
            // The syntax <expr>::<method>(<params>) and ::<function>(<params>) is reserved for
            // future use. This ensures we don't try to memoize this syntax without knowing
            // its meaning.

IrFunctionReference.valueArgumentsCount seems to currently refer to the parameters of the method being referenced, not the parameters in the calling syntax. I might be misunderstanding as I'm not very familiar with the internals, but this still seems premature at this moment.

For others reading this issue: a workaround is to replace direct bound function references:

- klass::method
+ remember(klass) { klass::method }
// or
+ { arg -> klass.method(arg) }

Can this issue be reopened or should I file a new issue with a more specific description?

as...@google.com <as...@google.com> #5Oct 4, 2023 08:09AM

Assigned to as...@google.com.

I see, this feature request sounds legit to me I think we should be able to remember reference based on receiver types.

el...@reaktor.com <el...@reaktor.com> #6Oct 4, 2023 09:18AM

Regarding #3, #5; based on the compiler source it seems function reference memoization is already implemented in the compiler, it just doesn't work correctly (see #4).

as...@google.com <as...@google.com> #7Oct 4, 2023 09:26AM

It is not quite the case, the line of code you linked prevents memoization in cases if someone does cls::invoke(0)

We'll need to support function references separately, as IR shape for them is different

ap...@google.com <ap...@google.com> #8Nov 13, 2023 05:30PM

Marked as fixed.

Project: platform/frameworks/support
Branch: androidx-main

commit c05d2bbe2b44dfead66315174a4d82951fd56b46
Author: Elias Benkhodja <elias.benkhodja@reaktor.com>
Date: Wed Oct 04 20:57:13 2023

Memoize function references with arguments

`ComposerLambdaMemoization` tried to prematurely prevent potential
syntax addition not breaking function reference optimization by checking
for amount of value parameters.

This inadvertently broke memoization of function references to methods
with more than zero arguments, since `IrFunctionReferenceImpl` currently
reflects the amount of value arguments in the referenced function.

Add a check to disable memoization for function references with context
receivers. Kotlin doesn't support them yet, and behavior might change in
the future.

This change also enables memoizing function references without checking
for stability if strong skipping is enabled.

Test: Added new tests to `LambdaMemoizationTransformTests`:
* `testNonComposableFunctionReferenceWithNoArgumentsMemoization`
* `testNonComposableFunctionReferenceWithArgumentsMemoization`
* `testNonComposableFunctionReferenceWithStableContextReceiverNotMemoized`
* `testUnstableReceiverFunctionReferenceNotMemoized`
* `testUnstableExtensionReceiverFunctionReferenceNotMemoized``

Added new tests to `StrongSkippingModeTransformTests`:
* `testUnstableReceiverFunctionReferenceMemoized`
* `testUnstableExtensionReceiverFunctionReferenceMemoized`

Fixes:

b/302680514

Relnote: Memoize stable function references with arguments
Change-Id: I4d7bf3061e60eaa6fd3af60662d6a7b02c1a83b2

M compose/compiler/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/LambdaMemoizationTransformTests.kt
M compose/compiler/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/StrongSkippingModeTransformTests.kt
A compose/compiler/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testUnstableExtensionReceiverFunctionReferenceNotMemoized[useFir = false].txt
A compose/compiler/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testUnstableExtensionReceiverFunctionReferenceNotMemoized[useFir = true].txt
A compose/compiler/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testUnstableReceiverFunctionReferenceNotMemoized[useFir = false].txt
A compose/compiler/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.LambdaMemoizationTransformTests/testUnstableReceiverFunctionReferenceNotMemoized[useFir = true].txt
A compose/compiler/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.StrongSkippingModeTransformTests/testUnstableExtensionReceiverFunctionReferenceMemoized[useFir = false].txt
A compose/compiler/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.StrongSkippingModeTransformTests/testUnstableExtensionReceiverFunctionReferenceMemoized[useFir = true].txt
A compose/compiler/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.StrongSkippingModeTransformTests/testUnstableReceiverFunctionReferenceMemoized[useFir = false].txt
A compose/compiler/compiler-hosted/integration-tests/src/test/resources/golden/androidx.compose.compiler.plugins.kotlin.StrongSkippingModeTransformTests/testUnstableReceiverFunctionReferenceMemoized[useFir = true].txt
M compose/compiler/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/ComposerLambdaMemoization.kt

https://android-review.googlesource.com/2822157

pr...@google.com <pr...@google.com> #9Nov 29, 2023 09:39PM

The following release(s) address this bug.It is possible this bug has only been partially addressed:

androidx.compose.compiler:compiler-hosted:1.5.5