Uses Arrow atomic

SafeMultishotContinuation.kt

/*
 * Copyright 2010-2018 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */
import CoroutineSingletons.UNDECIDED
import arrow.atomic.Atomic
import arrow.atomic.loop
import arrow.atomic.value
import kotlin.coroutines.Continuation
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.intrinsics.COROUTINE_SUSPENDED

@PublishedApi
internal class SafeMultishotContinuation<in T>
internal constructor(
  private val delegate: Continuation<T>,
  initialResult: Any?
) : Continuation<T> {
  @PublishedApi
  internal constructor(delegate: Continuation<T>) : this(delegate, UNDECIDED)

  override val context: CoroutineContext
    get() = delegate.context

  private var atomicResult: Atomic<Any?> = Atomic(initialResult)

  override fun resumeWith(result: Result<T>) {
    atomicResult.loop {
      when {
        it === UNDECIDED -> if (atomicResult.compareAndSet(it, result)) return
        else -> {
          delegate.resumeWith(result)
          return
        }
      }
    }
  }

  @PublishedApi
  internal fun getOrThrow(): Any? {
    var result = this.atomicResult.value // atomic read
    if (result === UNDECIDED) {
      if (atomicResult.compareAndSet(UNDECIDED, COROUTINE_SUSPENDED)) return COROUTINE_SUSPENDED
      result = this.atomicResult // reread volatile var
    }
    return when {
      result === COROUTINE_SUSPENDED -> COROUTINE_SUSPENDED // already called continuation, indicate COROUTINE_SUSPENDED upstream
      result is Result<*> -> result.getOrThrow()
      else -> error("Unexpected state with $result")
    }
  }

  override fun toString(): String =
    "SafeMultishotContinuation for $delegate"
}

internal enum class CoroutineSingletons { UNDECIDED }