#!/usr/bin/env python3
""" validate.py
Validates mnemonic labels match the numeric values given for opcodes,
register codes, etc; normalizes argument order and verifies  that there are
no missing arguments.

>>> from io import StringIO
>>> # doctest: +REPORT_NDIFF
... print(subv.join_all(validate(StringIO('''
... == code 0x80000000
... main:
... # load 0x10010000 (UART0) into t0
... 37/lui 5/rd/t0 0x10010/imm20
... # store 0x48 (H) in UART0+0
... 13/opi 0/subop/add 6/rd/t1 0/rs/x0 48/imm12
... 23/store 2/width/word 5/rs/t0 6/rs/t1 0/off12
... # store 0x65 (e) in UART0+0
... 13/opi 0/subop/add 6/rd/t1 0/rs/x0 65/imm12
... 23/store 2/width/word 5/rs/t0 6/rs/t1 0/off12
... # store 0x6c (l) in UART0+0
... 13/opi 0/subop/add 6/rd/t1 0/rs/x0 6c/imm12
... 23/store 2/width/word 5/rs/t0 6/rs/t1 0/off12
... # store 0x6c (l) in UART0+0
... 13/opi 0/subop/add 6/rd/t1 0/rs/x0 6c/imm12
... 23/store 2/width/word 5/rs/t0 6/rs/t1 0/off12
... # store 0x6f (o) in UART0+0
... 13/opi 0/subop/add 6/rd/t1 0/rs/x0 6f/imm12
... 23/store 2/width/word 5/rs/t0 6/rs/t1 0/off12
... # store 0x0a (\\\\n) in UART0+0
... 13/opi 0/subop/add 6/rd/t1 0/rs/x0 0a/imm12
... 23/store 2/width/word 5/rs/t0 6/rs/t1 0/off12
... # jump back up to the top
... 6f/jal 0/rd/x0 main/off21
... '''[1:-1]))))
== code 0x80000000
main:
# load 0x10010000 (UART0) into t0
37/u 5/rd 10010/imm20
# store 0x48 (H) in UART0+0
13/i 6/rd 0/funct3 0/rs 48/imm12
23/s 5/rs1 0/imm12 2/funct3 6/rs2
# store 0x65 (e) in UART0+0
13/i 6/rd 0/funct3 0/rs 65/imm12
23/s 5/rs1 0/imm12 2/funct3 6/rs2
# store 0x6c (l) in UART0+0
13/i 6/rd 0/funct3 0/rs 6c/imm12
23/s 5/rs1 0/imm12 2/funct3 6/rs2
# store 0x6c (l) in UART0+0
13/i 6/rd 0/funct3 0/rs 6c/imm12
23/s 5/rs1 0/imm12 2/funct3 6/rs2
# store 0x6f (o) in UART0+0
13/i 6/rd 0/funct3 0/rs 6f/imm12
23/s 5/rs1 0/imm12 2/funct3 6/rs2
# store 0x0a (\\n) in UART0+0
13/i 6/rd 0/funct3 0/rs a/imm12
23/s 5/rs1 0/imm12 2/funct3 6/rs2
# jump back up to the top
6f/j 0/rd main/off21
"""

import subv


def pop_piece(line, labels):
    for part in line:
        if part[1] in labels:
            line.remove(part)
            return (part[0], labels[0])

    raise ValueError("Expected a part with labels {}".format(labels))


def pop_immediate(line, sizes, allow_offset=False):
    if isinstance(sizes, int):
        sizes = [sizes]

    for part in line:
        imm = subv.parse_immediate(part[1])
        if imm and imm["size"] in sizes:
            line.remove(part)
            if not isinstance(part[0], str):
                imm["mode"] = "imm"

            return (part[0], subv.format_immediate(imm))

    raise ValueError(
        "Expected an immediate with size {}".format("/".join(str(s) for s in sizes))
    )


REG_NAMES = "zero,ra,sp,gp,tp,t0,t1,t2,s0,s1,a0,a1,a2,a3,a4,a5,a6,a7,s2,s3,s4,s5,s6,s7,s8,s9,s10,s11,t3,t4,t5,t6"
REG_NAMES = {name: i for i, name in enumerate(REG_NAMES.split(","))}
REG_NAMES.update({"x{}".format(i): i for i in range(32)})


