import { colorGreedy, createGraph, Graph, maxCardinalitySearch } from "./data/graph"
import { R8, StackOffset } from "./sm83/cpu"
import type { Loc } from "./ir/loc"
import type { SSA } from "./ir/ssa"

export type LivenessInfo = Array<Set<Loc>>

export const liveness = (block: Array<SSA>): LivenessInfo => {
    const info: LivenessInfo = []

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

        info[i] = new Set()

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

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

        last.forEach(loc => {
            if (loc.toString() === insn.dest.toString()) {
                return
            }

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

    return info
}

export const locations = (block: Array<SSA>): Set<Loc> => {
    const ls: Set<Loc> = new Set()

    block.forEach(ssa => {
        ls.add(ssa.dest)

        if (typeof ssa.source !== "number") {
            ls.add(ssa.source)
        }

        if ("source1" in ssa && typeof ssa.source1 !== "number") {
            ls.add(ssa.source1)
        }
    })

    return ls
}

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

        block.forEach((ssa, i) =>
            live[i + 1].forEach(u => {
                if (ssa.dest.toString() !== u.toString()) {
                    e(ssa.dest.toString(), u.toString())
                }
            })
        )
    })

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

export const allocateRegisters = (block: Array<SSA>): 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(Object.values(R8))
    const colorMap: { [c: number]: R8 | 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
}