#!/usr/bin/env python3
""" format.py
Bit-packs instructions and verifies value ranges according to instruction formats.
Every instruction-line in the code segment needs to start with the opcode,
which has to be tagged with the instruction-format name. Arguments have to be provided
in the correct order, and their first tag is verified also.

Some immediates can be given in different sizes and will be sliced to size automatically.
For example the "j" instruction-format can take an "imm20" pre-sliced immediate, or an
"imm21" immediate, truncating the lowest bit (and asserting that it is zero).

>>> from io import StringIO
>>> # doctest: +REPORT_NDIFF
... print(subv.join_all(format(StringIO('''
... == 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 -34/imm21
... '''[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 ff/8 1/1 3e6/10 1/1
"""

import subv
import bits


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


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

    >>> _test_format(pack_i([(0x13, 'i'), (6, 'rd'), (0, 'funct3'), (0, 'rs'), (0x65, 'imm12')]))
    '13/7 06/5 0/3 00/5 065/12'
    """
    (op, rd, funct, rs, imm) = instr
    op_tag = op[1]
    op = bits.u(subv.untag(op), 7)
    rd = bits.u(subv.untag(rd, "rd"), 5)
    funct = bits.u(subv.untag(funct, "funct3"), 3)
    rs = bits.u(subv.untag(rs, "rs"), 5)
    imm = bits.i(subv.untag(imm, "imm12"), 12)

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


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

    >>> _test_format(pack_s([(0x23, 's'), (5, 'rs1'), (0, 'imm12'), (2, 'funct3'), (6, 'rs2')]))
    '23/7 00/5 2/3 05/5 06/5 00/7'
    """
    (op, rs1, imm, funct, rs2) = instr
    op = bits.u(subv.untag(op), 7)
    rs1 = bits.u(subv.untag(rs1, "rs1"), 5)
    imm = bits.i(subv.untag(imm, "imm12"), 12)
    funct = bits.u(subv.untag(funct, "funct3"), 3)
    rs2 = bits.u(subv.untag(rs2, "rs2"), 5)

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


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

    >>> _test_format(pack_b([(0x63, 'branch'), (0, 'funct3'), (6, 'rs1'), (0, 'rs2'), (0, 'imm12')]))
    '63/7 0/1 0/4 0/3 06/5 00/5 00/6 0/1'

    >>> _test_format(pack_b([(0x63, 'branch'), (0, 'funct3'), (6, 'rs1'), (0, 'rs2'), (-2, 'imm12')]))
    '63/7 1/1 e/4 0/3 06/5 00/5 3f/6 1/1'
    """
    (op, funct, rs1, rs2, imm) = instr
    op = bits.u(subv.untag(op), 7)
    funct = bits.u(subv.untag(funct, "funct3"), 3)
    rs1 = bits.u(subv.untag(rs1, "rs1"), 5)
    rs2 = bits.u(subv.untag(rs2, "rs2"), 5)

    if imm[1] == "imm13":
        imm = bits.i(subv.untag(imm, "imm13"), 13)
        imm = imm[12:1]
    else:
        imm = bits.i(subv.untag(imm, "imm12"), 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, funct, rs1, rs2, imm_md, imm_12]


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'), (-0x1234, 'imm20')]))
    '37/7 05/5 fedcc/20'
    """
    (op, rd, imm) = instr
    op = bits.u(subv.untag(op), 7)
    rd = bits.u(subv.untag(rd, "rd"), 5)
    imm = bits.i(subv.untag(imm, "imm20"), 20)

    return [op, rd, imm]


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

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

    >>> _test_format(pack_j([(0x6f, 'jal'), (0, 'rd', 'x0'), (-2, 'imm20')]))
    '6f/7 00/5 ff/8 1/1 3fe/10 1/1'
    """
    (op, rd, imm) = instr
    op = bits.u(subv.untag(op), 7)
    rd = bits.u(subv.untag(rd, "rd"), 5)

    if imm[1] == "imm21":
        imm = bits.i(subv.untag(imm, "imm21"), 21)
        imm = imm[20:1]
    else:
        imm = bits.i(subv.untag(imm, "imm20"), 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]


format_map = {
    # "r": pack_r,
    "i": pack_i,
    "s": pack_s,
    "b": pack_b,
    "u": pack_u,
    "j": pack_j,
}


@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 format label: {}".format(op)

            (op, format) = op
            if format not in format_map:
                raise ValueError("unknown instruction format: {}".format(format))
            formatter = format_map[format]

            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)
