#!/usr/bin/env python3
# RISC-V assembly to fox32 translator
import argparse
import re

class Emitter:
    """
    The Emitter class emits fox32 assembly code.
    """

    def __init__(self, filename):
        self.filename = filename
        self.f = open(filename, 'w')

    def write(self, line):
        self.f.write(line + '\n')

    @staticmethod
    def _adjust_label(label):
        """
        Adjust a label to the requirements of fox32asm
        """
        l = label.replace('.', '_')
        if l.startswith('_'):
            return 'U' + l[1:]
        return l

    @staticmethod
    def _adjust_operand(op):
        """
        Adjust an operand to the requirements of fox32asm
        """
        if re.fullmatch(r'-?[0-9]x?[0-9a-fA-F]*', op):
            return f'{int(op, 0) & 0xffffffff:#x}'
        return op

    def insn(self, operation, *operands):
        """
        Emit an instruction
        """
        line = '\t' + operation
        if operands:
            line += '\t' + ', '.join([self._adjust_operand(op) for op in operands])
        self.write(line)

    def include(self, filename):
        """
        Emit an #include directive
        """
        self.write(f'#include "{filename}"')

    def comment(self, line):
        """
        Emit a comment
        """
        self.write(f'; {line}')

    def label(self, label):
        """
        Emit a label definition
        """
        self.write(self._adjust_label(label) + ':')

    def data(self, *data, size=32):
        """
        Emit data
        """
        for d in data:
            self.write(f'data.{size} {self._adjust_label(d)}')

    def zero(self, *data):
        """
        Emit zeros
        """
        for d in data:
            self.write(f'data.fill 0, {d}')

    def strz(self, *data):
        """
        Emit null-terminated string
        """
        for d in data:
            self.write(f'data.strz \"{d}\"')

    def nop(self):
        """
        Emit a nop instruction
        """
        self.insn('nop')

    def jump(self, label, condition=''):
        """
        Emit a jump to a label, optionally with a condition code
        """
        label = self._adjust_label(label)
        if condition:
            self.insn(f'{condition} rjmp', label)
        else:
            self.insn('rjmp', label)

    def jump_reg(self, reg):
        """
        Emit a jump to a register
        """
        self.insn('jmp', reg)

    def mov(self, dest, source, size=32):
        """
        Emit a mov instruction
        """
        self.insn(f'movz.{size}', dest, source)

    def store(self, rmem, rval, size=32):
        """
        Emit a mov instruction that performs a memory store
        """
        self.mov(f'[{rmem}]', rval, size)

    def load(self, rmem, rval, size=32):
        """
        Emit a mov instruction that performs a memory load
        """
        self.mov(rval, f'[{rmem}]', size)

    def rta(self, dest, source):
        """
        Emit an rta instruction
        """
        source = self._adjust_label(source)
        self.insn(f'rta', dest, source)

    def halt(self):
        """
        Emit an halt instruction
        """
        self.insn('halt')

