<commit_before>from itertools import zip_longest
from typing import Dict, List, Optional, Set, Tuple

from methodtools import lru_cache  # type: ignore

from gilgamesh.disassembly.parser import (
    EDITABLE_TOKEN_TYPES,
    EDITABLE_TOKENS,
    EQUIVALENT_TOKENS,
    HEADER_TOKENS,
    Parser,
    Token,
)
from gilgamesh.disassembly.parser import TokenType as T
from gilgamesh.disassembly.renames import apply_local_renames
from gilgamesh.errors import ParserError
from gilgamesh.snes.hw_registers import hw_registers
from gilgamesh.snes.instruction import Instruction, StackManipulation
from gilgamesh.snes.opcodes import Op
from gilgamesh.snes.state import StateChange
from gilgamesh.subroutine import Subroutine


class Disassembly:
    SEPARATOR_LINE = ";" + ("-" * 40)

    def __init__(
        self,
        subroutine: Subroutine,
        preview_excluded_labels: Optional[Set[str]] = None,
    ):
        self.log = subroutine.log
        self.subroutine = subroutine
        self.preview_excluded_labels = preview_excluded_labels or set()
        self.base_line_n = 1

    def get_html(self) -> str:
        return self._get_text(html=True)[1]

    def get_instruction_html(self, instruction: Instruction, verbose=False) -> str:
        tokens = self._instruction_to_tokens(instruction, verbose)
        return self._instruction_tokens_to_text(tokens, html=True)[1]

    def _get_text(self, html=False) -> Tuple[int, str, List[List[Token]]]:
        n_lines, s, tokens = 0, [], []

        for instruction in self.subroutine.instructions.values():
            instr_tokens = self._instruction_to_tokens(instruction)
            tokens.append(instr_tokens)

            instr_lines, text = self._instruction_tokens_to_text(instr_tokens, html)
            n_lines += instr_lines
            s.append(text)

        return n_lines, "".join(s), tokens

    def _instruction_to_tokens(self, instr: Instruction, verbose=True) -> List[Token]:
        """Convert an instruction into a list of token which describes it."""
        tokens = []

        def add(*args) -> None:
            tokens.append(Token(*args))

        def add_line(*args) -> None:
            add(*args)
            tokens.append(Token(T.NEWLINE))

        # Stack manipulation.
        if verbose and instr.stack_manipulation != StackManipulation.NONE:
            if instr.stack_manipulation == StackManipulation.CAUSES_UNKNOWN_STATE:
                add_line(T.FATAL_STACK_MANIPULATION_HEADER)
            else:
                add_line(T.STACK_MANIPULATION_HEADER)

        # Label.
        label = self.log.get_label(instr.pc, instr.subroutine_pc)
        if label:
            if instr.pc in self.log.jump_table_targets:
                add_line(T.JUMP_TABLE_LABEL, label)
            else:
                add_line(T.LABEL, label)
        # Operation + Operand.
        tokens.append(Token(T.OPERATION, instr.name))
        if instr.argument_alias:
            if instr.argument_alias in hw_registers:
                add(T.HW_REGISTER, instr.argument_alias)
            elif instr.argument_alias in self.preview_excluded_labels:
                add(T.HIGHLIGHTED_OPERAND_LABEL, instr.argument_alias)
            elif instr.pc in self.log.jump_table_targets:
                add(T.JUMP_TABLE_OPERAND_LABEL, instr.argument_alias)
            else:
                add(T.OPERAND_LABEL, instr.argument_alias)
        else:
            add(T.OPERAND, instr.argument_string)
        # PC + Comment.
        add(T.PC, "${:06X}".format(instr.pc))
        comment = self.log.comments.get(instr.pc, self._get_default_comment(instr))
        add_line(T.COMMENT, comment)

        # Don't show extra information on state in non-verbose mode.
        if not verbose:
            return tokens

        # Jump table.
        if instr.is_jump_table:
            default_typ = (
                T.JUMP_TABLE_COMPLETE_ENTRY
                if instr.pc in self.log.complete_jump_tables
                else T.JUMP_TABLE_ENTRY
            )
            add_line(T.JUMP_TABLE_HEADER)
            for x, target in self.log.sorted_jump_table(instr.pc):
                target_sub = self.log.subroutines.get(target, None)
                if target_sub is not None and target_sub.has_unknown_return_state:
                    typ = T.JUMP_TABLE_UNKNOWN_ENTRY
                else:
                    typ = default_typ
                add_line(
                    typ,
                    "x={}  ->  {}".format(
                        "????" if x is None else f"{x:04X}",
                        self.log.get_label(target, self.subroutine.pc),
                    ),
                )
            add_line(T.SEPARATOR_LINE)

        # Asserted or unknown state.
        state_change, assertion_type = self._get_unknown_state(instr)
        if assertion_type != "none" or state_change == "unknown":
            suggestion = self.log.suggest_assertion(instr)
            if state_change == "unknown":
                add_line(T.UNKNOWN_STATE_HEADER)
            else:
                add_line(T.ASSERTED_STATE_HEADER)
            add_line(T.LAST_KNOWN_STATE, str(instr.state))
            add_line(T.LAST_KNOWN_STATE_CHANGE, str(instr.state_change_before))
            if state_change == "unknown":
                add_line(
                    T.UNKNOWN_REASON, str(instr.state_change_after.unknown_reason_str)
                )
            add_line(T.SEPARATOR_LINE)
            if suggestion:
                add_line(T.SUGGESTED_ASSERTION_TYPE, suggestion[0])
                add_line(T.SUGGESTED_ASSERTION, str(suggestion[1]))
            else:
                add_line(T.ASSERTION_TYPE, assertion_type)
                add_line(T.ASSERTION, state_change)
            add_line(T.SEPARATOR_LINE)
        # Normal return state.
        elif instr.is_return and not instr.is_jump_table:
            add_line(T.KNOWN_STATE_HEADER)
            add_line(T.KNOWN_STATE_CHANGE, state_change)
            add_line(T.SEPARATOR_LINE)

        return tokens

    def _text_to_tokens(self, text: str) -> List[List[Token]]:
        """Parse a subroutines's disassembly into a list of lists of tokens."""

        p = Parser(text, self.base_line_n)
        while p.line_idx < len(p.lines):
            # Empty line.
            if not p.line:
                p.add_line()

            # Label line.
            elif len(p.words) == 1 and p.line[-1] == ":":
                p.new_instruction()
                p.add_line(T.LABEL, p.line[:-1])

            # Stack manipulation.
            elif p.maybe_match_line(self.string(T.STACK_MANIPULATION_HEADER)):
                p.new_instruction()
                p.add_line(T.STACK_MANIPULATION_HEADER)

            # Jump table.
            elif p.maybe_match_line(self.string(T.JUMP_TABLE_HEADER)):
                p.add_line(T.JUMP_TABLE_HEADER)
                while not p.maybe_match_line(self.SEPARATOR_LINE):
                    p.add_line(T.JUMP_TABLE_ENTRY)
                p.add_line(T.SEPARATOR_LINE)

            # Known return state.<commit_msg>Closes #128<commit_after>from itertools import zip_longest
