#!/usr/bin/env python3
""" survey.py
Resolves label references and substitutes them with literal values.

>>> from io import StringIO
>>> # doctest: +REPORT_NDIFF
... print(subv.join_all(survey(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 2/funct3 5/rs1 0/imm12 6/rs2
... # store 0x65 (e) in UART0+0
... 13/i 6/rd 0/funct3 0/rs 65/imm12
... 23/s 2/funct3 5/rs1 0/imm12 6/rs2
... # store 0x6c (l) in UART0+0
... 13/i 6/rd 0/funct3 0/rs 6c/imm12
... 23/s 2/funct3 5/rs1 0/imm12 6/rs2
... # store 0x6c (l) in UART0+0
... 13/i 6/rd 0/funct3 0/rs 6c/imm12
... 23/s 2/funct3 5/rs1 0/imm12 6/rs2
... # store 0x6f (o) in UART0+0
... 13/i 6/rd 0/funct3 0/rs 6f/imm12
... 23/s 2/funct3 5/rs1 0/imm12 6/rs2
... # store 0x0a (\\\\n) in UART0+0
... 13/i 6/rd 0/funct3 0/rs a/imm12
... 23/s 2/funct3 5/rs1 0/imm12 6/rs2
... # jump back up to the top
... 6f/j 0/rd 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 2/funct3 5/rs1 0/imm12 6/rs2
# store 0x65 (e) in UART0+0
13/i 6/rd 0/funct3 0/rs 65/imm12
23/s 2/funct3 5/rs1 0/imm12 6/rs2
# store 0x6c (l) in UART0+0
13/i 6/rd 0/funct3 0/rs 6c/imm12
23/s 2/funct3 5/rs1 0/imm12 6/rs2
# store 0x6c (l) in UART0+0
13/i 6/rd 0/funct3 0/rs 6c/imm12
23/s 2/funct3 5/rs1 0/imm12 6/rs2
# store 0x6f (o) in UART0+0
13/i 6/rd 0/funct3 0/rs 6f/imm12
23/s 2/funct3 5/rs1 0/imm12 6/rs2
# store 0x0a (\\n) in UART0+0
13/i 6/rd 0/funct3 0/rs a/imm12
23/s 2/funct3 5/rs1 0/imm12 6/rs2
# jump back up to the top
6f/j 0/rd -34/imm21

>>> # doctest: +REPORT_NDIFF
... print(subv.join_all(survey(StringIO('''
... == code 0x80000000
... # repeatedly print "Hi\\\\n"
... main:
... # load 0x10010000 (UART0) into t0
... 37/u 5/rd/t0 10010/imm20
... # store 0x48 (H) in UART0+0
... 13/i 0/subop/add 6/rd/t1 0/rs/x0 48/imm12
... 23/s 2/subop/word 5/rs/t0 6/rs/t1 0/off12
... # store 0x69 (i) in UART0+0
... 13/i 0/subop/add 6/rd/t1 0/rs/x0 69/imm12
... 23/s 2/subop/word 5/rs/t0 6/rs/t1 0/off12
... # store 0x0a (\\\\n) in UART0+0
... 13/i 0/subop/add 6/rd/t1 0/rs/x0 0a/imm12
... 23/s 2/subop/word 5/rs/t0 6/rs/t1 0/off12
... 17/u 7/rd/t2 main/off20hi
... 67/i 0/subop 0/rd/x0 7/rs/t2 main+4/off12lo
... '''[1:-1]))))
== code 0x80000000
# repeatedly print "Hi\\n"
# main:
# load 0x10010000 (UART0) into t0
37/u 5/rd/t0 10010/imm20
# store 0x48 (H) in UART0+0
13/i 0/subop/add 6/rd/t1 0/rs/x0 48/imm12
23/s 2/subop/word 5/rs/t0 6/rs/t1 0/off12
# store 0x69 (i) in UART0+0
13/i 0/subop/add 6/rd/t1 0/rs/x0 69/imm12
23/s 2/subop/word 5/rs/t0 6/rs/t1 0/off12
# store 0x0a (\\n) in UART0+0
13/i 0/subop/add 6/rd/t1 0/rs/x0 a/imm12
23/s 2/subop/word 5/rs/t0 6/rs/t1 0/off12
17/u 7/rd/t2 0/imm20
67/i 0/subop 0/rd/x0 7/rs/t2 -1c/imm12
"""
import subv


