I'm seeping super weird stuff again: I'm passing a...

# compose

[vide]

I'm seeping super weird stuff again: I'm passing a function reference around to some composable functions as an argument, and:
• debugger always shows it has changed when recomposing
• comparing to previous value with equals() returns true
• compose metrics claim the composable is stable. How can this be?

[Jussi Malinen]

To give more context (I am working on the same code), in the debugger when looking at the recomposition state on debugger it shows the value as different as that starts with a reference to the caller. {AudioLandingPageKt$AudioLandingPageWrapper$4@36054} on the first request and {AudioLandingPageKt$AudioLandingPageWrapper$4@36073} on second. When comparing the values in the function body, the == operation tells they are the same.

[dewildte]

I would consult a medical practitioner if I was you.

[vide]

Returning to this with a repro: a function reference to a stable class with a private property to an unstable class somehow disables skipping recompositions. Any idea why? 🧵

[vide]

class UnstableClass(var publicVar: Int)
class SomeClass(private val unstableClass: UnstableClass? = null) {
    fun method() {}
}

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

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

@Composable
fun Root(someClass: SomeClass) {
    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(someClass::method)
}

var previous: Any? = null

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

[vide]

This will never skip Skippable, but adding @Stable to either UnstableClass or SomeClass will allow skipping Skippable again. (Or removing private val unstableClass: UnstableClass? = null from SomeClass)

[vide]

Metrics indicate that Skippable is skippable and SomeClass is stable

[vide]

Equals is always true in the log line

[dewildte]

As far as i understand it, if you say a thing is stable or immutable then it's entire dependency stack should be either stable or immutable all the way down.

[dewildte]

The @Stable and @Immutable annotations are telling the compiler that.

[vide]

Yes. The compiler reports SomeClass to be stable, but the method reference still seems to act like it would be unstable. That is the issue

[vide]

Private properties should not affect the stability of the class afaik, from @Stable:

1) The result of [equals] will always return the same result for the same two instances.

2) When a public property of the type changes, composition will be notified.

3) All public property types are stable.

[dewildte]

Maybe. Function references are inherently unstable?

[dewildte]

Weird, maybe file a bug.

[vide]

Would need someone with compiler experience from the team to answer that 😅 I think I originally got the impression from this blog post when first learning compose, but it's not a guarantee that it's correct. It also might have changed in some compose/compiler verison...

https://multithreaded.stitchfix.com/blog/2022/08/05/jetpack-compose-recomposition/

[vide]

hmm, I might have been looking at compiler metrics for the wrong build variant

[vide]

restartable scheme("[androidx.compose.ui.UiComposable]") fun Root(
  unstable someClass: SomeClass
)
restartable skippable fun Skippable(
  stable functionRef: Function0<Unit>
)

and

unstable class UnstableClass {
  stable var publicVar: Int
  <runtime stability> = Unstable
}
unstable class SomeClass {
  unstable val unstableClass: UnstableClass?
  <runtime stability> = Unstable
}
stable class MainActivity2 {
  <runtime stability> = Stable
}

seem to be the correct ones, so SomeClass still surprisingly seems to be unstable, but the function reference is inferred to be stable

[vide]

I guess the function reference could reference the private var that might have changed even though the function reference is the same

[vide]

so! I tried a little more to match the production code closer and got a new result: a function reference to this method will never skip recomposition:

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

and this will skip recomposition:

class StableClass {
    fun method() {}
}

[dewildte]

Fascinating

[vide]

indeed! created an issue for this now: https://issuetracker.google.com/issues/302680514

[vide]

@Zach Klippenstein (he/him) [MOD] sorry to ping you directly, but could you check this issue and confirm if this is a real or am I just missing something on how skipping is supposed to behave? This feels like it could incur several perf degradation because large trees of composables are unable to skip: https://issuetracker.google.com/issues/302680514

[Zach Klippenstein (he/him) [MOD]]

Not sure what the expected behavior is, paging @Andrew Bailey

[vide]

Someone helpfully responded that function references are compared referentially (===) instead of .equals(), which explains the behavior. Normally the compiler generates remember expressions around function references, but this is explicitly disabled for more-than-0 parameters: https://cs.android.com/androidx/platform/frameworks/support/+/androidx-main:compose/[…]piler/plugins/kotlin/lower/ComposerLambdaMemoization.kt;l=449. I think this is still a bug, let's see if the issue can be reopened or if I should file a new bug

[Andrew Bailey]

