<commit_before># Copyright lowRISC contributors.
# Licensed under the Apache License, Version 2.0, see LICENSE for details.
# SPDX-License-Identifier: Apache-2.0

import random
from typing import List, Optional, Tuple

from shared.insn_yaml import Insn, InsnsFile
from shared.lsu_desc import LSUDesc
from shared.operand import ImmOperandType, OptionOperandType, RegOperandType

from ..config import Config
from ..program import ProgInsn, Program
from ..model import Model
from ..snippet import ProgSnippet, Snippet
from ..snippet_gen import GenCont, GenRet, SimpleGenRet, SnippetGen


class StraightLineInsn(SnippetGen):
    '''A super-simple snippet consisting of a single instruction'''
    def __init__(self, cfg: Config, insns_file: InsnsFile) -> None:
        super().__init__()

        # Find all the straight line, non-pseudo instructions in insns_file
        self.insns = []
        self.weights = []

        seen_insns = set()
        for insn in insns_file.insns:
            # Skip pseudo-ops
            if insn.python_pseudo_op or insn.literal_pseudo_op:
                continue

            seen_insns.add(insn.mnemonic)

            # Skip anything that isn't straight-line
            if not insn.straight_line:
                continue

            weight = cfg.insn_weights.get(insn.mnemonic)
            if weight > 0:
                self.insns.append(insn)
                self.weights.append(weight)

        # Check that the config's insn-weights dictionary didn't have any
        # non-existent instructions. Note that we even add jumps to seen_insns:
        # although weighting those won't make any difference here, this is the
        # code that makes sense of what's allowed.
        missing_insns = set(cfg.insn_weights.values.keys()) - seen_insns
        if missing_insns:
            raise ValueError('Config at {} defines weights for '
                             'non-existent instructions: {}'
                             .format(cfg.path,
                                     ', '.join(sorted(missing_insns))))

        # Check that at least one instruction has a positive weight
        assert len(self.insns) == len(self.weights)
        if not self.insns:
            raise ValueError('Config at {} defines a zero weight '
                             'for all instructions.'
                             .format(cfg.path))

    def gen(self,
            cont: GenCont,
            model: Model,
            program: Program) -> Optional[GenRet]:
        return SnippetGen._unsimple_genret(self._gen(model, program))

    def gen_some(self,
                 count: int,
                 model: Model,
                 program: Program) -> Optional[SimpleGenRet]:
        '''Generate a block of count straight-line instructions'''
        assert 0 < count

        # Take a copy of model, so that we fail atomically
        model = model.copy()

        children = []  # type: List[Snippet]
        for i in range(count):
            gen_ret = self._gen(model, program)
            if gen_ret is None:
                return None

            snippet, model = gen_ret
            children.append(snippet)

        return (Snippet.merge_list(children), model)

    def _gen(self,
             model: Model,
             program: Program) -> Optional[SimpleGenRet]:
        # Return None if this is the last instruction in the current gap
        # because we need to either jump or do an ECALL to avoid getting stuck.
        #
        # Note that we could do this by defining pick_weight, but we don't
        # expect it to happen very often so it's probably best (and cheaper)
        # just to disable ourselves on the rare occasions when it does.
        if program.get_insn_space_at(model.pc) <= 1:
            return None

        weights = self.weights
        prog_insn = None
        while prog_insn is None:
            idx = random.choices(range(len(self.insns)), weights=weights)[0]
            # At least one weight should be positive. This is guaranteed so
            # long as fill_insn doesn't fail for absolutely everything. We know
            # that doesn't happen, because we have some instructions (such as
            # addi) where it will not fail.
            assert weights[idx] > 0

            # Try to fill out the instruction. On failure, clear the weight for
            # this index and go around again. We take the copy here, rather
            # than outside the loop, because we don't expect this to happen
            # very often.
            prog_insn = self.fill_insn(self.insns[idx], model)
            if prog_insn is None:
                weights = self.weights.copy()
                weights[idx] = 0
                continue

        # Success! We have generated an instruction. Put it in a snippet and
        # add that to the program
        snippet = ProgSnippet(model.pc, [prog_insn])
        snippet.insert_into_program(program)

        # Then update the model with the instruction and update the model PC
        model.update_for_insn(prog_insn)
        model.pc += 4

        return (snippet, model)

    def fill_insn(self, insn: Insn, model: Model) -> Optional[ProgInsn]:
        '''Try to fill out an instruction

        This might fail if, for example, the model doesn't have enough
        registers with architectural values. In that case, return None.

        Note that there are some instructions that will never return None. For
        example, addi, which can always expand to "addi x0, x0, 0" (also known
        as nop).

        '''

        # Special-case BN load/store instructions by mnemonic. These use
        # complicated indirect addressing, so it's probably more sensible to
        # give them special code.
        if insn.mnemonic in ['bn.lid', 'bn.sid']:
            return self._fill_bn_xid(insn, model)
        if insn.mnemonic == 'bn.movr':
            return self._fill_bn_movr(insn, model)

        # If this is not an LSU operation, or it is an LSU operation that
        # operates on CSR/WSRs, we can pick operands independently.
        if insn.lsu is None:
            return self._fill_non_lsu_insn(insn, model)

        return self._fill_lsu_insn(insn, model)

    def _fill_non_lsu_insn(self,
                           insn: Insn,
                           model: Model) -> Optional[ProgInsn]:
        '''Fill out an instruction with no LSU component'''
        assert insn.lsu is None
        # For each operand, pick a value that's allowed by the model (i.e.
        # one that won't trigger any undefined behaviour)
        op_vals = []
        for operand in insn.operands:
            op_val = model.pick_operand_value(operand.op_type)
            if op_val is None:
                return None

            op_vals.append(op_val)

        assert len(op_vals) == len(insn.operands)
        return ProgInsn(insn, op_vals, None)

    def _pick_gpr_for_indirect_wdr(self,
                                   known_regs: List[Tuple[int, int]],
                                   model: Model,
                                   wdr_is_src: bool) -> Optional[int]:
        '''Return index of a GPR pointing to a WDR with architectural value

        Here, the GPR value must be known and must be at most 31. If wdr_is_src
        then the WDR that it points to must also have an architectural value.

        '''
        arch_wdrs = set(model.regs_with_architectural_vals('wdr'))
        valid_gprs = []
        for gpr, gpr_val in known_regs:
            if gpr_val is None or gpr_val > 31:
                continue
            if wdr_is_src and gpr_val not in arch_wdrs:
                continue

            valid_gprs.append(gpr)

        if not valid_gprs:
            return None

        return random.choices(valid_gprs)[0]

    def _pick_inc_vals(self,
                       idx0: int,
                       idx1: int,
                       model: Model) -> Tuple[int, int]:
        '''Pick two values in 0, 1 that aren't both 1

        These are appropriate to use as the increment flags for
        BN.LID/BN.SID/BN.MOVR. idx0 and idx1 are the indices of the GPRs in
        question.

        '''
        options = [(0, 0), (1, 0), (0, 1)]
        wt10 = 0.0 if model.is_const('gpr', idx0) else 1.0
        wt01 = 0.0 if model.is_const('gpr', idx1) else 1.0
        return random.choices(options, weights=[1.0, wt10, wt01])[0]

    def _fill_bn_xid(self, insn: Insn, model: Model) -> Optional[ProgInsn]:
        '''Fill out a BN.LID or BN.SID instruction'''<commit_msg>[otbn,dv] Don't write to const registers with BN.LID in RIG<commit_after># Copyright lowRISC contributors.
