import avl
from avl_riscv_coverage._coverage import COVERPACKAGES, CoverPackage
from avl_riscv_coverage._elf import EXTENSIONS
from avl_riscv_coverage._trace import Trace


class Jump(CoverPackage):

    def __init__(self) -> None:
        super().__init__("Jump")

        self.mnemonics = ["jal", "jalr"]

        if "C" in EXTENSIONS:
            self.mnemonics.extend(["c_jal", "c_jalr"])

        self.mnemonic = None
        self.distance = 0

        self.cg = avl.Covergroup("jump", None)

        cp_mnemonic = self.cg.add_coverpoint("mnemonic", lambda: self.mnemonic)
        for m in self.mnemonics:
            cp_mnemonic.add_bin(m, lambda x,m=m: x == m)

        cp_distance = self.cg.add_coverpoint("distance", lambda: self.distance)
        cp_distance.add_bin("short", range(-64,64), stats=True)
        cp_distance.add_bin("medium", range(-1024,1024), stats=True)
        cp_distance.add_bin("any", lambda x: True, stats=True)

        self.cg.add_covercross("mnemonicXdistance", cp_mnemonic, cp_distance)

    def sample(self, trace : Trace) -> None:
        if trace.instr.encoding.mnemonic not in self.mnemonics:
            return
        self.mnemonic = trace.instr.encoding.mnemonic
        self.distance = trace.next.instr.pc - trace.instr.pc
        self.cg.sample()

COVERPACKAGES.append(Jump())