from typing import Dict, List, Optional, Set, Tuple

from methodtools import lru_cache  # type: ignore

from gilgamesh.disassembly.parser import (
    EDITABLE_TOKEN_TYPES,
    EDITABLE_TOKENS,
    EQUIVALENT_TOKENS,
    HEADER_TOKENS,
    Parser,
    Token,
)
from gilgamesh.disassembly.parser import TokenType as T
from gilgamesh.disassembly.renames import apply_local_renames
from gilgamesh.errors import ParserError
from gilgamesh.snes.hw_registers import hw_registers
from gilgamesh.snes.instruction import Instruction, StackManipulation
from gilgamesh.snes.opcodes import Op
from gilgamesh.snes.state import StateChange
from gilgamesh.subroutine import Subroutine


class Disassembly:
    SEPARATOR_LINE = ";" + ("-" * 40)

    def __init__(
        self,
        subroutine: Subroutine,
        preview_excluded_labels: Optional[Set[str]] = None,
    ):
        self.log = subroutine.log
        self.subroutine = subroutine
        self.preview_excluded_labels = preview_excluded_labels or set()
        self.base_line_n = 1

    def get_html(self) -> str:
        return self._get_text(html=True)[1]

    def get_instruction_html(self, instruction: Instruction, verbose=False) -> str:
        tokens = self._instruction_to_tokens(instruction, verbose)
        return self._instruction_tokens_to_text(tokens, html=True)[1]

    def _get_text(self, html=False) -> Tuple[int, str, List[List[Token]]]:
        n_lines, s, tokens = 0, [], []

        for instruction in self.subroutine.instructions.values():
            instr_tokens = self._instruction_to_tokens(instruction)
            tokens.append(instr_tokens)

            instr_lines, text = self._instruction_tokens_to_text(instr_tokens, html)
            n_lines += instr_lines
            s.append(text)

        return n_lines, "".join(s), tokens

    def _instruction_to_tokens(self, instr: Instruction, verbose=True) -> List[Token]:
        """Convert an instruction into a list of token which describes it."""
        tokens = []

        def add(*args) -> None:
            tokens.append(Token(*args))

        def add_line(*args) -> None:
            add(*args)
            tokens.append(Token(T.NEWLINE))

        # Label.
        label = self.log.get_label(instr.pc, instr.subroutine_pc)
        if label:
            if instr.pc in self.log.jump_table_targets:
                add_line(T.JUMP_TABLE_LABEL, label)
            else:
                add_line(T.LABEL, label)

        # Stack manipulation.
        if verbose and instr.stack_manipulation != StackManipulation.NONE:
            if instr.stack_manipulation == StackManipulation.CAUSES_UNKNOWN_STATE:
                add_line(T.FATAL_STACK_MANIPULATION_HEADER)
            else:
                add_line(T.STACK_MANIPULATION_HEADER)

        # Operation + Operand.
        tokens.append(Token(T.OPERATION, instr.name))
        if instr.argument_alias:
            if instr.argument_alias in hw_registers:
                add(T.HW_REGISTER, instr.argument_alias)
            elif instr.argument_alias in self.preview_excluded_labels:
                add(T.HIGHLIGHTED_OPERAND_LABEL, instr.argument_alias)
            elif instr.pc in self.log.jump_table_targets:
                add(T.JUMP_TABLE_OPERAND_LABEL, instr.argument_alias)
            else:
                add(T.OPERAND_LABEL, instr.argument_alias)
        else:
            add(T.OPERAND, instr.argument_string)
        # PC + Comment.
        add(T.PC, "${:06X}".format(instr.pc))
        comment = self.log.comments.get(instr.pc, self._get_default_comment(instr))
        add_line(T.COMMENT, comment)

        # Don't show extra information on state in non-verbose mode.
        if not verbose:
            return tokens

        # Jump table.
        if instr.is_jump_table:
            default_typ = (
                T.JUMP_TABLE_COMPLETE_ENTRY
                if instr.pc in self.log.complete_jump_tables
                else T.JUMP_TABLE_ENTRY
            )
            add_line(T.JUMP_TABLE_HEADER)
            for x, target in self.log.sorted_jump_table(instr.pc):
                target_sub = self.log.subroutines.get(target, None)
                if target_sub is not None and target_sub.has_unknown_return_state:
                    typ = T.JUMP_TABLE_UNKNOWN_ENTRY
                else:
                    typ = default_typ
                add_line(
                    typ,
                    "x={}  ->  {}".format(
                        "????" if x is None else f"{x:04X}",
                        self.log.get_label(target, self.subroutine.pc),
                    ),
                )
            add_line(T.SEPARATOR_LINE)

        # Asserted or unknown state.
        state_change, assertion_type = self._get_unknown_state(instr)
        if assertion_type != "none" or state_change == "unknown":
            suggestion = self.log.suggest_assertion(instr)
            if state_change == "unknown":
                add_line(T.UNKNOWN_STATE_HEADER)
            else:
                add_line(T.ASSERTED_STATE_HEADER)
            add_line(T.LAST_KNOWN_STATE, str(instr.state))
            add_line(T.LAST_KNOWN_STATE_CHANGE, str(instr.state_change_before))
            if state_change == "unknown":
                add_line(
                    T.UNKNOWN_REASON, str(instr.state_change_after.unknown_reason_str)
                )
            add_line(T.SEPARATOR_LINE)
            if suggestion:
                add_line(T.SUGGESTED_ASSERTION_TYPE, suggestion[0])
                add_line(T.SUGGESTED_ASSERTION, str(suggestion[1]))
            else:
                add_line(T.ASSERTION_TYPE, assertion_type)
                add_line(T.ASSERTION, state_change)
            add_line(T.SEPARATOR_LINE)
        # Normal return state.
        elif instr.is_return and not instr.is_jump_table:
            add_line(T.KNOWN_STATE_HEADER)
            add_line(T.KNOWN_STATE_CHANGE, state_change)
            add_line(T.SEPARATOR_LINE)

        return tokens

    def _text_to_tokens(self, text: str) -> List[List[Token]]:
        """Parse a subroutines's disassembly into a list of lists of tokens."""

        p = Parser(text, self.base_line_n)
        while p.line_idx < len(p.lines):
            # Empty line.
            if not p.line:
                p.add_line()

            # Label line.
            elif len(p.words) == 1 and p.line[-1] == ":":
                p.new_instruction()
                p.add_line(T.LABEL, p.line[:-1])

            # Stack manipulation.
            elif p.maybe_match_line(self.string(T.STACK_MANIPULATION_HEADER)):
                p.maybe_new_instruction()
                p.add_line(T.STACK_MANIPULATION_HEADER)

            # Jump table.
            elif p.maybe_match_line(self.string(T.JUMP_TABLE_HEADER)):
                p.add_line(T.JUMP_TABLE_HEADER)
                while not p.maybe_match_line(self.SEPARATOR_LINE):
                    p.add_line(T.JUMP_TABLE_ENTRY)
                p.add_line(T.SEPARATOR_LINE)

            # Known return state.
{'fraction_non_alphanumeric': 0.0887708111937655, 'fraction_numerical': 0.0023379383634431457, 'mean_word_length': 2.0468379020073386, 'pattern_counts': {'":': 10, '<': 5, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4652116', 'n_tokens_mistral': 4458, 'n_tokens_neox': 4206, 'n_words': 796}