Fctx.kt

fun Assembler.preserve(regs: Array<Reg>): FunctionContext {
    val ctx = FunctionContext(this, regs)
    ctx.preserve()
    return ctx
}

class FunctionContext(private val assembler: Assembler, private val regs: Array<Reg>) {
    fun byteSize() = regs.size * 4

    fun preserve() = with(assembler) {
        addi(sp, sp, -regs.size * 4)
        regs.forEachIndexed { idx, reg ->
            sw(reg, idx * 4, sp)
        }
    }

    fun restore() = with(assembler) {
        restoreRegs()
        addi(sp, sp, regs.size * 4)
    }

    fun exit() = with(assembler) {
        jr(ra)
        nop()
    }

    fun restoreAndExit() = with(assembler) {
        restoreRegs()
        jr(ra)
        addi(sp, sp, regs.size * 4)
        nop()
    }

    private fun restoreRegs() = with(assembler) {
        regs.forEachIndexed { idx, reg ->
            lw(reg, idx * 4, sp)
        }
    }
}