class Converter:
    """
    This class converts RISC-V assembly code to fox32 assembly code
    """
    def __init__(self, input, output, *, debug=False, start_file='start.asm', additional_includes=[]):
        self.input_file = open(input, 'r')
        self.e = Emitter(output)

        self.debug = debug
        self.start_file = start_file
        self.additional_includes = additional_includes

        self.xregs = [f'x{n}' for n in range(32)]
        self.regs = ['zero', 'ra', 'sp', 'gp', 'tp']
        self.regs += [f't{n}' for n in range(3)]
        self.regs += [f's{n}' for n in range(2)]
        self.regs += [f'a{n}' for n in range(8)]
        self.regs += [f's{n}' for n in range(2, 12)]
        self.regs += [f't{n}' for n in range(3, 7)]
        self.regidx = { r: i for i, r in enumerate(self.regs) }
        self.xregidx = { r: i for i, r in enumerate(self.xregs) }

        self.tmp = 'rfp' # temporary register
        self.branches = ['bge', 'bgeu', 'bgt', 'bgtu', 'blt', 'bltu', 'ble', 'bleu', 'bne', 'bneu', 'beq']

        self.label_gen = iter(range(1000000))

    def unique_label(self):
        return f'rv2fox_{(next(self.label_gen))}'

    @staticmethod
    def ldst_to_size(insn):
        size_map = {
                'lb': 8, 'lh': 16, 'lw': 32,
                'lbu': 8, 'lhu': 16,
                'sb': 8, 'sh': 16, 'sw': 32,
        }
        if insn in size_map:
            return size_map[insn]

    def ldst(self, insn, rval, rbase, offset):
        size = self.ldst_to_size(insn)
        if int(offset, 0) != 0:
            self.e.mov(self.tmp, rbase)
            self.e.insn('add', self.tmp, offset)
            rmem = self.tmp
        else:
            rmem = rbase
        if insn.startswith('s'):
            self.e.store(rmem, rval, size)
        else:
            self.e.load(rmem, rval, size)

    def convert_directive(self, line):
        if line in ['text', 'data']:
            pass # ignore
        elif m := re.fullmatch(r'word\s+([^\s]+)', line):
            self.e.data(*m.groups())
        elif m := re.fullmatch(r'4byte\s+([^\s]+)', line):
            self.e.data(*m.groups())
        elif m := re.fullmatch(r'zero\s+([^\s]+)', line):
            self.e.zero(*m.groups())
        elif m := re.fullmatch(r'byte\s+([^\s]+)', line):
            self.e.zero(*m.groups())
        elif m := re.fullmatch(r'string\s+\"(.*)\"', line):
            self.e.strz(*m.groups())
        elif m := re.fullmatch(r'asciz\s+\"(.*)\"', line):
            self.e.strz(*m.groups())
        elif m := re.match(r'set\s([^,]*),', line):
            self.e.label(*m.groups())

    def reg(self, reg):
        """
        Convert a RISC-V register to a fox32 register
        """
        if reg == 'sp':
            return 'rsp'
        elif reg == 'zero':
            return f'0';
        elif reg in self.regidx:
            return f'r{self.regidx[reg]}'
        elif reg in self.xregidx:
            return f'r{self.xregidx[reg]}'
        else:
            raise Exception(f'Unknown register {reg}')

    @staticmethod
    def convert_arithmetic_op(op):
        if op in ['add', 'addi']: return 'add'
        if op in ['sub', 'subi']: return 'sub'
        if op in ['mul']: return 'mul'
        if op in ['div']: return 'idiv'
        if op in ['divu']: return 'div'
        if op in ['rem']: return 'irem'
        if op in ['remu']: return 'rem'
        if op in ['and', 'andi']: return 'and'
        if op in ['or', 'ori']: return 'or'
        if op in ['xor', 'xori']: return 'xor'
        if op in ['sll', 'slli']: return 'sla'
        if op in ['srl', 'srli']: return 'srl'
        else:
            raise Exception(f'Unknown op {op}')

    @staticmethod
    def convert_condition(op):
        if op in ['blt', 'bltu']: return 'iflt', False
        if op in ['bne', 'bnez']: return 'ifnz', False
        if op in ['beq', 'beqz']: return 'ifz', False
        if op in ['ble', 'bleu']: return 'ifgt', True
        if op in ['bge', 'bgeu']: return 'ifgteq', False
        if op in ['bgt', 'bgtu']: return 'ifgt', False

    def convert_branch(self, insn, rs1, rs2, label):
        cond, reverse = self.convert_condition(insn)
        if reverse:
            rs1, rs2 = rs2, rs1
        self.e.insn('cmp', self.reg(rs1), self.reg(rs2))
        self.e.jump(label, cond)

    def convert_branch_imm(self, insn, rs, label, imm):
        cond, reverse = self.convert_condition(insn)
        if reverse:
            rs1, rs2 = rs2, rs1
        self.e.insn('cmp', self.reg(rs), str(imm))
        self.e.jump(label, cond)

    @staticmethod
    def iter_operands(operands):
        if operands:
            for op in operands.split(','):
                op = op.strip()
                if op != '':
                    yield op

    def convert_insn(self, insn, operands=None):
        tmp = self.tmp
        operands = list(self.iter_operands(operands))

        if insn == 'nop':
            self.e.nop()
        elif insn in ['addi', 'subi', 'div', 'rem', 'andi', 'ori', 'xori', 'slli', 'srli']:
            rd, rs, imm = operands
            op = self.convert_arithmetic_op(insn)
            if rd == rs:
                self.e.insn(op, self.reg(rd), imm)
            else:
                self.e.mov(tmp, self.reg(rs))
                self.e.insn(op, tmp, imm)
                self.e.mov(self.reg(rd), tmp)
        elif insn in ['add', 'sub', 'mul', 'divu', 'remu', 'and', 'or', 'nor', 'xor', 'sll', 'srl']:
            rd, rs1, rs2 = operands
            if rd == rs1:
                op = self.convert_arithmetic_op(insn)
                self.e.insn(op, self.reg(rd), self.reg(rs2))
            else:
                op = self.convert_arithmetic_op(insn)
                self.e.mov(tmp, self.reg(rs1))
                self.e.insn(op, tmp, self.reg(rs2))
                self.e.mov(self.reg(rd), tmp)
        elif insn in ['sw', 'lw', 'sh', 'lh', 'lhu', 'sb', 'lb', 'lbu']:
            rval, address = operands
            offset, rbase = re.fullmatch(r'(-?[0-9]+)\(([a-z0-9]+)\)', address).groups()
            self.ldst(insn, self.reg(rval), self.reg(rbase), offset)
        elif insn in ['j']:
            self.e.jump(*operands)
        elif insn in ['jr']:
            self.e.jump_reg(self.reg(*operands))
        elif insn in ['ret']:
            self.e.jump_reg(self.reg('ra'))
        elif insn in ['lla']:
            rd, label = operands
            self.e.rta(self.reg(rd), label)
        elif insn in ['mv']:
            rd, rs = operands
            self.e.mov(self.reg(rd), self.reg(rs))
        elif insn in ['call']:
            label = self.unique_label()
            self.e.mov(self.reg('ra'), label)
            self.e.jump(*operands)
            self.e.label(label)
        elif insn in self.branches:
            self.convert_branch(insn, *operands)
        elif insn in ['bnez']:
            self.convert_branch_imm(insn, *operands, 0)
        elif insn in ['li']:
            rd, imm = operands
            self.e.mov(self.reg(rd), imm)
        elif insn in ['lui']:
            rd, imm = operands
            self.e.mov(self.reg(rd), str(int(imm, 0) << 12))
        elif insn == 'wfi':
            self.e.halt()
        else:
            raise Exception(f'Unknown instruction {insn} {operands}')

    def convert(self):
        self.e.comment(f'Generated with rv2fox')

        self.e.include(self.start_file)
        for file in self.additional_includes:
            self.e.include(file)

        for i, line in enumerate(self.input_file):
            num = i + 1
            line = line.strip()

            if self.debug:
                self.e.write(f'\t; {line}')

            if line == '' or re.fullmatch(r'#.*', line):
                pass
            elif re.fullmatch(r'\.[a-z0-9]+.*', line):
                self.convert_directive(line[1:])
            elif m := re.fullmatch(r'([a-zA-Z0-9_.]+):', line):
                self.e.label(*m.groups())
            elif m := re.fullmatch(r'([a-z]+)', line) or \
                      re.fullmatch(r'([a-z]+)\s+([^\s].*)', line):
                self.convert_insn(*m.groups())
            else:
                raise Exception(f'Unknown line: {line}')

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Convert RISC-V assembly to fox32 assembly')
    parser.add_argument('input', help='RISC-V input')
    parser.add_argument('-o', '--output', help='fox32 output')
    parser.add_argument('--debug', help='produce verbose output to allow debugging this tool', action='store_true')
    args = parser.parse_args()
    c = Converter(args.input, args.output, debug=args.debug)
    c.convert()