Compose KMP For Result

RememberNavResult.kt

import androidx.compose.runtime.Composable
import androidx.compose.runtime.CompositionLocalProvider
import androidx.compose.runtime.Stable
import androidx.compose.runtime.State
import androidx.compose.runtime.compositionLocalOf
import androidx.compose.runtime.derivedStateOf
import androidx.compose.runtime.mutableStateMapOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.saveable.Saver
import androidx.compose.runtime.saveable.SaverScope
import androidx.compose.runtime.saveable.rememberSaveable
import androidx.compose.runtime.snapshots.SnapshotStateMap
import cafe.adriel.voyager.core.screen.Screen
import core.platform.PlatformParcelable
import core.platform.PlatformParcelize
import core.platform.PlatformRawValue
import core.kotlin.UUID
import core.platform.PlatformIgnoredOnParcel
import kotlinx.serialization.Serializable
import kotlin.native.HiddenFromObjC

private val LocalNavResultStateHolder = compositionLocalOf<SnapshotStateMap<String, NavResult<*>>> {
    throw AssertionError()
}

private object StateMapSaver :
    Saver<SnapshotStateMap<String, NavResult<*>>, List<StateMapSaver.KeyValue>> {

    @PlatformParcelize
    class KeyValue(val key: String, val value: NavResult<*>) : PlatformParcelable

    override fun restore(
        value: List<KeyValue>
    ): SnapshotStateMap<String, NavResult<*>>? {
        return mutableStateMapOf(*value.map { it.key to it.value }.toTypedArray())
    }

    override fun SaverScope.save(
        value: SnapshotStateMap<String, NavResult<*>>
    ): List<KeyValue> {
        return value.entries.map { KeyValue(it.key, it.value) }
    }
}


@Composable
internal fun ProvideNavResultStateHolder(
    content: @Composable () -> Unit
) {
    val stateHolder = rememberSaveable(
        saver = StateMapSaver
    ) {
        mutableStateMapOf()
    }

    CompositionLocalProvider(
        LocalNavResultStateHolder provides stateHolder,
        content = content
    )
}

@HiddenFromObjC
internal sealed class NavResult<out T> : PlatformParcelable {

    @HiddenFromObjC
    @PlatformParcelize
    data class Value<T>(val value: @PlatformRawValue T) : NavResult<T>()

    @HiddenFromObjC
    @PlatformParcelize
    data object Empty : NavResult<Nothing>()
}

@Stable
@Serializable
@PlatformParcelize
internal class NavResultRequest<T>(
    val key: String
) : PlatformParcelable

@Stable
internal class NavResultRequestContract<T>(
    private val key: String,
    private val stateHolder: SnapshotStateMap<String, NavResult<*>>,
) {
    val request: NavResultRequest<T> = NavResultRequest(key)

    val response: State<NavResult<T>> = derivedStateOf {
        val currentValue = stateHolder[key]
        if (currentValue != null && currentValue is NavResult.Value<*>) {
            @Suppress("UNCHECKED_CAST")
            currentValue as NavResult<T>
        } else {
            NavResult.Empty
        }
    }

    fun consume() {
        stateHolder.remove(key)
    }
}

@Stable
internal class NavResultResponseContract<T>(
    private val key: String,
    private val stateHolder: SnapshotStateMap<String, NavResult<*>>,
) {
    fun send(value: T) {
        stateHolder[key] = NavResult.Value(value)
    }
}

@Composable
internal fun <T> rememberNavResultContract(
    request: NavResultRequest<T>
): NavResultResponseContract<T> {
    val stateHolder = LocalNavResultStateHolder.current
    return remember(stateHolder, request.key) {
        NavResultResponseContract(
            key = request.key,
            stateHolder = stateHolder
        )
    }
}


@Composable
internal fun <T> rememberNavResultContract(
    key: String
): NavResultRequestContract<T> {
    val stateHolder = LocalNavResultStateHolder.current
    val contractKey = rememberSaveable { "$key+${UUID.randomUUID()}" }
    return remember(stateHolder, key) {
        NavResultRequestContract(
            key = contractKey,
            stateHolder = stateHolder
        )
    }
}

// Root
ProvideNavResultStateHolder() {
   Navigator(Screen1())
}


// Screen1

class Screen1: Screen {
    
    @Composable 
    fun Content() {
        val contract = rememberNavResultContract<Int>(key)
        val response = contract.response.value
        val isCurrentScreen = navigator.isCurrentScreen(key)

        LaunchedEffect(response, isCurrentScreen) {
          if(!isCurrentScreen) return@LaunchedEffect
          if(response !is NavResult.Value) return@LaunchedEffect
          // Do something
          print(response.value)
          contract.consume()
        }

        Button(
          onClick = { 
              navigator.push(Screen2(contract.request))
          }
        ) {
           Text("Push")
        }
    }
}


// Screen2

class Screen1(val request: NavResultRequest<Int>) : Screen {
    @Composable 
    fun Content() {
        val contract = rememberNavResultContract(request)

        Button(
          onClick = { 
              contract.send(0)
              navigator.pop()
          }
        ) {
              Text("Pop")
        }
    }
}