Inter.kt

package com.d360.oam


class Inter(val dependencies: Code, val code: Code, val prims: Prims) {
    constructor(code: Code, prims: Prims) : this(listOf(), code, prims)
    private var instance: Instance = Instance(0, mutableListOf())
    private var values = mutableListOf<Value>()
    private var coeffects = mutableListOf<Coeffect>()

    private var queue = mutableListOf<Token>()

    sealed class Realized {
        data class Values(val vals: List<Value>) : Realized()
        data class Pend(val pending: Pending) : Realized()
        object Stopped : Realized()
    }

    private fun getCode(pc: Int) : CodeFun {
        return if (pc < 0) {
            code[Math.abs(pc) - 1]
        } else {
            dependencies[pc]
        }
    }


    private tailrec fun tick(pc: Pc, stack: Stack, env: Env) {
        val op = getCode(pc.pc).code[pc.epc]
        fun realized(args: List<Int>): Realized {
            val values = mutableListOf<Value>()
            for (arg in args) {
                val argV = env[arg]
                when (argV) {
                    is EnvPending -> {
                        val v = argV.pend
                        v.value.let {
                            when (it) {
                                is Pend -> {
                                    return Realized.Pend(argV.pend)
                                }
                                is PendStopped -> {
                                    return Realized.Stopped
                                }
                                is PendVal -> values.add(it.v)
                            }
                        }
                    }
                    is EnvVal -> {
                        values.add(argV.v)
                    }
                }
            }
            return Realized.Values(values)
        }

        when (op) {
            is OpConst -> {
                publish(stack, env, VConst(op.const))
                halt(stack, env)
            }
            is OpLabel -> {
                publish(stack, env, VLabel(op.pc))
                halt(stack, env)
            }
            is OpStop -> halt(stack, env)
            is OpParallel -> {
                incrementInstances(stack)
                tick(Pc(pc.pc, op.left), stack, env.toMutableList())
                tick(Pc(pc.pc, op.right), stack, env)
            }
            is OpOtherwise -> {
                val frame = FOtherwise(false, 1, Pc(pc.pc, op.right))
                tick(Pc(pc.pc, op.left), Cons(frame, stack), env)
            }
            is OpPruning -> {
                val pending = Pending(Pend, mutableListOf())
                val frame = FPruning(1, pending)
                op.index?.let { env.set(it, EnvPending(pending)) }
                tick(Pc(pc.pc, op.right), Cons(frame, stack), env)
                tick(Pc(pc.pc, op.left), stack, env)
            }
            is OpSequential -> {
                val frame = FSequential(op.index, Pc(pc.pc, op.right))
                tick(Pc(pc.pc, op.left), Cons(frame, stack), env)
            }
            is OpClosure -> {
                publish(stack, env, VClosure(op.pc, op.toCopy, env))
                halt(stack, env)
            }
            is OpTailCall -> {
                val t = op.target as? TFun ?: throw RuntimeException()
                val f = getCode(t.pc)
                var i = 0
                for (arg in op.args) {
                    env[i] = env[arg]
                    i++
                }
                queue.add(Token(Pc(t.pc, f.code.lastIndex), env, stack))
            }
            is OpCall -> {
                val t = op.target
                when (t) {
                    is TPrim -> {
                        val impl = prims[t.impl]
                        val args = realized(op.args)
//                        System.err.println("----CALL PRIM ${op.target} ${args}")
                        when (args) {
                            is Realized.Values -> {
                                val res = impl(args.vals)
                                when (res) {
                                    is PrimVal -> {
                                        publish(stack, env, res.v)
                                        halt(stack, env)
                                    }
                                    is PrimHalt -> halt(stack, env)
                                    is PrimUnsupported -> throw UnsupportedOperationException()
                                }
                            }
                            is Realized.Stopped -> halt(stack, env)
                            is Realized.Pend ->
                                args.pending.waiters.add(Token(pc, env, stack))
                        }
                    }
                    is TFun -> {
                        val f = getCode(t.pc)
                        val newEnv = MutableList(f.envSize) { index ->
                            if (index < op.args.size) {
                                env[op.args[index]]
                            } else {
                                EnvVal(VConst(ConstNull))
                            }
                        }
                        val frame = FCall(env)
                        tick(Pc(t.pc, f.code.lastIndex), Cons(frame, stack), newEnv)
                    }
                    is TDynamic -> {
                        val realized = realized(listOf(t.index))
                        when (realized) {
                            is Realized.Values -> {
                                val v = realized.vals[0]
                                when (v) {
                                    is VClosure -> {
                                        val f = getCode(v.pc)
                                        val newEnv = MutableList(f.envSize) { index ->
                                            if (index < v.toCopy) {
                                                v.env[index]
                                            } else if (index < op.args.size + v.toCopy) {
                                                env[op.args[index - v.toCopy]]
                                            } else {
                                                EnvVal(VConst(ConstNull))
                                            }
                                        }
                                        val frame = FCall(env)
                                        tick(Pc(v.pc, f.code.lastIndex), Cons(frame, stack), newEnv)
                                    }
                                    is VLabel -> {
                                        val f = getCode(v.pc)
                                        val newEnv = MutableList(f.envSize) { index ->
                                            if (index < op.args.size) {
                                                env[op.args[index]]
                                            }  else {
                                                EnvVal(VConst(ConstNull))
                                            }
                                        }
                                        val frame = FCall(env)
                                        tick(Pc(v.pc, f.code.lastIndex), Cons(frame, stack), newEnv)
                                    }
                                    is VTuple -> {
                                        val index = realized(listOf(op.args[0]))

                                        when (index) {
                                            is Realized.Values -> {
                                                val indexV = index.vals[0]
//                                                System.err.println("---TUPLE ACCESS ${v} ${indexV}")
                                                when (indexV) {
                                                    is VConst -> {
                                                        if (indexV.const is ConstInt) {
                                                            publish(stack, env, v.v[indexV.const.v])
                                                            halt(stack, env)
                                                        } else {
                                                            throw RuntimeException()
                                                        }
                                                    }
                                                    else -> throw RuntimeException()
                                                }
                                            }
                                            is Realized.Stopped -> halt(stack, env)
                                            is Realized.Pend ->
                                                index.pending.waiters.add(Token(pc, env, stack))
                                        }
                                    }
                                }
                            }
                            is Realized.Stopped -> halt(stack, env)
                            is Realized.Pend ->
                                realized.pending.waiters.add(Token(pc, env, stack))
                        }
                    }
                }
            }
            is OpCoeffect -> {
                val realized = realized(listOf(op.index))
                when (realized) {
                    is Realized.Values -> {
                        val v = realized.vals[0]
                        val id = instance.currentCoeffect
                        instance.currentCoeffect += 1
                        instance.blocks.add(Block(id, Token(pc, env.toMutableList(), stack)))
                        coeffects.add(Coeffect(id, v))
                    }
                    is Realized.Stopped -> halt(stack, env)
                    is Realized.Pend ->
                        realized.pending.waiters.add(Token(pc, env, stack))
                }
            }
        }
    }

