#!/usr/bin/env python3
from subx import clean, classify, parse_segment, parse_instr, is_lref, untag, format_instr
from riscv import u, i, bit_concat, bit_slice

instr_map = {
    'opr': ('r', 0x33),
    'load': ('i', 0x03),
    'opi': ('i', 0x13),
    'jalr': ('i', 0x67),
    'store': ('s', 0x23),
    'branch': ('b', 0x63),
    'lui': ('u', 0x37),
    'auipc': ('u', 0x17),
    'jal': ('j', 0x6f),
}

def pack_u(instr):
    (op, rd, imm) = instr
    op = u(untag(op), 7)
    rd = u(untag(rd, 'rd'), 5)
    if not is_lref(imm):
        imm = i(untag(imm, 'imm20'), 20)

    return [op, rd, imm]

def pack_i(instr):
    (op, sub, rd, rs, imm) = instr
    op = u(untag(op), 7)
    sub = u(untag(sub, 'subop'), 3)
    rd = u(untag(rd, 'rd'), 5)
    rs = u(untag(rs, 'rs'), 5)
    if not is_lref(imm):
        imm = i(untag(imm, 'imm12'), 12)

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

def pack_s(instr):
    (op, sub, rs1, rs2, imm) = instr
    op = u(untag(op), 7)
    sub = u(untag(sub, 'subop'), 3)
    rs1 = u(untag(rs1, 'rs'), 5)
    rs2 = u(untag(rs2, 'rs'), 5)
    if not is_lref(imm):
        imm = i(untag(imm, 'off12'), 12)
        imm_lo = bit_slice(imm, 4, 0)
        imm_hi = bit_slice(imm, 11, 5)
    else:
        label = untag(imm, 'off12')
        imm_lo = (label, 5, 'off12')
        imm_hi = (label, 7, 'off12>>5')
    return [op, imm_lo, sub, rs1, rs2, imm_hi]

def pack_j(instr):
    (op, rd, imm) = instr
    op = u(untag(op), 7)
    rd = u(untag(rd, 'rd'), 5)

    if not is_lref(imm):
        imm = i(untag(imm, 'off21'), 21)
        if imm & 0b1 == 1:
            raise ArgumentError("J-type offsets have to be halfword-aligned")
        imm_lo = bit_slice(imm, 10, 1)
        imm_11 = bit_slice(imm, 11, 11)
        imm_hi = bit_slice(imm, 19, 12)
        imm_20 = bit_slice(imm, 20, 20)
    else:
        label = untag(imm, 'off21')
        imm_lo = (label, 10, 'off21>>1')
        imm_11 = (label, 1, 'off21>>11')
        imm_hi = (label, 8, 'off21>>12')
        imm_20 = (label, 1, 'off21>>20')

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

def pack_b(instr):
    (op, sub, rs1, rs2, imm) = instr
    op = u(untag(op), 7)
    sub = u(untag(sub, 'subop'), 3)
    rs1 = u(untag(rs1, 'rs'), 5)
    rs2 = u(untag(rs2, 'rs'), 5)
    if not is_lref(imm):
        imm = i(untag(imm, 'off13'), 13)
        if imm & 0b1 == 1:
            raise ArgumentError("B-type offsets have to be halfword-aligned")
        imm_lo = bit_slice(imm, 4, 1)
        imm_md = bit_slice(imm, 10, 5)
        imm_11 = bit_slice(imm, 11, 11)
        imm_12 = bit_slice(imm, 12, 12)
    else:
        label = untag(imm, 'off13')
        imm_lo = (label, 4, 'off13>>1')
        imm_md = (label, 6, 'off13>>5')
        imm_11 = (label, 1, 'off13>>11')
        imm_12 = (label, 1, 'off13>>12')

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

def format(iter):
    for line in iter:
        line = clean(line)
        if line == '':
            continue

        type = classify(line)

        if type == 'segment' or type == 'label':
            yield line
        else:
            instr = parse_instr(line)
            op = instr[0]
            if len(op) != 2:
                raise ValueError("instruction without op label")

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

            if format == 'u':
                out = pack_u(instr)
            elif format == 'i':
                out = pack_i(instr)
            elif format == 's':
                out = pack_s(instr)
            elif format == 'j':
                out = pack_j(instr)
            else:
                raise NotImplementedError()

            yield format_instr(out)

if __name__ == '__main__':
    import sys
    for line in format(sys.stdin):
        print(line)

import unittest
class TestPackers(unittest.TestCase):
    def test_pack_u(self):
        final = pack_u([(0x37, 'lui'), (5, 'rd', 't0'), (0x10010, 'imm20')])
        self.assertEqual(
            final,
            [(0x37, 7), (0x5, 5), (0x10010, 20)]
        )
        self.assertEqual(bit_concat(*final)[1], 32)

        label = pack_u([(0x37, 'lui'), (5, 'rd', 't0'), ('pos', 'imm20')])
        self.assertEqual(
            label,
            [(0x37, 7), (0x5, 5), ('pos', 'imm20')]
        )

    def test_pack_i(self):
        final = pack_i([
            (0x13, 'opi'),
            (0, 'subop', 'add'),
            (6, 'rd', 't1'),
            (0, 'rs', 'x0'),
            (0x65, 'imm12'),
        ])
        self.assertEqual(
            final,
            [(0x13, 7), (6, 5), (0, 3), (0, 5), (0x65, 12)]
        )
        self.assertEqual(bit_concat(*final)[1], 32)

    def test_pack_s(self):
        final = pack_s([
            (0x23, 'store'),
            (2, 'subop', 'word'),
            (5, 'rs', 't0'),
            (6, 'rs', 't1'),
            (0, 'off12'),
        ])
        self.assertEqual(
            final,
            [(0x23, 7), (0, 5), (2, 3), (5, 5), (6, 5), (0, 7)]
        )
        self.assertEqual(bit_concat(*final)[1], 32)

        label = pack_s([
            (0x23, 'store'),
            (2, 'subop', 'word'),
            (5, 'rs', 't0'),
            (6, 'rs', 't1'),
            ('home', 'off12'),
        ])
        self.assertEqual(label, [
            (0x23, 7),
            ('home', 5, 'off12'),
            (2, 3),
            (5, 5),
            (6, 5),
            ('home', 7, 'off12>>5'),
        ])

class TestE2E(unittest.TestCase):
    def test_e2e(self):
        from io import StringIO
        from textwrap import dedent
        inp = dedent('''\
            == 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
        ''')

        out = dedent('''\
            == code 0x80000000
            main:
            37/7 05/5 10010/20
            13/7 06/5 00/3 00/5 48/12
            23/7 00/5 02/3 05/5 06/5 00/7
            13/7 06/5 00/3 00/5 65/12
            23/7 00/5 02/3 05/5 06/5 00/7
            13/7 06/5 00/3 00/5 6c/12
            23/7 00/5 02/3 05/5 06/5 00/7
            13/7 06/5 00/3 00/5 6c/12
            23/7 00/5 02/3 05/5 06/5 00/7
            13/7 06/5 00/3 00/5 6f/12
            23/7 00/5 02/3 05/5 06/5 00/7
            13/7 06/5 00/3 00/5 0a/12
            23/7 00/5 02/3 05/5 06/5 00/7
            6f/7 00/5 main/8/off21>>12 main/1/off21>>11 main/10/off21>>1 main/1/off21>>20
        ''')

        got = ''
        for line in format(StringIO(inp)):
            got += line + '\n'

        self.assertEqual(got, out)