I think Andrei's answer on that bug sums up what you're seeing fairly well. Regarding the original question, we use reference equality (===) to check whether a lambda has changed or not. And we infer that all lambdas are stable. We also implicitly memoize lambdas, so you'll get a new instance in composition if the dispatch receiver changes or if you get a new capture scope that requires us to make a new lambda for you. So that would explain how you're seeing an .equals() lambda that changes during recomposition and is also stable.

[vide]

Thanks Andrew. To me it however looks like memoization is already implemented for function references too, not only lambdas. It just doesn't work as intended(?). @shikasd could you double check the newest comment? Thank you!

[shikasd]

It is likely that Kotlin optimizes function references somehow as well, making them single instance for example if you don't have any parameters
You can check the generated bytecode, if you are interested

[shikasd]

We might unintentionally memoize function references, but it was never intended to work iirc Nvm, I am mistaken

[vide]

I'm somewhat interested in learning enough about the internals to be able to contribute myself to the compiler as well as I feel like I keep on finding these edge cases a lot. Is it open for contribution from outside?

[vide]

Thanks again for the insight Andrei, I can check the actual generated bytecode.

[shikasd]

Yeah, we are definitely open for contribution, if you want to implement this

[shikasd]

The easiest way to test if your theory about valueArgumentCount is correct is to debug our compiler test in androidx (e.g. LambdaMemoizationTransformTests), might be a good point to start

[vide]

It was definitely an experience to clone and build androidx blob sweat smile I created a draft PR as it seemed to be the issue as I suspected: https://github.com/androidx/androidx/pull/615. I'll work on the CLA tomorrow with our legal. I have never touched the Kotlin compiler before but I think this is where the value parameters are coming from:

// FunctionReferenceLowering.kt
...
    private inner class FunctionReferenceBuilder(...)
...
        private fun JvmIrBuilder.generateSignature(...)
...
                    IrFunctionReferenceImpl(
                        UNDEFINED_OFFSET, UNDEFINED_OFFSET, irFunctionReference.type, target,
                        irFunctionReference.typeArgumentsCount, target.owner.valueParameters.size,
                        irFunctionReference.reflectionTarget, null
                    )

[shikasd]

Interesting, maybe you should also check that all valueArguments are null then?
I am not sure why it has any, but it kinda makes sense. I still want to guard against remembering unexpected shapes, e.g. in cases where context receivers will be involved (they are represented as regular args)

[vide]

It seems that getValueArgument always returns null for any index regardless of presence of context receivers. Should I just check for expression.symbol.owner.contextReceiverParametersCount == 0?
We could also potentially check for stability of the context receivers and use those as additional keys to remember? I'm not sure how I would implement that though.

[vide]

Protyping checking for contextReceiverParametersCount:
https://github.com/androidx/androidx/commit/d44308ab7056b8b4dcb268c88fe2faa51531a9c3

[vide]

Actually, it seems like function references to functions with context receivers are kind of unsupported in Kotlin at the moment. The context receiver must be passed as a regular parameter when calling the function reference. With the current behaviour it wouldn't then matter if the context receiver types are stable or unstable? So the reference could be memoized regardless.
• https://pl.kotl.in/VUlMl7ztT
• https://github.com/Kotlin/KEEP/blob/master/proposals/context-receivers.md#callable-references-to-contextual-functions

[shikasd]

We want to treat them the same way as current receivers, that's why I am referencing it
In the current state of things, if Kotlin was ever to allow function references with context receivers, we'll get a stale reference (correctness bug), while not memoizing it is performance bug, much less critical

[vide]

Alright. Should I just disable memoization completely for now if the function referenced has context receivers?

[shikasd]

That + I'd check if all of value arguments are null, so we are only processing known IR shapes

[vide]

Hi, our legal got the CLA approved while I was on vacation and I'm back on this. I updated the patch and uploaded the changes directly to gerrit: https://android-review.googlesource.com/c/platform/frameworks/support/+/2822157. I added some logic to check if all value arguments are null, is this what you meant?

[vide]

This also doesn't now take into account the new strong skipping mode, but no changes to function reference handling were yet made as part of that anyways as far as I can tell. Maybe it could also enable memoizing function references without caring for stability here?

[shikasd]

Thanks, I'll take a look
Re: strong skipping mode, I am currently changing that behavior to work with another feature, but we should definitely stop caring about the stability of captures if strong skipping is enabled.
It is also interesting to see if we can make it work with @DontMemoize annotation

[vide]

I saw you voted +2 on Gerrit, thanks for doing the review. What is the next step for me? It seems like Review-Enforcement is still not satisfied, should I add a second reviewer?

[shikasd]

Yeah, we need two reviewers for external contributions

[shikasd]

I'll ask someone from the team to review later today

[vide]

Thanks!