#!/usr/bin/env python3
""" format.py
Verifies instruction formats, orders and pre-bit-packs arguments.
Every instruction-line in the code segment needs to start with the opcode,
which has to be correctly tagged. Arguments have to be provided in the correct
order currently, and their first tag is verified also.

Labels can be referenced in the immXX/offXX fields of most instructions. If no
bitslice is given for the labels, the default slice is picked based on the
instruction.

>>> from io import StringIO
>>> # doctest: +REPORT_NDIFF
... print(subv.join_all(format(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/subop/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/subop/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/subop/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/subop/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/subop/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/subop/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/7 05/5 10010/20
# store 0x48 (H) in UART0+0
13/7 06/5 0/3 00/5 048/12
23/7 00/5 2/3 05/5 06/5 00/7
# store 0x65 (e) in UART0+0
13/7 06/5 0/3 00/5 065/12
23/7 00/5 2/3 05/5 06/5 00/7
# store 0x6c (l) in UART0+0
13/7 06/5 0/3 00/5 06c/12
23/7 00/5 2/3 05/5 06/5 00/7
# store 0x6c (l) in UART0+0
13/7 06/5 0/3 00/5 06c/12
23/7 00/5 2/3 05/5 06/5 00/7
# store 0x6f (o) in UART0+0
13/7 06/5 0/3 00/5 06f/12
23/7 00/5 2/3 05/5 06/5 00/7
# store 0x0a (\\n) in UART0+0
13/7 06/5 0/3 00/5 00a/12
23/7 00/5 2/3 05/5 06/5 00/7
# jump back up to the top
6f/7 00/5 main/off21/[19:12] main/off21/[11:11] main/off21/[10:1] main/off21/[20:20]
"""

import subv
import bits


def _test_format(words):
    return " ".join(map(str, words))


def ref_slice(ref, hi, lo):
    if hi < lo:
        raise ValueError("cant slice reverse range")
    elif hi >= ref["size"]:
        raise ValueError(
            "cant slice [{}:{}] from {} bit value".format(hi, lo, ref["size"])
        )

    return {
        **ref,
        "hi": hi,
        "lo": lo,
        "size": hi - lo + 1,
    }


def sign_trunc(val, size):
    """truncate and sign-shrink a number literal.

    >>> sign_trunc(bits.i(4, 32), 8)
    Bitfield(0x4, 8)
    >>> sign_trunc(bits.i(-4, 32), 8)
    Bitfield(0xfc, 8)
    >>> sign_trunc(bits.i(0x11ff, 32), 8)
    Bitfield(0x7f, 8)
    >>> sign_trunc(bits.i(-0x11ff, 32), 8)
    Bitfield(0x81, 8)
    """
    sign = val[val.size - 1]
    rest = val[size - 2 : 0]
    return sign & rest


def pack_u(instr):
    """verify & pack U-type instructions.

    >>> _test_format(pack_u([(0x37, 'lui'), (5, 'rd', 't0'), (0x10010, 'imm20')]))
    '37/7 05/5 10010/20'

    >>> _test_format(pack_u([(0x37, 'lui'), (5, 'rd', 't0'), ('pos', 'imm20')]))
    '37/7 05/5 pos/imm20/[31:12]'

    >>> _test_format(pack_u([(0x37, 'lui'), (5, 'rd', 't0'), ('pos', 'imm20', '[19:0]')]))
    '37/7 05/5 pos/imm20/[19:0]'

    >>> _test_format(pack_u([(0x37, 'lui'), (5, 'rd', 't0'), ('pos', 'imm20', '[31:0]')]))
    Traceback (most recent call last):
        ...
    AssertionError: expected 20 bit slice, got pos/imm20/[31:0]
    """
    (op, rd, imm) = instr
    op = bits.u(subv.untag(op), 7)
    rd = bits.u(subv.untag(rd, "rd"), 5)
    if subv.is_reference(imm):
        imm = bits.ref(imm, slice(31, 12))
    else:
        imm = bits.i(subv.untag(imm, "imm20"), 20)

    return [op, rd, imm]


def pack_i(instr):
    """verify & pack I-type instructions.

    >>> _test_format(pack_i([(0x13, 'opi'), (0, 'subop', 'add'), (6, 'rd', 't1'), (0, 'rs', 'x0'), (0x65, 'imm12')]))
    '13/7 06/5 0/3 00/5 065/12'

    >>> _test_format(pack_i([(0x13, 'opi'), (0, 'subop', 'add'), (6, 'rd', 't1'), (0, 'rs', 'x0'), ('label', 'imm12')]))
    '13/7 06/5 0/3 00/5 label/imm12/[11:0]'
    """
    (op, sub, rd, rs, imm) = instr
    op_tag = op[1]
    op = bits.u(subv.untag(op), 7)
    sub = bits.u(subv.untag(sub, "subop"), 3)
    rd = bits.u(subv.untag(rd, "rd"), 5)
    rs = bits.u(subv.untag(rs, "rs"), 5)
    if subv.is_reference(imm):
        imm = bits.ref(imm, slice(11, 0))
    else:
        imm = bits.i(subv.untag(imm, "imm12"), 12)

    return [op, rd, sub, rs, imm]