    private tailrec fun publish(stack: Stack, env: Env, v: Value) {
        val s = stack.value
        when (s) {
            is FResult -> values.add(v)
            is FSequential -> {
                incrementInstances(stack)
                s.index?.let { env.set(it, EnvVal(v)) }
                tick(s.right, stack.next!!, env)
            }
            is FOtherwise -> {
                s.first_value = true
                publish(stack.next!!, env, v)
            }
            is FPruning -> {
                if (s.pending.value is PendVal) {
                    return
                }
                s.pending.value = PendVal(v)
                for (w in s.pending.waiters) {
                    tick(w.pc, w.stack, w.env)
                }
            }
            is FCall -> {
                publish(stack.next!!, s.env, v)
            }
        }
    }

    private fun halt(stack: Stack, env: Env) {
        val s = stack.value
        when (s) {
            is FResult -> {
            }
            is FOtherwise -> {
                if (!s.first_value && s.instances == 1) {
                    tick(s.otherwise, stack.next!!, env)
                } else if (s.instances == 1) {
                    halt(stack.next!!, env)
                } else {
                    s.instances -= 1
                }
            }
            is FPruning -> {
                if (s.pending.value is Pend && s.instances == 1) {
                    s.pending.value = PendStopped
                } else {
                    s.instances -= 1
                }
            }
            is FSequential, is FCall -> {
                halt(stack.next!!, env)
            }
        }
    }

    private fun incrementInstances(stack: Stack) {
        loop@ for (s in stack) {
            when (s) {
                is FOtherwise -> {
                    s.instances += 1; break@loop
                }
                is FPruning -> {
                    s.instances += 1; break@loop
                }
            }
        }
    }

    fun isAlive(stack: Stack) : Boolean {
        for (s in stack) {
            if (s is FPruning && !(s.pending.value is Pend)) {
                return false
            }
        }
        return true
    }

    fun checkKilled(justUnblocked: Int) : List<Int> {
        val killed = mutableListOf<Int>()
        val newBlocks = mutableListOf<Block>()
        for (block in instance.blocks) {
            when {
                block.id == justUnblocked -> {}
                !isAlive(block.token.stack) -> killed.add(block.id)
                else -> newBlocks.add(block)
            }
        }
        instance.blocks.clear()
        instance.blocks.addAll(newBlocks)
        return killed
    }

    fun runLoop() {
        while (queue.isNotEmpty()) {
            val token = queue.removeAt(queue.size - 1)
            tick(token.pc, token.stack, token.env)
        }
    }

    fun run(): Snapshot {
        val l = code.last()
        val pc = Pc(-code.lastIndex - 1, l.code.lastIndex)
        val stack = Cons(FResult, null)
        val env: MutableList<EnvValue> = MutableList(l.envSize, { EnvVal(VConst(ConstNull)) })
        values = mutableListOf<Value>()
        coeffects = mutableListOf<Coeffect>()
        queue.add(Token(pc, env, stack))
        runLoop()
        return Snapshot(values, coeffects, listOf(), instance)
    }

    fun unblock(instance: Instance, id: Int, value: Value): Snapshot {
        this.instance = instance
        values = mutableListOf<Value>()
        coeffects = mutableListOf<Coeffect>()
        val block = instance.blocks.find { it.id == id }
        if (block == null) {
            throw RuntimeException("Unknown coeffect")
        } else {
            publish(block.token.stack, block.token.env, value)
            runLoop()
            val killed = checkKilled(id)
            return Snapshot(values, coeffects, killed, instance)
        }
    }

    fun unblock(id: Int, value: Value): Snapshot {
        return unblock(instance, id, value)
    }
}