Context receivers – new extension functions

I think I won’t reveal the big secret that Ozon has developed a certain number of mobile applications: for buyers, for sellers, a bank, etc. Each of them requires authorization. For this purpose, our Ozon ID team exists with its own SDK. I am part of the Ozon ID team – an Android developer with an exorbitant love for the syntactic sugar of Kotlin.

Introduction

Let's talk today about context receivers – a Kotlin feature that I learned about a long time ago, but was able to find application only a couple of months ago. I’ll tell you what context receivers are, where they can be used, and, of course, about “success” – minus 60% of the self-written DI in the Ozon ID SDK. But first things first.

Extension functions

The Ozon ID SDK code contains the following extension for ComponentActvity.

inline fun <T> ComponentActvity.collectWhenStarted(
    data: Flow<T>,
    crossinline collector: (T) -> Unit
) {
    lifecycleScope.launchWhenStarted { 
        data.collect {
            collector.invoke(it)
        } 
    }
}

Extension implemented collectWhenStarted solely for the convenience of collecting data from a variety of Flow from ViewModel inside Activity. Below is an example of using this extension.

class AuthFlowActivity : AppCompatActivity() {
    override fun onCreate(savedInstanceState: Bundle?) {
        ...
        collectWhenStarted(viewModel.backStack) { onBackStackChanged(it) }
        collectWhenStarted(viewModel.navigationEvents) { onNavigationEvent(it) }
        ...
    }
}

Comfortable? In general, yes. The call is undoubtedly shorter than without using the extension. Perfect? Not at all. Personally, I'd like to see a challenge collectWhenStarted something like this:

class AuthFlowActivity : AppCompatActivity() {
    override fun onCreate(savedInstanceState: Bundle?) {
        ...
        viewModel.backStack.collectWhenStarted { onBackStackChanged(it) }
        viewModel.navigationEvents.collectWhenStarted { onNavigationEvent(it) }
        ...
    }
}

Call collect* at Flow more intuitive than collect* at Activityis not it?

In order to implement an improved extension option collectWhenStarted, technically we need the mechanism of extension methods to be able to accept 2 receiver parameters. Unfortunately, JetBrains has not yet implemented this feature in the Kotlin language… Or has it?

Implementation via scope

In fact, there is such an approach as scope. With this approach, the extension method is implemented in a new class. In case of Activity The scope solution can be applied as follows.

interface LifecycleOwnerScope : LifecycleOwner {
    fun <T> Flow<T>.collectWhenStarted(collector: (T) -> Unit) {
        lifecycleScope.launchWhenStarted { collect { collector.invoke(it) } }
    }
}

We implement in Activity interface LifecycleOwnerScopeand you're done Flow extension appears collectWhenStarted. In this case, the implementation of the interface LifecycleOwnerScope inside Activity is not required due to the “default implementation”.

class AuthFlowActivity : AppCompatActivity(), LifecycleOwnerScope {
                                            // ^ Добавили scope ^
    override fun onCreate(savedInstanceState: Bundle?) {
        ...             // collectWhenStarted берётся из LifecycleOwnerScope
        viewModel.backStack.collectWhenStarted { onBackStackChanged(it) }
        viewModel.navigationEvents.collectWhenStarted { onNavigationEvent(it) }
        ...
    }
}

The approach is quite viable. But, as for me, it is not without its shortcomings. For example, to expand collectWhenStarted it works, you need to inherit from it yourself LifecycleOwnerScope. This is not as convenient as global extensions, which the IDE helpfully suggests when you enter the name.

Let's move on to context receivers.

Context receivers come to the rescue

Context receivers are a concept, a feature of the Kotlin language. Context receivers were added to the language as a tool to overcome the limitations of extension functions. Technically, context receivers, like extension functions, are compiled into a static method with an additional this-parameter. That is, context receivers are just another syntactic sugar of Kotlin, but, of course, cooler and “sweeter” than extension.

Context receivers appeared in Kotlin 1.6.20. Along with context receivers, a new keyword was added to the language context. The feature in Kotlin 1.9.22 is still experimental, so if you try to use it right away, the IDE will display the following message.

The feature “context receivers” is experimental and should be enabled explicitly

In order for context receivers to work, you need to add the file build.gradle module in which it will be used context.

tasks.withType<org.jetbrains.kotlin.gradle.tasks.KotlinCompile>().configureEach {
    kotlinOptions {
        freeCompilerArgs = freeCompilerArgs + "-Xcontext-receivers"
    }
}

Now you can use context in the code. Let's get started. We will modify the extension step by step.

Step 1: Transfer receiver parameter ComponentActvity into arguments context.

context (ComponentActvity)
inline fun <T> /*ComponentActvity.*/collectWhenStarted(
    data: Flow<T>,
    crossinline collector: (T) -> Unit
) {
    lifecycleScope.launchWhenStarted { 
        data.collect {
            collector.invoke(it)
        } 
    }
}

public static final void collectWhenStarted(
    @NotNull ComponentActivity $this$collectWhenStarted, 
    @NotNull final Flow data, 
    @NotNull final Function1 collector
) 

The code will be compiled and executed, as in the previous implementation through the extension. But there is one nuance that is worth dwelling on. The fact is that context receivers are not an extension (your captain Obviousness). Functions with context receivers cannot be called as if they were a class method. An example in the snippet below.

// Реализация через extension
activity.collectWhenStarted() // компилятор позволяет вызвать функцию-расширение `collectWhenStarted`, как будто это метод класса