# Licensed under the Apache License, Version 2.0, see LICENSE for details.
# SPDX-License-Identifier: Apache-2.0

import random
from typing import List, Optional, Tuple

from shared.insn_yaml import Insn, InsnsFile
from shared.lsu_desc import LSUDesc
from shared.operand import ImmOperandType, OptionOperandType, RegOperandType

from ..config import Config
from ..program import ProgInsn, Program
from ..model import Model
from ..snippet import ProgSnippet, Snippet
from ..snippet_gen import GenCont, GenRet, SimpleGenRet, SnippetGen


class StraightLineInsn(SnippetGen):
    '''A super-simple snippet consisting of a single instruction'''
    def __init__(self, cfg: Config, insns_file: InsnsFile) -> None:
        super().__init__()

        # Find all the straight line, non-pseudo instructions in insns_file
        self.insns = []
        self.weights = []

        seen_insns = set()
        for insn in insns_file.insns:
            # Skip pseudo-ops
            if insn.python_pseudo_op or insn.literal_pseudo_op:
                continue

            seen_insns.add(insn.mnemonic)

            # Skip anything that isn't straight-line
            if not insn.straight_line:
                continue

            weight = cfg.insn_weights.get(insn.mnemonic)
            if weight > 0:
                self.insns.append(insn)
                self.weights.append(weight)

        # Check that the config's insn-weights dictionary didn't have any
        # non-existent instructions. Note that we even add jumps to seen_insns:
        # although weighting those won't make any difference here, this is the
        # code that makes sense of what's allowed.
        missing_insns = set(cfg.insn_weights.values.keys()) - seen_insns
        if missing_insns:
            raise ValueError('Config at {} defines weights for '
                             'non-existent instructions: {}'
                             .format(cfg.path,
                                     ', '.join(sorted(missing_insns))))

        # Check that at least one instruction has a positive weight
        assert len(self.insns) == len(self.weights)
        if not self.insns:
            raise ValueError('Config at {} defines a zero weight '
                             'for all instructions.'
                             .format(cfg.path))

    def gen(self,
            cont: GenCont,
            model: Model,
            program: Program) -> Optional[GenRet]:
        return SnippetGen._unsimple_genret(self._gen(model, program))

    def gen_some(self,
                 count: int,
                 model: Model,
                 program: Program) -> Optional[SimpleGenRet]:
        '''Generate a block of count straight-line instructions'''
        assert 0 < count

        # Take a copy of model, so that we fail atomically
        model = model.copy()

        children = []  # type: List[Snippet]
        for i in range(count):
            gen_ret = self._gen(model, program)
            if gen_ret is None:
                return None

            snippet, model = gen_ret
            children.append(snippet)

        return (Snippet.merge_list(children), model)

    def _gen(self,
             model: Model,
             program: Program) -> Optional[SimpleGenRet]:
        # Return None if this is the last instruction in the current gap
        # because we need to either jump or do an ECALL to avoid getting stuck.
        #
        # Note that we could do this by defining pick_weight, but we don't
        # expect it to happen very often so it's probably best (and cheaper)
        # just to disable ourselves on the rare occasions when it does.
        if program.get_insn_space_at(model.pc) <= 1:
            return None

        weights = self.weights
        prog_insn = None
        while prog_insn is None:
            idx = random.choices(range(len(self.insns)), weights=weights)[0]
            # At least one weight should be positive. This is guaranteed so
            # long as fill_insn doesn't fail for absolutely everything. We know
            # that doesn't happen, because we have some instructions (such as
            # addi) where it will not fail.
            assert weights[idx] > 0

            # Try to fill out the instruction. On failure, clear the weight for
            # this index and go around again. We take the copy here, rather
            # than outside the loop, because we don't expect this to happen
            # very often.
            prog_insn = self.fill_insn(self.insns[idx], model)
            if prog_insn is None:
                weights = self.weights.copy()
                weights[idx] = 0
                continue

        # Success! We have generated an instruction. Put it in a snippet and
        # add that to the program
        snippet = ProgSnippet(model.pc, [prog_insn])
        snippet.insert_into_program(program)

        # Then update the model with the instruction and update the model PC
        model.update_for_insn(prog_insn)
        model.pc += 4

        return (snippet, model)

    def fill_insn(self, insn: Insn, model: Model) -> Optional[ProgInsn]:
        '''Try to fill out an instruction

        This might fail if, for example, the model doesn't have enough
        registers with architectural values. In that case, return None.

        Note that there are some instructions that will never return None. For
        example, addi, which can always expand to "addi x0, x0, 0" (also known
        as nop).

        '''

        # Special-case BN load/store instructions by mnemonic. These use
        # complicated indirect addressing, so it's probably more sensible to
        # give them special code.
        if insn.mnemonic in ['bn.lid', 'bn.sid']:
            return self._fill_bn_xid(insn, model)
        if insn.mnemonic == 'bn.movr':
            return self._fill_bn_movr(insn, model)

        # If this is not an LSU operation, or it is an LSU operation that
        # operates on CSR/WSRs, we can pick operands independently.
        if insn.lsu is None:
            return self._fill_non_lsu_insn(insn, model)

        return self._fill_lsu_insn(insn, model)

    def _fill_non_lsu_insn(self,
                           insn: Insn,
                           model: Model) -> Optional[ProgInsn]:
        '''Fill out an instruction with no LSU component'''
        assert insn.lsu is None
        # For each operand, pick a value that's allowed by the model (i.e.
        # one that won't trigger any undefined behaviour)
        op_vals = []
        for operand in insn.operands:
            op_val = model.pick_operand_value(operand.op_type)
            if op_val is None:
                return None

            op_vals.append(op_val)

        assert len(op_vals) == len(insn.operands)
        return ProgInsn(insn, op_vals, None)

    def _pick_gpr_for_indirect_wdr(self,
                                   known_regs: List[Tuple[int, int]],
                                   model: Model,
                                   wdr_is_src: bool) -> Optional[int]:
        '''Return index of a GPR pointing to a WDR with architectural value

        Here, the GPR value must be known and must be at most 31. If wdr_is_src
        then the WDR that it points to must also have an architectural value.
        Otherwise, the WDR is a destination and must not be marked const.

        '''
        wdr_is_dest = not wdr_is_src
        arch_wdrs = set(model.regs_with_architectural_vals('wdr'))
        valid_gprs = []
        for gpr, gpr_val in known_regs:
            if gpr_val is None or gpr_val > 31:
                continue
            if wdr_is_src and gpr_val not in arch_wdrs:
                continue
            if wdr_is_dest and model.is_const('wdr', gpr_val):
                continue

            valid_gprs.append(gpr)

        if not valid_gprs:
            return None

        return random.choices(valid_gprs)[0]

    def _pick_inc_vals(self,
                       idx0: int,
                       idx1: int,
                       model: Model) -> Tuple[int, int]:
        '''Pick two values in 0, 1 that aren't both 1

        These are appropriate to use as the increment flags for
        BN.LID/BN.SID/BN.MOVR. idx0 and idx1 are the indices of the GPRs in
        question.

        '''
        options = [(0, 0), (1, 0), (0, 1)]
        wt10 = 0.0 if model.is_const('gpr', idx0) else 1.0
        wt01 = 0.0 if model.is_const('gpr', idx1) else 1.0
        return random.choices(options, weights=[1.0, wt10, wt01])[0]

    def _fill_bn_xid(self, insn: Insn, model: Model) -> Optional[ProgInsn]:
        '''Fill out a BN.LID or BN.SID instruction'''
{'fraction_non_alphanumeric': 0.0761315682987456, 'fraction_numerical': 0.00612752182207064, 'mean_word_length': 1.9158941513568177, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 27, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3131105', 'n_tokens_mistral': 4935, 'n_tokens_neox': 4738, 'n_words': 1691}