def pop_register(line, labels):
    for part in line:
        if part[1] in labels:
            line.remove(part)

            value = part[0]
            if value < 0 or value > 31:
                raise ValueError(
                    "Invalid register value {} (expected 0...31)".format(value)
                )

            if len(part) > 2:
                name = part[2]
                try:
                    assert value == REG_NAMES[name], ValueError(
                        "Invalid register value {} for '{}', expected {}".format(
                            value, name, REG_NAMES[name]
                        )
                    )
                except KeyError:
                    raise ValueError("Unknown register name '{}'".format(name))

            return (part[0], labels[0])

    raise ValueError("Expected a register with labels {}".format(labels))


def validate_r(inputs):
    op = inputs.pop(0)

    dest = pop_register(inputs, ["rd", "dest"])
    funct = pop_piece(inputs, ["funct7", "funct", "subop"])
    rs1 = pop_register(inputs, ["rs1", "rs", "src1", "src"])
    rs2 = pop_register(inputs, ["rs2", "rs", "src2", "src"])

    assert len(inputs) == 0, ValueError("Extra arguments: {}".format(inputs))

    return [
        (op[0], "r"),
        dest,
        funct,
        rs1,
        rs2,
    ]


def validate_i(inputs):
    op = inputs.pop(0)

    rs_name = "src"
    if op[1] in ["load", "jalr"]:
        rs_name = "base"

    dest = pop_register(inputs, ["rd", "dest"])
    funct = pop_piece(inputs, ["funct3", "funct", "subop"])
    rs = pop_register(inputs, ["rs", rs_name])
    imm = pop_immediate(inputs, 12)

    assert len(inputs) == 0, ValueError("Extra arguments: {}".format(inputs))

    return [
        (op[0], "i"),
        dest,
        funct,
        rs,
        imm,
    ]


def validate_s(inputs):
    op = inputs.pop(0)
    base = pop_register(inputs, ["rs1", "rs", "base"])
    offset = pop_immediate(inputs, 12, allow_offset=True)
    width = pop_piece(inputs, ["funct3", "funct", "width"])
    src = pop_register(inputs, ["rs2", "rs", "src"])

    assert len(inputs) == 0, ValueError("Extra arguments: {}".format(inputs))

    return [
        (op[0], "s"),
        base,
        offset,
        width,
        src,
    ]


def validate_b(inputs):
    op = inputs.pop(0)

    funct = pop_piece(inputs, ["funct3", "funct", "subop"])
    rs1 = pop_register(inputs, ["rs1", "rs", "src1", "src"])
    rs2 = pop_register(inputs, ["rs2", "rs", "src2", "src"])
    offset = pop_immediate(inputs, [12, 13], allow_offset=True)

    assert len(inputs) == 0, ValueError("Extra arguments: {}".format(inputs))

    return [
        (op[0], "b"),
        funct,
        rs1,
        rs2,
        offset,
    ]


def validate_u(inputs):
    op = inputs.pop(0)
    rd = pop_register(inputs, ["rd", "dest"])
    imm = pop_immediate(inputs, 20, allow_offset=True)

    assert len(inputs) == 0, ValueError("Extra arguments: {}".format(inputs))

    return [
        (op[0], "u"),
        rd,
        imm,
    ]


def validate_j(inputs):
    op = inputs.pop(0)

    dest = pop_register(inputs, ["rd", "dest"])
    offset = pop_immediate(inputs, [21, 20], allow_offset=True)

    assert len(inputs) == 0, ValueError("Extra arguments: {}".format(inputs))

    return [
        (op[0], "j"),
        dest,
        offset,
    ]


instr_map = {
    "opr": (validate_r, 0x33),
    "opi": (validate_i, 0x13),
    "load": (validate_i, 0x03),
    "jalr": (validate_i, 0x67),
    "system": (validate_i, 0x73),
    "store": (validate_s, 0x23),
    "branch": (validate_b, 0x63),
    "lui": (validate_u, 0x37),
    "auipc": (validate_u, 0x17),
    "jal": (validate_j, 0x6F),
}


@subv.with_parsed_lines
def validate(iter):
    for segment, line in iter:
        if line["type"] == "instr" and segment == "code":
            op = line["instr"][0]
            assert len(op) == 2, "instruction without op label: {}".format(op)

            (op, label) = op
            if label not in instr_map:
                raise ValueError("unknown op label: {}".format(label))
            (validator, expected) = instr_map[label]
            if op != expected:
                raise ValueError(
                    "opcode {} doesn't match label {} (expected {})".format(
                        op, label, expected
                    )
                )

            line["instr"] = validator(line["instr"][:])
            line = yield subv.format(line)
        else:
            line = yield line["raw"]


if __name__ == "__main__":
    import sys

    for line in validate(sys.stdin):
        print(line)