import type { RegAlloc } from '../regalloc'
import { AbsInsn2, CopyInsn, GotoInsn, insnReduce2, LabelInsn, Operand, UnaryInsn } from '../ir/insn'
import { Loc, LocType } from '../ir/loc'
import { UnaryOp, BinaryOp } from '../ast'
import { BasicBlock } from '../ir/block'

export const allocate = (alloc: RegAlloc, block: BasicBlock): BasicBlock =>
    new BasicBlock(
        block.insns.map(insn => allocateInsn(alloc, insn))
    )

export const allocateInsn = (alloc: RegAlloc, insn: AbsInsn2): AbsInsn2 => insnReduce2<AbsInsn2>(
    (dest: Loc, source: Operand) => CopyInsn(
        allocateInsnDest(alloc, dest),
        typeof source === 'number' ? source : allocateInsnDest(alloc, source),
    ),
    (name: string) => LabelInsn(name),
    (name: string) => GotoInsn(name),
    (dest: Loc, op: UnaryOp | BinaryOp, source: Operand) => UnaryInsn(
        allocateInsnDest(alloc, dest),
        op,
        typeof source === 'number' ? source : allocateInsnDest(alloc, source),
    ),
    insn,
)

export const allocateInsnDest = (alloc: RegAlloc, loc: Loc): Loc =>
    loc.type === LocType.TEMPORARY
        ? lookupAllocation(alloc, loc.ppName())
        : loc

const lookupAllocation = (alloc: RegAlloc, name: string): Loc => {
    const loc = alloc[name]
    if (typeof loc === 'object') {
        throw new Error('stack offsets not supported yet')
    }
    return Loc.reg(loc)
}