// Реализация через context receivers
activity.collectWhenStarted() // Ошибка: Unresolved reference: collectWhenStarted
with(activity) { // apply, run тоже подойдут
    collectWhenStarted() // Функцию с context можно вызвать только внутри "контекста" 
}

Step 2: Move argument data: Flow in receiver function

context (ComponentActvity)
inline fun <T> Flow<T>.collectWhenStarted(
    /*data: Flow<T>,*/
    crossinline collector: (T) -> Unit
) {
    lifecycleScope.launchWhenStarted { 
        /*data.*/collect {
            collector.invoke(it)
        } 
    }
}

Voila! Ready. Now Flow inside (in other words, “in context”) ComponentActvity the method will appear collectWhenStarted.

class AuthFlowActivity : AppCompatActivity() {
    override fun onCreate(savedInstanceState: Bundle?) {
        ...            // powered by context receivers
        viewModel.backStack.collectWhenStarted { onBackStackChanged(it) }
        viewModel.navigationEvents.collectWhenStarted { onNavigationEvent(it) }
        ...
    }
}

We need to go deeper

I'll reveal a little more information regarding the keyword context:

1. context can take more than one argument. For example, context (classA, classB).

2. context can be written above classes.

Just like annotation @Component(dependencies = [...]) in Dagger, isn't it?

I’ll start analyzing these points from the last one – context above the class. As a basis for examples, I will take DI from the Ozon ID SDK.

internal class RootModule(
    val application: Application
)

context(RootModule)
internal class ChildModule {

    val dependency by lazy {
        Dependency(
            /*this@RootModule.*/application,
            ...
        )
    }
}

As you can see from the example above, inside ChildModule you can contact applicationas if it were a property ChildModule. In this case, the property can also be accessed via this — this@RootModule.application. Useful in case of a name conflict between context receivers or to clarify the question “where did this dependency come from?”

Let's continue the dive. Let's look at how to create dependencies with context receivers.

context(RootModule)
internal class SubChildModule

context(RootModule)
internal class ChildModule {
    ...

    val subChildModule by lazy {
        // with не нужен, уже в нужном контексте
        SubChildModule()
    }
}

val childModule = with(rootModule) {
    // Нужен with для создания
    ChildModule()
}

In order to create a class object with context receivers, the constructor must be called in the required context. The context can be set as follows:

1. create an instance inside the class required in the context receiver;

2. create an instance inside an independent class whose context receivers contain the required class. For example, how the module is created SubChildModule inside ChildModule;

3. create inside a scope function with a lambda extension as an argument (with, apply, run). Custom functions with a similar parameter are also suitable.

Now let's look at creating objects with several context receivers.

context(RootModule, RepositoryModule, NetworkModule, CookieModule)
internal class MultiContextModule

val multiContextModule = with(rootModule) {
    with(repositoryModule) {
        with(networkModule) {
            with(cookieModule) {
                MultiContextModule()
            }
        }
    }
}

As you can see, the code for creating an instance of a class with several context receivers may not be as elegant as we would like. But this can be fixed because context can be applied to lambda arguments. Let's take the code as a basis with from the Kotlin standard library and enrich it context.

// Пример из stdlib
@kotlin.internal.InlineOnly
public inline fun <T, R> with(receiver: T, block: T.() -> R): R {
    contract {
        callsInPlace(block, InvocationKind.EXACTLY_ONCE)
    }
    return receiver.block()
}

// with context receivers
@OptIn(ExperimentalContracts::class)
@Suppress("SUBTYPING_BETWEEN_CONTEXT_RECEIVERS")
internal inline fun <T1, T2, R> with(c1: T1, c2: T2, block: context(T1, T2) () -> R): R {
    contract {                                           // ^^^^^^^
        callsInPlace(block, InvocationKind.EXACTLY_ONCE) // см. сюда
    }
    return block(c1, c2)
}

In the example above to the lambda type block keyword added context. Options context – generic types. Similar with you can write more for the required number of context receivers. In turn, this allows us to significantly reduce code shift to the right when creating MultiContextModule.

context(RootModule, RepositoryModule, NetworkModule, CookieModule)
internal class MultiContextModule

val multiContextModule = with(
    rootModule,
    repositoryModule,
    networkModule,
    cookieModule
) {
    MultiContextModule()
}

Results

Let's sum up what we learned today. The file will help us with this KEEP by context receivers.

• Context receivers are a mechanism designed to expand the capabilities of extension functions. Moreover, as we saw in the example with Flownamely to expand, not replace.

• Context receivers can be written both over functions and properties, and over classes.

• Context receivers allow you to use more than one receiver argument.

• Context receivers are unclear where and when to use. At the very least, I have not yet formed a clear opinion so that I can unambiguously say “here is an extension, here is to pass it as an argument, and here is the context.” In the video at the end of the article you can find thoughts on this topic. For now, in this matter I will stick exclusively to the technical point.

Conclusion

I learned about context receivers almost 2 years ago from video from Jetbrains, could not come up with any useful use for them and put them on the far shelf of knowledge about Kotlin. However, a couple of months ago I had a chance to watch report from Droidcon, which helped open my eyes to the full power of this mechanism. And then it started: refactoring DI into the Ozon ID SDK, a report within the mobile team, and as a result this article. I hope that I was able to convey to readers the power of context receivers and encourage them to further search for the applicability of this feature.