def pack_s(instr):
    """verify & pack S-type instructions.

    >>> _test_format(pack_s([(0x23, 'store'), (2, 'subop', 'word'), (5, 'rs', 't0'), (6, 'rs', 't1'), (0, 'off12')]))
    '23/7 00/5 2/3 05/5 06/5 00/7'

    >>> _test_format(pack_s([(0x23, 'store'), (2, 'subop', 'word'), (5, 'rs', 't0'), (6, 'rs', 't1'), ('home', 'off12')]))
    '23/7 home/off12/[4:0] 2/3 05/5 06/5 home/off12/[11:5]'
    """
    (op, sub, rs1, rs2, imm) = instr
    op = bits.u(subv.untag(op), 7)
    sub = bits.u(subv.untag(sub, "subop"), 3)
    rs1 = bits.u(subv.untag(rs1, "rs"), 5)
    rs2 = bits.u(subv.untag(rs2, "rs"), 5)

    if subv.is_reference(imm):
        imm = bits.ref(imm, slice(11, 0))
    else:
        imm = bits.i(subv.untag(imm, "off12"), 12)

    imm_lo = imm[4:0]
    imm_hi = imm[11:5]
    return [op, imm_lo, sub, rs1, rs2, imm_hi]


def pack_j(instr):
    """verify & pack J-type instructions.

    >>> _test_format(pack_j([(0x6f, 'jal'), (0, 'rd', 'x0'), (0, 'off20')]))
    '6f/7 00/5 00/8 0/1 000/10 0/1'

    >>> _test_format(pack_j([(0x6f, 'jal'), (0, 'rd', 'x0'), (-2, 'off20')]))
    '6f/7 00/5 ff/8 1/1 3fe/10 1/1'

    >>> _test_format(pack_j([(0x6f, 'jal'), (2, 'rd', 'x2'), ('home', 'off20')]))
    '6f/7 02/5 home/off20/[19:12] home/off20/[11:11] home/off20/[10:1] home/off20/[20:20]'
    """
    (op, rd, imm) = instr
    op = bits.u(subv.untag(op), 7)
    rd = bits.u(subv.untag(rd, "rd"), 5)

    if subv.is_reference(imm):
        imm = bits.ref(imm, slice(20, 1))
    else:
        imm = bits.i(subv.untag(imm, "off20"), 20)

    imm_lo = imm[9:0]
    imm_11 = imm[10]
    imm_hi = imm[18:11]
    imm_20 = imm[19]

    return [op, rd, imm_hi, imm_11, imm_lo, imm_20]


def pack_b(instr):
    """verify & pack B-type instructions.

    >>> _test_format(pack_b([(0x63, 'branch'), (0, 'subop', '=='), (6, 'rs'), (0, 'rs'), (0, 'off12')]))
    '63/7 0/1 0/4 0/3 06/5 00/5 00/6 0/1'

    >>> _test_format(pack_b([(0x63, 'branch'), (0, 'subop', '=='), (6, 'rs'), (0, 'rs'), (-2, 'off12')]))
    '63/7 1/1 e/4 0/3 06/5 00/5 3f/6 1/1'

    >>> _test_format(pack_b([(0x63, 'branch'), (0, 'subop', '=='), (6, 'rs'), (0, 'rs'), ('home', 'off12')]))
    '63/7 home/off12/[11:11] home/off12/[4:1] 0/3 06/5 00/5 home/off12/[10:5] home/off12/[12:12]'
    """
    (op, sub, rs1, rs2, imm) = instr
    op = bits.u(subv.untag(op), 7)
    sub = bits.u(subv.untag(sub, "subop"), 3)
    rs1 = bits.u(subv.untag(rs1, "rs"), 5)
    rs2 = bits.u(subv.untag(rs2, "rs"), 5)

    if subv.is_reference(imm):
        imm = bits.ref(imm, slice(12, 1))
    else:
        imm = bits.i(subv.untag(imm, "off12"), 12)

    imm_lo = imm[3:0]
    imm_md = imm[9:4]
    imm_11 = imm[10]
    imm_12 = imm[11]

    return [op, imm_11, imm_lo, sub, rs1, rs2, imm_md, imm_12]


instr_map = {
    # 'opr': (pack_r, 0x33),
    "load": (pack_i, 0x03),
    "opi": (pack_i, 0x13),
    "jalr": (pack_i, 0x67),
    "store": (pack_s, 0x23),
    "branch": (pack_b, 0x63),
    "lui": (pack_u, 0x37),
    "auipc": (pack_u, 0x17),
    "jal": (pack_j, 0x6F),
}


@subv.with_parsed_lines
def format(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 instruction label: {}".format(label))
            (formatter, expected) = instr_map[label]
            if op != expected:
                raise ValueError(
                    "opcode {} doesn't match label {} (expected {})".format(
                        op, label, expected
                    )
                )

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


if __name__ == "__main__":
    import sys

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