def observe(word, pc, map):
    """resolve a label reference.

    >>> observe((3, '3'), 0, {})
    (3, '3')
    >>> subv.format_part(observe(('label', 'imm32'), 0, {'label': 0x1234}))
    '1234/imm32'
    >>> subv.format_part(observe(('label-2', 'imm20',), 0, {'label': 0x1234}))
    '1232/imm20'
    >>> subv.format_part(observe(('label', 'off32'), 0x1000, {'label': 0x1234}))
    '234/imm32'
    >>> subv.format_part(observe(('label+4', 'off32'), 0x1000, {'label': 0x1234}))
    '238/imm32'
    >>> subv.format_part(observe(('label', 'off32'), 0x1234, {'label': 0x1000}))
    '-234/imm32'
    >>> subv.format_part(observe(('label', 'off20hi'), 0x1000, {'label': 0x3234}))
    '2/imm20'
    >>> subv.format_part(observe(('label', 'off12lo'), 0x1000, {'label': 0x3234}))
    '234/imm12'
    >>> subv.format_part(observe(('label', 'off20hi'), 0x3234, {'label': 0x1000}))
    '-2/imm20'
    >>> subv.format_part(observe(('label', 'off12lo'), 0x3234, {'label': 0x1000}))
    '-234/imm12'

    >>> observe(('label', 'imm32'), 0, {})
    Traceback (most recent call last):
        ...
    ValueError: undefined label 'label'
    >>> observe(('label', 'imm8'), 0, {'label': 0x1234})
    Traceback (most recent call last):
        ...
    ValueError: label reference out of range (8 bit): 'label' (0x1234)
    """
    if not subv.is_reference(word):
        if len(word) != 2:
            return word

        ref = subv.parse_immediate(word[1])
        if not ref or "split" not in ref:
            return word

        value = word[0]

        sgn = -1 if value < 0 else 1
        value *= sgn

        if ref["split"] == "hi" and ref["size"] == 20:
            value = value >> 12
        elif ref["split"] == "lo" and ref["size"] == 12:
            value = value & 0xFFF
        else:
            raise ValueError(
                "unknown label split mode {mode}{size}{split}".format(**ref)
            )

        return (value, "imm{size}".format(**ref))

    ref = subv.parse_reference(word)

    if ref["label"] not in map:
        raise ValueError("undefined label '{}'".format(ref["label"]))

    value = map[ref["label"]] + ref["offset"]

    if ref["mode"] == "off":
        value -= pc

    if "split" in ref:
        sgn = -1 if value < 0 else 1
        value *= sgn

        if ref["split"] == "hi" and ref["size"] == 20:
            value = value >> 12
        elif ref["split"] == "lo" and ref["size"] == 12:
            value = value & 0xFFF
        else:
            raise ValueError(
                "unknown label split mode {mode}{size}{split}".format(**ref)
            )

        value *= sgn

    if value > 0:
        if value >= (1 << ref["size"]):
            raise ValueError(
                "label reference out of range ({size} bit): '{label}' (0x{:x})".format(
                    value, **ref
                )
            )
    else:
        if value < -(1 << (ref["size"] - 1)):
            raise ValueError(
                "label reference out of range ({size} bit): '{label}' (-0x{:x})".format(
                    abs(value), **ref
                )
            )

    return (value, "imm{size}".format(**ref))


@subv.with_parsed_lines
def survey(iter):
    queue = []
    map = {}
    addr, bitcount = -1, 0
    for segment, line in iter:
        line["addr"] = addr
        queue.append(line)

        # step forward addr
        type = line["type"]
        if type == "segment":
            addr, bitcount = line["segment"][1], 0
        elif type == "label":
            assert bitcount == 0, "label isn't byte aligned"
            map[line["label"]] = addr
        elif type == "instr":
            if len(line["instr"]) == 0:
                continue

            if segment == "data":
                for part in line["instr"]:
                    bitcount += int(part[1])

                if bitcount >= 8:
                    addr += bitcount // 8
                    bitcount = bitcount % 8
            elif segment == "code":
                assert (
                    bitcount == 0
                ), "instruction isn't byte aligned ({} bitcount left)".format(
                    8 - bitcount
                )
                assert addr % 2 == 0, "instruction isn't 2-byte aligned"

                op = line["instr"][0]
                instr_size = 0
                if op[1] in "usirjb":
                    addr += 4
                else:
                    raise ValueError("Unknown instruction format '{}'".format(op[1]))
            else:
                raise ValueError("unknown segment '{}'".format(segment))

    for i, line in enumerate(queue, start=1):
        try:
            type = line["type"]
            if type == "instr":
                instr = []
                for part in line[type]:
                    instr.append(observe(part, line["addr"], map))
                line[type] = instr
                yield subv.format(line)
            elif type == "label":
                yield "# {}".format(line["raw"])
            else:
                yield line["raw"]
        except Exception as e:
            msg = "failed to survey line"
            msg = msg + "\n{}:{}:  {}".format("(unknown)", i, subv.format(line))
            msg = msg + "\nparsed as {}".format(subv.dump(line))
            raise subv.SubVException(msg) from e


if __name__ == "__main__":
    import sys

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