my_code_3

import pyparsing as pp
import re
import struct
from typing import List, Dict, Union

debug_mode = True

class PDP11Assembler:
    def __init__(self):
        # Инициализация парсера pyparsing (как в вашем коде)
        self.command_name = pp.Optional(pp.Suppress(' ')) + pp.Word(pp.alphas)
        self.argument_name = pp.Optional(pp.Suppress(' ')) + pp.Word(pp.printables, exclude_chars=';,')
        self.arguments = pp.Optional(pp.Suppress(' ')) + pp.Optional(self.argument_name) + pp.ZeroOrMore(pp.Suppress(',') + self.argument_name)
        self.label = pp.Combine(pp.Optional(pp.Suppress(' ')) + pp.Word(pp.alphanums) + pp.Optional(pp.Suppress(' ')) + pp.Literal(':'))
        self.comment = pp.Optional(pp.Suppress(' ')) + pp.Suppress(';') + pp.Optional(pp.Suppress(' ')) + pp.restOfLine()
        self.rule = pp.Optional(self.label) + pp.Optional(self.command_name) + pp.Optional(
            self.arguments) + pp.Optional(self.comment)

        # Таблица команд (расширенная версия)
        self.command_table = {
            'halt': {'opcode': 0o000000, 'args': [], 'handler': self._handle_no_args},
            'mov': {'opcode': 0o010000, 'args': ['ss', 'dd'], 'handler': self._handle_ss_dd},
            'add': {'opcode': 0o060000, 'args': ['ss', 'dd'], 'handler': self._handle_ss_dd},
            'sub': {'opcode': 0o160000, 'args': ['ss', 'dd'], 'handler': self._handle_ss_dd},
            'cmp': {'opcode': 0o020000, 'args': ['ss', 'dd'], 'handler': self._handle_ss_dd},
            'br': {'opcode': 0o000400, 'args': ['xx'], 'handler': self._handle_branch},
            'bne': {'opcode': 0o001000, 'args': ['xx'], 'handler': self._handle_branch},
        }

        self.labels = {}
        self.current_address = 0o1000
        self.output_blocks = []

    def parse_line(self, line: str) -> List[str]:
        """Аналог вашей функции parse с улучшениями"""
        #if debug_mode:
        #    print("line:", line)
        line = line.strip()
        if not line:
            return []

        if line.startswith('.'):
            if line.startswith('.='):
                return ['origin', line[2:].strip()]
            return []

        try:
            return self.rule.parseString(line).asList()
        except pp.ParseException:
            return []

    def _handle_no_args(self, cmd: Dict, args: List[str], words: List[Union[int, str]]):
        """Обработка команд без аргументов"""
        words.append(cmd['opcode'])

    def _handle_branch(self, cmd: Dict, args: List[str], words: List[Union[int, str]]):
        """Обработка команд ветвления"""
        words.append(cmd['opcode'])
        words.append(args[0])  # Сохраняем метку для разрешения позже

    def _handle_ss_dd(self, cmd: Dict, args: List[str], words: List[Union[int, str]]):
        """Обработка двухоперандных команд"""
        src = self._parse_operand(args[0])
        dst = self._parse_operand(args[1])

        instruction = cmd['opcode'] | (src['mode'] << 6) | (src['reg'] << 3) | dst['mode'] | dst['reg']
        words.append(instruction)

        if src['mode'] == 0o27:  # Немедленная адресация
            words.append(src['value'])

    def _parse_operand(self, op: str) -> Dict[str, Union[int, str]]:
        """Улучшенный парсер операндов"""
        if not op:
            return {'mode': 0, 'reg': 0, 'value': 0}

        # Регистровая адресация (R0-R7)
        if re.match(r'^[rR][0-7]$', op):
            return {'mode': 0, 'reg': int(op[1]), 'value': 0}

        # Немедленная адресация (#n)
        if op.startswith('#'):
            return {'mode': 0o27, 'reg': 0, 'value': int(op[1:])}

        return {'mode': 0, 'reg': 0, 'value': 0}

    def cmd_to_raw_machine_code(self, parsed: List[str]) -> List[Union[int, str]]:
        """Аналог вашей функции, но с улучшенной структурой"""
        if not parsed:
            return []

        """if isinstance(parsed[0], list) and len(parsed[0]) == 2 and parsed[0][1] == ':':
            # Обработка метки
            label = parsed[0][0]
            self.labels[label] = self.current_address
            parsed = parsed[1:]

        if not parsed:
            return []"""

        if parsed[0] == 'origin':
            self.current_address = int(parsed[1], 8)
            return []

        cmd_name = parsed[0].lower()
        if cmd_name not in self.command_table:
            return []

        cmd = self.command_table[cmd_name]
        args = parsed[1:]

        if len(args) != len(cmd['args']):
            return []

        words = []
        cmd['handler'](cmd, args, words)
        return words


    def recycle_lines(self, lines: List[str]):
        """Второй проход: генерация машинного кода"""
        self.current_address = 0o1000
        current_block = {'start': 0o1000, 'data': []}

        if debug_mode:
            print()
            print("recycle_lines")
            print()

        for line in lines:

            if debug_mode:
                print("line:", line)

            parsed = self.parse_line(line)

            if debug_mode:
                print("parsed:", parsed)

            if not parsed:
                continue

            if parsed[0] == 'origin':
                if current_block['data']:
                    self.output_blocks.append(current_block)
                addr = int(parsed[1], 8)
                current_block = {'start': addr, 'data': []}
                self.current_address = addr
                continue

            """if isinstance(parsed[0], list) and len(parsed[0]) == 2 and parsed[0][1] == ':':
                parsed = parsed[1:]
                if not parsed:
                    continue"""

            words = self.cmd_to_raw_machine_code(parsed)
            for item in words:
                if isinstance(item, str):  # Метка для ветвления
                    offset = (self.labels[item] - self.current_address - 2) // 2
                    current_block['data'][-1] |= (offset & 0o377)
                else:
                    current_block['data'].append(item)
                    self.current_address += 2

        if current_block['data']:
            self.output_blocks.append(current_block)

    def write_s_record_format(self, filename: str):
        """Запись результата в требуемом формате"""
        if not self.output_blocks:
            return

        block = self.output_blocks[0]
        data_bytes = bytearray()

        for word in block['data']:
            # Правильный порядок байтов (как в вашем примере)
            data_bytes.append(word & 0xFF)
            data_bytes.append((word >> 8) & 0xFF)

        with open(filename, 'w') as f:
            # Первая строка: адрес и размер
            f.write(f"0200 {len(data_bytes):04x}\n")

            # Байты машинного кода
            for byte in data_bytes:
                f.write(f"{byte:02x}\n")

    def assemble(self, input_filename: str, output_filename: str):
        """Основной метод ассемблирования"""
        with open(input_filename, 'r') as f:
            lines = f.readlines()

        self.recycle_lines(lines)

        # Запись результата
        self.write_s_record_format(output_filename)



assembler = PDP11Assembler()

test_code = """
    .=1000
    mov #2, r0
    mov #3, r1
    add r0, r1
    halt
"""

with open("test.asm", "w") as f:
    f.write(test_code)

assembler.assemble("asm_code.txt", "output.txt")
print("Ассемблирование завершено. Результат в output.txt")