Screeps restorable coroutines

RestorableContext.kt

package util.coroutines.restorable

import screeps.api.MutableRecord
import screeps.api.get
import screeps.api.keys
import screeps.api.set
import screeps.utils.unsafe.delete
import kotlin.coroutines.Continuation
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.intrinsics.COROUTINE_SUSPENDED
import kotlin.coroutines.intrinsics.suspendCoroutineUninterceptedOrReturn

@Suppress("NOTHING_TO_INLINE")
inline fun Any.asRecord() = this.unsafeCast<MutableRecord<String, Any?>>()

class RestorableContext(val restoreId: String, val memory: MutableRecord<String, out Any?>) : CoroutineContext.Element {

    object Key :
        CoroutineContext.Key<RestorableContext>

    override val key: CoroutineContext.Key<RestorableContext> =
        Key
}

suspend fun saveAndSuspend(): Unit = suspendCoroutineUninterceptedOrReturn { originalContinuation ->
    saveCoroutine(originalContinuation)
    // TODO we don't *have* to suspend, we could just continue instead so we have a checkpoint for globalreset
    println("suspending")
    COROUTINE_SUSPENDED
}

/**
 * @param evalFn the eval function for your module. Just insert `::eval`
 */
suspend fun restoreIfNecessary(evalFn: (String) -> Any?) = suspendCoroutineUninterceptedOrReturn<Unit> {
    restoreCoroutineIfNecessary(it, evalFn)
}

private fun saveCoroutine(continuation: Continuation<Unit>) {
    val restorableContext = continuation.context[RestorableContext.Key]
    requireNotNull(restorableContext) { "no restorable coroutine" }

    val plain = js("{}")
    plain.state_0 = continuation.asDynamic().state_0
    plain.original_name_1 = continuation.asDynamic().constructor.name
    plain.exceptionState_0 = continuation.asDynamic().exceptionState_0
//  TODO could use something like this for logging
//  plain.createdSavePointAt = Game.time

    continuation.asRecord().keys.filter { it.startsWith("local") && !it.contains("this") }.forEach {
        val value = continuation.asRecord()[it].asDynamic()
        plain[it] = value
        if (js("typeof value") == "object" && value.constructor.name != "Object") {
            plain["$it\$prototype\$name"] = value.constructor.name
        }
    }

    val serialized = JSON.stringify(plain)
    restorableContext.memory[restorableContext.restoreId] = serialized.asDynamic()
//    println("storing checkpoint for ${restorableContext.restoreId} - ${JSON.stringify(restorableContext.memory)}")
}

private fun restoreCoroutineIfNecessary(continuation: Continuation<Unit>, evalfn: (String) -> Any?): Boolean {
    val restorableContext = continuation.context[RestorableContext.Key]
    if (restorableContext != null) {
        val saved = restorableContext.memory[restorableContext.restoreId].unsafeCast<String?>()
        if (saved != null) {
            val deserializedCont = JSON.parse<Any>(saved)

            val resumeFromState = deserializedCont.asDynamic().state_0
            val currentState = continuation.asDynamic().state_0

            if(currentState >= resumeFromState){
                return false // we would go backwards in time
            }

            println("restoring ${restorableContext.restoreId} from checkpoint. Skipping from $currentState to $resumeFromState")
            restoreCoroutineToState(continuation, deserializedCont, evalfn)
            delete(restorableContext.memory[restorableContext.restoreId])
            return true
        }
    }
    return false
}

private fun restoreCoroutineToState(continuation: Continuation<Unit>, deserializedCont: Any, evalfn: (String) -> dynamic) {

    js("Object").assign(continuation, deserializedCont)

    deserializedCont.asRecord().keys.forEach {
        if (it.endsWith("\$prototype\$name")) {
            delete(continuation.asDynamic()[it])
            val value = deserializedCont.asRecord()[it]
            val constructor = evalfn("$value")
            val propertyName = it.removeSuffix("\$prototype\$name")
            if (constructor != null && constructor.prototype != null) {
                val realValue = continuation.asDynamic()[propertyName]
                js("Object").setPrototypeOf(realValue, constructor.prototype)
            } else {
                println("unable to restore prototype $value to $propertyName")
            }
        }
    }
}






/*******************
 * TESTS below
 *******************/


data class SmallClass(val value: Int)

suspend fun noActualSuspend() {
    println("not suspending")
}

suspend fun fancyCalculation(input: Int): Int {
    yield()
    println("fancyTest() called at TestGame.cpu=${++TestGame.cpu} TestGame.tick=${TestGame.tick} result=$input + 1")
    return input + 1
}

object TestGame {
    var cpu = 0
    var tick = 0
    val memory = mutableRecordOf<String, String>()
}

class RestorableCoroutinesTest {

    fun runTest(tick: Int) {
        TestGame.cpu = 0
        TestGame.tick = tick
        Scheduler.initProcess(RestorableContext(restoreId = "123", TestGame.memory)) {
            restoreIfNecessary(::eval)
            var myNumber = 0
            println("inside initProcess")
            myNumber = fancyCalculation(myNumber)
            val mySmallClass = SmallClass(myNumber)
            println("myNumber=${myNumber}")
            myNumber = fancyCalculation(myNumber)
            println("myNumber=${myNumber}")
            noActualSuspend()
            noActualSuspend()
            saveAndSuspend()
            println("after yieldAndSave")
            noActualSuspend()
            myNumber = fancyCalculation(myNumber)
            println("myNumber=${myNumber}")
            println(JSON.stringify(mySmallClass))
            println(mySmallClass)
            assertEquals(SmallClass(1), mySmallClass)
        }
    }

    @Test
    fun shouldSerialize() {
        runTest(1)
        runTest(2)
        runTest(3)
        runTest(4)
        runTest(5)
        runTest(6)
    }
}