import { colorGreedy, createGraph, Graph, maxCardinalitySearch } from "./data/graph"
import { setEquals, union } from "./data/set"
import { BasicBlock } from "./ir/block"
import { Loc } from "./ir/loc"
import { StackOffset } from "./sm83/cpu"

export type LivenessInfo = Array<Set<Loc>>

export const liveness = (block: BasicBlock): LivenessInfo => {
    const info: LivenessInfo = []

    info[block.insns.length] = new Set()
    for (let i = block.insns.length - 1; i >= 0; --i) {
        const insn = block.insns[i]
        const last = info[i + 1]

        info[i] = new Set()

        if (insn.type === "goto") {
            continue
        }

        if ("source" in insn && typeof insn.source !== "number") {
            info[i].add(insn.source)
        }

        last.forEach(loc => {
            if (insn.dest instanceof Loc && loc.ppName() === insn.dest.ppName()) {
                return
            }

            info[i].add(loc)
        })
    }

    let dirty = true
    while (dirty) {
        dirty = false
        for (let i = block.insns.length - 1; i >= 0; --i) {
            if (block.insns[i].type !== "goto") {
                continue
            }

            const labelSet = info[block.labels[block.insns[i].dest as string]]
            if (setEquals(info[i], labelSet)) {
                continue
            }

            dirty = true
            info[i] = union(info[i], labelSet)
        }
    }

    return info
}

export const interference = (block: BasicBlock, live: LivenessInfo): Graph<Loc> =>
    createGraph((v, e) => {
        block.locs().forEach(loc => {
            v(loc.ppName(), loc)
        })

        block.insns.forEach((insn, i) =>
            live[i + 1].forEach(u => {
                if (insn.dest instanceof Loc && insn.dest.ppName() !== u.ppName()) {
                    e(insn.dest.ppName(), u.ppName())
                }
            })
        )
    })

export type RegAlloc = {
    [s: string]: string | StackOffset,
}

export const allocateRegisters = (block: BasicBlock, registers: Array<string>): RegAlloc => {
    const info = liveness(block)
    const graph = interference(block, info)
    const ordering = maxCardinalitySearch(graph)
    const coloring = colorGreedy(graph, ordering)

    const allocation: RegAlloc = {}
    const availableRegisters = new Set(registers)
    const colorMap: { [c: number]: string | StackOffset } = {}
    let nextStackOffset = 0

    Object.entries(coloring.colors).forEach(([vertex, color]) => {
        if (typeof colorMap[color] !== 'undefined') {
            allocation[vertex] = colorMap[color]
            return
        }

        let value = null
        if (availableRegisters.size == 0) {
            value = { offset: nextStackOffset++ }
        } else {
            const result = availableRegisters.values().next()
            value = result.value
            availableRegisters.delete(value)
        }

        allocation[vertex] = value
        colorMap[color] = value
    })

    return allocation
}