LifecycleEffect.kt

import androidx.compose.runtime.Composable
import androidx.compose.runtime.InternalComposeApi
import androidx.compose.runtime.NonRestartableComposable
import androidx.compose.runtime.RememberObserver
import androidx.compose.runtime.currentComposer
import androidx.compose.runtime.remember
import kotlin.coroutines.CoroutineContext
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Job
import kotlinx.coroutines.cancel
import kotlinx.coroutines.launch

internal interface LifecycleEffectHandle : CoroutineScope {
  fun whenDispose(dispose: () -> Unit)
}

private class LifecycleEffectScope(
  parentCoroutineContext: CoroutineContext,
  private val lifecycleTask: suspend LifecycleEffectHandle.() -> Unit,
) : RememberObserver {
  private val scope = CoroutineScope(parentCoroutineContext)
  private var job: Job? = null
  private var onDispose: (() -> Unit)? = null

  private val handle = object : LifecycleEffectHandle, CoroutineScope by scope {
    override fun whenDispose(dispose: () -> Unit) {
      onDispose = dispose
    }
  }

  override fun onRemembered() {
    // This should never happen but is left here for safety
    job?.cancel("Old job was still running!")
    job = scope.launch { handle.lifecycleTask() }
  }

  override fun onForgotten() {
    job?.cancel(LeftCompositionCancellationException())
    onDispose?.invoke()

    job = null
    onDispose = null
  }

  override fun onAbandoned() {
    onForgotten()
  }
}

private class LeftCompositionCancellationException : CancellationException("The coroutine scope left the composition") {
  override fun fillInStackTrace(): Throwable {
    // Avoid null.clone() on Android <= 6.0 when accessing stackTrace
    stackTrace = emptyArray()
    return this
  }
}

@[Composable NonRestartableComposable]
@OptIn(InternalComposeApi::class)
internal fun LifecycleEffect(
  key1: Any?,
  block: suspend LifecycleEffectHandle.() -> Unit,
) {
  val applyContext = currentComposer.applyCoroutineContext
  remember(key1) { LifecycleEffectScope(applyContext, block) }
}

@[Composable NonRestartableComposable]
@OptIn(InternalComposeApi::class)
internal fun LifecycleEffect(
  key1: Any?,
  key2: Any?,
  block: suspend LifecycleEffectHandle.() -> Unit,
) {
  val applyContext = currentComposer.applyCoroutineContext
  remember(key1, key2) { LifecycleEffectScope(applyContext, block) }
}

@[Composable NonRestartableComposable]
@OptIn(InternalComposeApi::class)
internal fun LifecycleEffect(
  key1: Any?,
  key2: Any?,
  key3: Any?,
  block: suspend LifecycleEffectHandle.() -> Unit,
) {
  val applyContext = currentComposer.applyCoroutineContext
  remember(key1, key2, key3) { LifecycleEffectScope(applyContext, block) }
}

@[Composable NonRestartableComposable]
@OptIn(InternalComposeApi::class)
internal fun LifecycleEffect(
  vararg keys: Any?,
  block: suspend LifecycleEffectHandle.() -> Unit,
) {
  val applyContext = currentComposer.applyCoroutineContext
  remember(*keys) { LifecycleEffectScope(applyContext, block) }
}