[bigint-presentation-code.git] / src / bigint_presentation_code / compiler_ir.py

from contextlib import contextmanager
import enum
from abc import ABCMeta, abstractmethod
from enum import Enum, unique
from functools import lru_cache, total_ordering
from io import StringIO
from typing import (AbstractSet, Any, Callable, Generic, Iterable, Iterator,
                    Mapping, Sequence, TypeVar, Union, overload)
from weakref import WeakValueDictionary as _WeakVDict

from cached_property import cached_property
from nmutil.plain_data import fields, plain_data

from bigint_presentation_code.type_util import (Literal, Self, assert_never,
                                                final)
from bigint_presentation_code.util import (BitSet, FBitSet, FMap, InternedMeta,
                                           OFSet, OSet)

GPR_SIZE_IN_BYTES = 8
BITS_IN_BYTE = 8
GPR_SIZE_IN_BITS = GPR_SIZE_IN_BYTES * BITS_IN_BYTE
GPR_VALUE_MASK = (1 << GPR_SIZE_IN_BITS) - 1


@final
class Fn:
    def __init__(self):
        self.ops = []  # type: list[Op]
        self.__op_names = _WeakVDict()  # type: _WeakVDict[str, Op]
        self.__next_name_suffix = 2

    def _add_op_with_unused_name(self, op, name=""):
        # type: (Op, str) -> str
        if op.fn is not self:
            raise ValueError("can't add Op to wrong Fn")
        if hasattr(op, "name"):
            raise ValueError("Op already named")
        orig_name = name
        while True:
            if name != "" and name not in self.__op_names:
                self.__op_names[name] = op
                return name
            name = orig_name + str(self.__next_name_suffix)
            self.__next_name_suffix += 1

    def __repr__(self):
        # type: () -> str
        return "<Fn>"

    def ops_to_str(self, as_python_literal=False, wrap_width=63,
                   python_indent="            ", indent="    "):
        # type: (bool, int, str, str) -> str
        l = []  # type: list[str]
        for op in self.ops:
            l.append(op.__repr__(wrap_width=wrap_width, indent=indent))
        retval = "\n".join(l)
        if as_python_literal:
            l = [python_indent + "\""]
            for ch in retval:
                if ch == "\n":
                    l.append(f"\\n\"\n{python_indent}\"")
                elif ch in "\"\\":
                    l.append("\\" + ch)
                elif ch.isascii() and ch.isprintable():
                    l.append(ch)
                else:
                    l.append(repr(ch).strip("\"'"))
            l.append("\"")
            retval = "".join(l)
            empty_end = f"\"\n{python_indent}\"\""
            if retval.endswith(empty_end):
                retval = retval[:-len(empty_end)]
        return retval

    def append_op(self, op):
        # type: (Op) -> None
        if op.fn is not self:
            raise ValueError("can't add Op to wrong Fn")
        self.ops.append(op)

    def append_new_op(self, kind, input_vals=(), immediates=(), name="",
                      maxvl=1):
        # type: (OpKind, Iterable[SSAVal], Iterable[int], str, int) -> Op
        retval = Op(fn=self, properties=kind.instantiate(maxvl=maxvl),
                    input_vals=input_vals, immediates=immediates, name=name)
        self.append_op(retval)
        return retval

    def sim(self, state):
        # type: (BaseSimState) -> None
        for op in self.ops:
            op.sim(state)

    def gen_asm(self, state):
        # type: (GenAsmState) -> None
        for op in self.ops:
            op.gen_asm(state)

    def pre_ra_insert_copies(self):
        # type: () -> None
        orig_ops = list(self.ops)
        copied_outputs = {}  # type: dict[SSAVal, SSAVal]
        setvli_outputs = {}  # type: dict[SSAVal, Op]
        self.ops.clear()
        for op in orig_ops:
            for i in range(len(op.input_vals)):
                inp = copied_outputs[op.input_vals[i]]
                if inp.ty.base_ty is BaseTy.I64:
                    maxvl = inp.ty.reg_len
                    if inp.ty.reg_len != 1:
                        setvl = self.append_new_op(
                            OpKind.SetVLI, immediates=[maxvl],
                            name=f"{op.name}.inp{i}.setvl")
                        vl = setvl.outputs[0]
                        mv = self.append_new_op(
                            OpKind.VecCopyToReg, input_vals=[inp, vl],
                            maxvl=maxvl, name=f"{op.name}.inp{i}.copy")
                    else:
                        mv = self.append_new_op(
                            OpKind.CopyToReg, input_vals=[inp],
                            name=f"{op.name}.inp{i}.copy")
                    op.input_vals[i] = mv.outputs[0]
                elif inp.ty.base_ty is BaseTy.CA \
                        or inp.ty.base_ty is BaseTy.VL_MAXVL:
                    # all copies would be no-ops, so we don't need to copy,
                    # though we do need to rematerialize SetVLI ops right
                    # before the ops VL
                    if inp in setvli_outputs:
                        setvl = self.append_new_op(
                            OpKind.SetVLI,
                            immediates=setvli_outputs[inp].immediates,
                            name=f"{op.name}.inp{i}.setvl")
                        inp = setvl.outputs[0]
                    op.input_vals[i] = inp
                else:
                    assert_never(inp.ty.base_ty)
            self.ops.append(op)
            for i, out in enumerate(op.outputs):
                if op.kind is OpKind.SetVLI:
                    setvli_outputs[out] = op
                if out.ty.base_ty is BaseTy.I64:
                    maxvl = out.ty.reg_len
                    if out.ty.reg_len != 1:
                        setvl = self.append_new_op(
                            OpKind.SetVLI, immediates=[maxvl],
                            name=f"{op.name}.out{i}.setvl")
                        vl = setvl.outputs[0]
                        mv = self.append_new_op(
                            OpKind.VecCopyFromReg, input_vals=[out, vl],
                            maxvl=maxvl, name=f"{op.name}.out{i}.copy")
                    else:
                        mv = self.append_new_op(
                            OpKind.CopyFromReg, input_vals=[out],
                            name=f"{op.name}.out{i}.copy")
                    copied_outputs[out] = mv.outputs[0]
                elif out.ty.base_ty is BaseTy.CA \
                        or out.ty.base_ty is BaseTy.VL_MAXVL:
                    # all copies would be no-ops, so we don't need to copy
                    copied_outputs[out] = out
                else:
                    assert_never(out.ty.base_ty)


@final
@unique
@total_ordering
class OpStage(Enum):
    value: Literal[0, 1]  # type: ignore

    def __new__(cls, value):
        # type: (int) -> OpStage
        value = int(value)
        if value not in (0, 1):
            raise ValueError("invalid value")
        retval = object.__new__(cls)
        retval._value_ = value
        return retval

    Early = 0
    """ early stage of Op execution, where all input reads occur.
    all output writes with `write_stage == Early` occur here too, and therefore
    conflict with input reads, telling the compiler that it that can't share
    that output's register with any inputs that the output isn't tied to.

    All outputs, even unused outputs, can't share registers with any other
    outputs, independent of `write_stage` settings.
    """
    Late = 1
    """ late stage of Op execution, where all output writes with
    `write_stage == Late` occur, and therefore don't conflict with input reads,
    telling the compiler that any inputs can safely use the same register as
    those outputs.

    All outputs, even unused outputs, can't share registers with any other
    outputs, independent of `write_stage` settings.
    """

    def __repr__(self):
        # type: () -> str
        return f"OpStage.{self._name_}"

    def __lt__(self, other):
        # type: (OpStage | object) -> bool
        if isinstance(other, OpStage):
            return self.value < other.value
        return NotImplemented


assert OpStage.Early < OpStage.Late, "early must be less than late"


@plain_data(frozen=True, unsafe_hash=True, repr=False)
@final
@total_ordering
class ProgramPoint(metaclass=InternedMeta):
    __slots__ = "op_index", "stage"

    def __init__(self, op_index, stage):
        # type: (int, OpStage) -> None
        self.op_index = op_index
        self.stage = stage

    @property
    def int_value(self):
        # type: () -> int
        """ an integer representation of `self` such that it keeps ordering and
        successor/predecessor relations.
        """
        return self.op_index * 2 + self.stage.value

    @staticmethod
    def from_int_value(int_value):
        # type: (int) -> ProgramPoint
        op_index, stage = divmod(int_value, 2)
        return ProgramPoint(op_index=op_index, stage=OpStage(stage))

    def next(self, steps=1):
        # type: (int) -> ProgramPoint
        return ProgramPoint.from_int_value(self.int_value + steps)

    def prev(self, steps=1):
        # type: (int) -> ProgramPoint
        return self.next(steps=-steps)

    def __lt__(self, other):
        # type: (ProgramPoint | Any) -> bool
        if not isinstance(other, ProgramPoint):
            return NotImplemented
        if self.op_index != other.op_index:
            return self.op_index < other.op_index
        return self.stage < other.stage

    def __repr__(self):
        # type: () -> str
        return f"<ops[{self.op_index}]:{self.stage._name_}>"


@plain_data(frozen=True, unsafe_hash=True, repr=False)
@final
class ProgramRange(Sequence[ProgramPoint], metaclass=InternedMeta):
    __slots__ = "start", "stop"

    def __init__(self, start, stop):
        # type: (ProgramPoint, ProgramPoint) -> None
        self.start = start
        self.stop = stop

    @cached_property
    def int_value_range(self):
        # type: () -> range
        return range(self.start.int_value, self.stop.int_value)

    @staticmethod
    def from_int_value_range(int_value_range):
        # type: (range) -> ProgramRange
        if int_value_range.step != 1:
            raise ValueError("int_value_range must have step == 1")
        return ProgramRange(
            start=ProgramPoint.from_int_value(int_value_range.start),
            stop=ProgramPoint.from_int_value(int_value_range.stop))

    @overload
    def __getitem__(self, __idx):
        # type: (int) -> ProgramPoint
        ...

    @overload
    def __getitem__(self, __idx):
        # type: (slice) -> ProgramRange
        ...

    def __getitem__(self, __idx):
        # type: (int | slice) -> ProgramPoint | ProgramRange
        v = range(self.start.int_value, self.stop.int_value)[__idx]
        if isinstance(v, int):
            return ProgramPoint.from_int_value(v)
        return ProgramRange.from_int_value_range(v)

    def __len__(self):
        # type: () -> int
        return len(self.int_value_range)

    def __iter__(self):
        # type: () -> Iterator[ProgramPoint]
        return map(ProgramPoint.from_int_value, self.int_value_range)

    def __repr__(self):
        # type: () -> str
        start = repr(self.start).lstrip("<").rstrip(">")
        stop = repr(self.stop).lstrip("<").rstrip(">")
        return f"<range:{start}..{stop}>"


@plain_data(frozen=True, unsafe_hash=True, repr=False)
@final
class SSAValSubReg(metaclass=InternedMeta):
    __slots__ = "ssa_val", "reg_idx"

    def __init__(self, ssa_val, reg_idx):
        # type: (SSAVal, int) -> None
        if reg_idx < 0 or reg_idx >= ssa_val.ty.reg_len:
            raise ValueError("reg_idx out of range")
        self.ssa_val = ssa_val
        self.reg_idx = reg_idx

    def __repr__(self):
        # type: () -> str
        return f"{self.ssa_val}[{self.reg_idx}]"


@plain_data(frozen=True, eq=False, repr=False)
@final
class FnAnalysis:
    __slots__ = ("fn", "uses", "op_indexes", "live_ranges", "live_at",
                 "def_program_ranges", "use_program_points",
                 "all_program_points")

    def __init__(self, fn):
        # type: (Fn) -> None
        self.fn = fn
        self.op_indexes = FMap((op, idx) for idx, op in enumerate(fn.ops))
        self.all_program_points = ProgramRange(
            start=ProgramPoint(op_index=0, stage=OpStage.Early),
            stop=ProgramPoint(op_index=len(fn.ops), stage=OpStage.Early))
        def_program_ranges = {}  # type: dict[SSAVal, ProgramRange]
        use_program_points = {}  # type: dict[SSAUse, ProgramPoint]
        uses = {}  # type: dict[SSAVal, OSet[SSAUse]]
        live_range_stops = {}  # type: dict[SSAVal, ProgramPoint]
        for op in fn.ops:
            for use in op.input_uses:
                uses[use.ssa_val].add(use)
                use_program_point = self.__get_use_program_point(use)
                use_program_points[use] = use_program_point
                live_range_stops[use.ssa_val] = max(
                    live_range_stops[use.ssa_val], use_program_point.next())
            for out in op.outputs:
                uses[out] = OSet()
                def_program_range = self.__get_def_program_range(out)
                def_program_ranges[out] = def_program_range
                live_range_stops[out] = def_program_range.stop
        self.uses = FMap((k, OFSet(v)) for k, v in uses.items())
        self.def_program_ranges = FMap(def_program_ranges)
        self.use_program_points = FMap(use_program_points)
        live_ranges = {}  # type: dict[SSAVal, ProgramRange]
        live_at = {i: OSet[SSAVal]() for i in self.all_program_points}
        for ssa_val in uses.keys():
            live_ranges[ssa_val] = live_range = ProgramRange(
                start=self.def_program_ranges[ssa_val].start,
                stop=live_range_stops[ssa_val])
            for program_point in live_range:
                live_at[program_point].add(ssa_val)
        self.live_ranges = FMap(live_ranges)
        self.live_at = FMap((k, OFSet(v)) for k, v in live_at.items())
        self.copies  # initialize
        self.const_ssa_vals  # initialize
        self.const_ssa_val_sub_regs  # initialize

    def __get_def_program_range(self, ssa_val):
        # type: (SSAVal) -> ProgramRange
        write_stage = ssa_val.defining_descriptor.write_stage
        start = ProgramPoint(
            op_index=self.op_indexes[ssa_val.op], stage=write_stage)
        # always include late stage of ssa_val.op, to ensure outputs always
        # overlap all other outputs.
        # stop is exclusive, so we need the next program point.
        stop = ProgramPoint(op_index=start.op_index, stage=OpStage.Late).next()
        return ProgramRange(start=start, stop=stop)

    def __get_use_program_point(self, ssa_use):
        # type: (SSAUse) -> ProgramPoint
        assert ssa_use.defining_descriptor.write_stage is OpStage.Early, \
            "assumed here, ensured by GenericOpProperties.__init__"
        return ProgramPoint(
            op_index=self.op_indexes[ssa_use.op], stage=OpStage.Early)

    def __eq__(self, other):
        # type: (FnAnalysis | Any) -> bool
        if isinstance(other, FnAnalysis):
            return self.fn == other.fn
        return NotImplemented

    def __hash__(self):
        # type: () -> int
        return hash(self.fn)

    def __repr__(self):
        # type: () -> str
        return "<FnAnalysis>"

    @cached_property
    def copies(self):
        # type: () -> FMap[SSAValSubReg, SSAValSubReg]
        """ map from SSAValSubRegs to the original SSAValSubRegs that they are
        a copy of, looking through all layers of copies. The map excludes all
        SSAValSubRegs that aren't copies of other SSAValSubRegs.
        This ignores inputs of copy Ops that aren't actually being copied
        (e.g. the VL input of VecCopyToReg).
        """
        retval = {}  # type: dict[SSAValSubReg, SSAValSubReg]
        for op in self.op_indexes.keys():
            if not op.properties.is_copy:
                continue
            copy_reg_len = op.properties.copy_reg_len
            copy_inputs = []  # type: list[SSAValSubReg]
            for inp in op.input_vals[:op.properties.copy_inputs_len]:
                for inp_sub_reg in inp.ssa_val_sub_regs:
                    # propagate copies of copies
                    inp_sub_reg = retval.get(inp_sub_reg, inp_sub_reg)
                    copy_inputs.append(inp_sub_reg)
            assert len(copy_inputs) == copy_reg_len, "logic error"
            copy_outputs = []  # type: list[SSAValSubReg]
            for out in op.outputs[:op.properties.copy_outputs_len]:
                copy_outputs.extend(out.ssa_val_sub_regs)
            assert len(copy_outputs) == copy_reg_len, "logic error"
            for inp, out in zip(copy_inputs, copy_outputs):
                retval[out] = inp
        return FMap(retval)

    @cached_property
    def copy_related_ssa_vals(self):
        # type: () -> FMap[SSAVal, OFSet[SSAVal]]
        """ map from SSAVals to the full set of SSAVals that are related by
        being sources/destinations of copies, transitively looking through all
        copies.
        This ignores inputs of copy Ops that aren't actually being copied
        (e.g. the VL input of VecCopyToReg).
        """
        sets_map = {i: OSet([i]) for i in self.uses.keys()}
        for k, v in self.copies.items():
            k_set = sets_map[k.ssa_val]
            v_set = sets_map[v.ssa_val]
            # merge k_set and v_set
            if k_set is v_set:
                continue
            k_set |= v_set
            for i in k_set:
                sets_map[i] = k_set
        # this way we construct each OFSet only once rather than
        # for each SSAVal
        sets_set = {id(i): i for i in sets_map.values()}
        retval = {}  # type: dict[SSAVal, OFSet[SSAVal]]
        for v in sets_set.values():
            v = OFSet(v)
            for k in v:
                retval[k] = v
        return FMap(retval)

    @cached_property
    def const_ssa_vals(self):
        # type: () -> FMap[SSAVal, tuple[int, ...]]
        state = ConstPropagationState(
            ssa_vals={}, memory={}, skipped_ops=OSet())
        self.fn.sim(state)
        return FMap(state.ssa_vals)

    @cached_property
    def const_ssa_val_sub_regs(self):
        # type: () -> FMap[SSAValSubReg, int]
        retval = {}  # type: dict[SSAValSubReg, int]
        for ssa_val, const_val in self.const_ssa_vals.items():
            assert ssa_val.ty.reg_len == len(const_val), "logic error"
            for reg_idx, v in enumerate(const_val):
                retval[SSAValSubReg(ssa_val, reg_idx)] = v
        return FMap(retval)

    def is_always_equal(self, a, b):
        # type: (SSAValSubReg, SSAValSubReg) -> bool
        """check if a and b are known to be always equal to each other.
        This means they can be allocated to the same location if other
        constraints don't prevent that.

        this can happen for a number of reasons, such as:
        * a and b are copies of the same thing
        * a and b are known to be constants and they have the same value
        """
        if a.ssa_val.base_ty != b.ssa_val.base_ty:
            return False  # can't be equal, they have different types
        # look through copies
        a = self.copies.get(a, a)
        b = self.copies.get(b, b)
        if a == b:
            return True
        # check if they have the same constant value
        try:
            a_const_val = self.const_ssa_val_sub_regs[a]
            b_const_val = self.const_ssa_val_sub_regs[b]
            if a_const_val == b_const_val:
                return True
        except KeyError:
            pass
        return False


@unique
@final
class BaseTy(Enum):
    I64 = enum.auto()
    CA = enum.auto()
    VL_MAXVL = enum.auto()

    @cached_property
    def only_scalar(self):
        # type: () -> bool
        if self is BaseTy.I64:
            return False
        elif self is BaseTy.CA or self is BaseTy.VL_MAXVL:
            return True
        else:
            assert_never(self)

    @cached_property
    def max_reg_len(self):
        # type: () -> int
        if self is BaseTy.I64:
            return 128
        elif self is BaseTy.CA or self is BaseTy.VL_MAXVL:
            return 1
        else:
            assert_never(self)

    def __repr__(self):
        return "BaseTy." + self._name_


@plain_data(frozen=True, unsafe_hash=True, repr=False)
@final
class Ty(metaclass=InternedMeta):
    __slots__ = "base_ty", "reg_len"

    @staticmethod
    def validate(base_ty, reg_len):
        # type: (BaseTy, int) -> str | None
        """ return a string with the error if the combination is invalid,
        otherwise return None
        """
        if base_ty.only_scalar and reg_len != 1:
            return f"can't create a vector of an only-scalar type: {base_ty}"
        if reg_len < 1 or reg_len > base_ty.max_reg_len:
            return "reg_len out of range"
        return None

    def __init__(self, base_ty, reg_len):
        # type: (BaseTy, int) -> None
        msg = self.validate(base_ty=base_ty, reg_len=reg_len)
        if msg is not None:
            raise ValueError(msg)
        self.base_ty = base_ty
        self.reg_len = reg_len

    def __repr__(self):
        # type: () -> str
        if self.reg_len != 1:
            reg_len = f"*{self.reg_len}"
        else:
            reg_len = ""
        return f"<{self.base_ty._name_}{reg_len}>"


@unique
@final
class LocKind(Enum):
    GPR = enum.auto()
    StackI64 = enum.auto()
    CA = enum.auto()
    VL_MAXVL = enum.auto()

    @cached_property
    def base_ty(self):
        # type: () -> BaseTy
        if self is LocKind.GPR or self is LocKind.StackI64:
            return BaseTy.I64
        if self is LocKind.CA:
            return BaseTy.CA
        if self is LocKind.VL_MAXVL:
            return BaseTy.VL_MAXVL
        else:
            assert_never(self)

    @cached_property
    def loc_count(self):
        # type: () -> int
        if self is LocKind.StackI64:
            return 512
        if self is LocKind.GPR or self is LocKind.CA \
                or self is LocKind.VL_MAXVL:
            return self.base_ty.max_reg_len
        else:
            assert_never(self)

    def __repr__(self):
        return "LocKind." + self._name_


@final
@unique
class LocSubKind(Enum):
    BASE_GPR = enum.auto()
    SV_EXTRA2_VGPR = enum.auto()
    SV_EXTRA2_SGPR = enum.auto()
    SV_EXTRA3_VGPR = enum.auto()
    SV_EXTRA3_SGPR = enum.auto()
    StackI64 = enum.auto()
    CA = enum.auto()
    VL_MAXVL = enum.auto()

    @cached_property
    def kind(self):
        # type: () -> LocKind
        # pyright fails typechecking when using `in` here:
        # reported: https://github.com/microsoft/pyright/issues/4102
        if self in (LocSubKind.BASE_GPR, LocSubKind.SV_EXTRA2_VGPR,
                    LocSubKind.SV_EXTRA2_SGPR, LocSubKind.SV_EXTRA3_VGPR,
                    LocSubKind.SV_EXTRA3_SGPR):
            return LocKind.GPR
        if self is LocSubKind.StackI64:
            return LocKind.StackI64
        if self is LocSubKind.CA:
            return LocKind.CA
        if self is LocSubKind.VL_MAXVL:
            return LocKind.VL_MAXVL
        assert_never(self)

    @property
    def base_ty(self):
        return self.kind.base_ty

    @lru_cache()
    def allocatable_locs(self, ty):
        # type: (Ty) -> LocSet
        if ty.base_ty != self.base_ty:
            raise ValueError("type mismatch")
        if self is LocSubKind.BASE_GPR:
            starts = range(32)
        elif self is LocSubKind.SV_EXTRA2_VGPR:
            starts = range(0, 128, 2)
        elif self is LocSubKind.SV_EXTRA2_SGPR:
            starts = range(64)
        elif self is LocSubKind.SV_EXTRA3_VGPR \
                or self is LocSubKind.SV_EXTRA3_SGPR:
            starts = range(128)
        elif self is LocSubKind.StackI64:
            starts = range(LocKind.StackI64.loc_count)
        elif self is LocSubKind.CA or self is LocSubKind.VL_MAXVL:
            return LocSet([Loc(kind=self.kind, start=0, reg_len=1)])
        else:
            assert_never(self)
        retval = []  # type: list[Loc]
        for start in starts:
            loc = Loc.try_make(kind=self.kind, start=start, reg_len=ty.reg_len)
            if loc is None:
                continue
            conflicts = False
            for special_loc in SPECIAL_GPRS:
                if loc.conflicts(special_loc):
                    conflicts = True
                    break
            if not conflicts:
                retval.append(loc)
        return LocSet(retval)

    def __repr__(self):
        return "LocSubKind." + self._name_


@plain_data(frozen=True, unsafe_hash=True)
@final
class GenericTy(metaclass=InternedMeta):
    __slots__ = "base_ty", "is_vec"

    def __init__(self, base_ty, is_vec):
        # type: (BaseTy, bool) -> None
        self.base_ty = base_ty
        if base_ty.only_scalar and is_vec:
            raise ValueError(f"base_ty={base_ty} requires is_vec=False")
        self.is_vec = is_vec

    def instantiate(self, maxvl):
        # type: (int) -> Ty
        # here's where subvl and elwid would be accounted for
        if self.is_vec:
            return Ty(self.base_ty, maxvl)
        return Ty(self.base_ty, 1)

    def can_instantiate_to(self, ty):
        # type: (Ty) -> bool
        if self.base_ty != ty.base_ty:
            return False
        if self.is_vec:
            return True
        return ty.reg_len == 1


@plain_data(frozen=True, unsafe_hash=True)
@final
class Loc(metaclass=InternedMeta):
    __slots__ = "kind", "start", "reg_len"

    @staticmethod
    def validate(kind, start, reg_len):
        # type: (LocKind, int, int) -> str | None
        msg = Ty.validate(base_ty=kind.base_ty, reg_len=reg_len)
        if msg is not None:
            return msg
        if reg_len > kind.loc_count:
            return "invalid reg_len"
        if start < 0 or start + reg_len > kind.loc_count:
            return "start not in valid range"
        return None

    @staticmethod
    def try_make(kind, start, reg_len):
        # type: (LocKind, int, int) -> Loc | None
        msg = Loc.validate(kind=kind, start=start, reg_len=reg_len)
        if msg is not None:
            return None
        return Loc(kind=kind, start=start, reg_len=reg_len)

    def __init__(self, kind, start, reg_len):
        # type: (LocKind, int, int) -> None
        msg = self.validate(kind=kind, start=start, reg_len=reg_len)
        if msg is not None:
            raise ValueError(msg)
        self.kind = kind
        self.reg_len = reg_len
        self.start = start

    def conflicts(self, other):
        # type: (Loc) -> bool
        return (self.kind == other.kind
                and self.start < other.stop and other.start < self.stop)

    @staticmethod
    def make_ty(kind, reg_len):
        # type: (LocKind, int) -> Ty
        return Ty(base_ty=kind.base_ty, reg_len=reg_len)

    @cached_property
    def ty(self):
        # type: () -> Ty
        return self.make_ty(kind=self.kind, reg_len=self.reg_len)

    @property
    def stop(self):
        # type: () -> int
        return self.start + self.reg_len

    def try_concat(self, *others):
        # type: (*Loc | None) -> Loc | None
        reg_len = self.reg_len
        stop = self.stop
        for other in others:
            if other is None or other.kind != self.kind:
                return None
            if stop != other.start:
                return None
            stop = other.stop
            reg_len += other.reg_len
        return Loc(kind=self.kind, start=self.start, reg_len=reg_len)

    def get_subloc_at_offset(self, subloc_ty, offset):
        # type: (Ty, int) -> Loc
        if subloc_ty.base_ty != self.kind.base_ty:
            raise ValueError("BaseTy mismatch")
        if offset < 0 or offset + subloc_ty.reg_len > self.reg_len:
            raise ValueError("invalid sub-Loc: offset and/or "
                             "subloc_ty.reg_len out of range")
        return Loc(kind=self.kind,
                   start=self.start + offset, reg_len=subloc_ty.reg_len)


SPECIAL_GPRS = (
    Loc(kind=LocKind.GPR, start=0, reg_len=1),
    Loc(kind=LocKind.GPR, start=1, reg_len=1),
    Loc(kind=LocKind.GPR, start=2, reg_len=1),
    Loc(kind=LocKind.GPR, start=13, reg_len=1),
)


@final
class LocSet(OFSet[Loc], metaclass=InternedMeta):
    def __init__(self, __locs=()):
        # type: (Iterable[Loc]) -> None
        super().__init__(__locs)
        if isinstance(__locs, LocSet):
            self.__starts = __locs.starts
            self.__ty = __locs.ty
            return
        starts = {i: BitSet() for i in LocKind}
        ty = None  # type: None | Ty
        for loc in self:
            if ty is None:
                ty = loc.ty
            if ty != loc.ty:
                raise ValueError(f"conflicting types: {ty} != {loc.ty}")
            starts[loc.kind].add(loc.start)
        self.__starts = FMap(
            (k, FBitSet(v)) for k, v in starts.items() if len(v) != 0)
        self.__ty = ty

    @property
    def starts(self):
        # type: () -> FMap[LocKind, FBitSet]
        return self.__starts

    @property
    def ty(self):
        # type: () -> Ty | None
        return self.__ty

    @cached_property
    def stops(self):
        # type: () -> FMap[LocKind, FBitSet]
        if self.ty is None:
            return FMap()
        sh = self.ty.reg_len
        return FMap(
            (k, FBitSet(bits=v.bits << sh)) for k, v in self.starts.items())

    @property
    def kinds(self):
        # type: () -> AbstractSet[LocKind]
        return self.starts.keys()

    @property
    def reg_len(self):
        # type: () -> int | None
        if self.ty is None:
            return None
        return self.ty.reg_len

    @property
    def base_ty(self):
        # type: () -> BaseTy | None
        if self.ty is None:
            return None
        return self.ty.base_ty

    def concat(self, *others):
        # type: (*LocSet) -> LocSet
        if self.ty is None:
            return LocSet()
        base_ty = self.ty.base_ty
        reg_len = self.ty.reg_len
        starts = {k: BitSet(v) for k, v in self.starts.items()}
        for other in others:
            if other.ty is None:
                return LocSet()
            if other.ty.base_ty != base_ty:
                return LocSet()
            for kind, other_starts in other.starts.items():
                if kind not in starts:
                    continue
                starts[kind].bits &= other_starts.bits >> reg_len
                if starts[kind] == 0:
                    del starts[kind]
                    if len(starts) == 0:
                        return LocSet()
            reg_len += other.ty.reg_len

        def locs():
            # type: () -> Iterable[Loc]
            for kind, v in starts.items():
                for start in v:
                    loc = Loc.try_make(kind=kind, start=start, reg_len=reg_len)
                    if loc is not None:
                        yield loc
        return LocSet(locs())

    @lru_cache(maxsize=None, typed=True)
    def max_conflicts_with(self, other):
        # type: (LocSet | Loc) -> int
        """the largest number of Locs in `self` that a single Loc
        from `other` can conflict with
        """
        if isinstance(other, LocSet):
            return max(self.max_conflicts_with(i) for i in other)
        else:
            return sum(other.conflicts(i) for i in self)

    def __repr__(self):
        return f"LocSet(starts={self.starts!r}, ty={self.ty!r})"


@plain_data(frozen=True, unsafe_hash=True)
@final
class GenericOperandDesc(metaclass=InternedMeta):
    """generic Op operand descriptor"""
    __slots__ = ("ty", "fixed_loc", "sub_kinds", "tied_input_index", "spread",
                 "write_stage")

    def __init__(
        self, ty,  # type: GenericTy
        sub_kinds,  # type: Iterable[LocSubKind]
        *,
        fixed_loc=None,  # type: Loc | None
        tied_input_index=None,  # type: int | None
        spread=False,  # type: bool
        write_stage=OpStage.Early,  # type: OpStage
    ):
        # type: (...) -> None
        self.ty = ty
        self.sub_kinds = OFSet(sub_kinds)
        if len(self.sub_kinds) == 0:
            raise ValueError("sub_kinds can't be empty")
        self.fixed_loc = fixed_loc
        if fixed_loc is not None:
            if tied_input_index is not None:
                raise ValueError("operand can't be both tied and fixed")
            if not ty.can_instantiate_to(fixed_loc.ty):
                raise ValueError(
                    f"fixed_loc has incompatible type for given generic "
                    f"type: fixed_loc={fixed_loc} generic ty={ty}")
            if len(self.sub_kinds) != 1:
                raise ValueError(
                    "multiple sub_kinds not allowed for fixed operand")
            for sub_kind in self.sub_kinds:
                if fixed_loc not in sub_kind.allocatable_locs(fixed_loc.ty):
                    raise ValueError(
                        f"fixed_loc not in given sub_kind: "
                        f"fixed_loc={fixed_loc} sub_kind={sub_kind}")
        for sub_kind in self.sub_kinds:
            if sub_kind.base_ty != ty.base_ty:
                raise ValueError(f"sub_kind is incompatible with type: "
                                 f"sub_kind={sub_kind} ty={ty}")
        if tied_input_index is not None and tied_input_index < 0:
            raise ValueError("invalid tied_input_index")
        self.tied_input_index = tied_input_index
        self.spread = spread
        if spread:
            if self.tied_input_index is not None:
                raise ValueError("operand can't be both spread and tied")
            if self.fixed_loc is not None:
                raise ValueError("operand can't be both spread and fixed")
            if self.ty.is_vec:
                raise ValueError("operand can't be both spread and vector")
        self.write_stage = write_stage

    @cached_property
    def ty_before_spread(self):
        # type: () -> GenericTy
        if self.spread:
            return GenericTy(base_ty=self.ty.base_ty, is_vec=True)
        return self.ty

    def tied_to_input(self, tied_input_index):
        # type: (int) -> Self
        return GenericOperandDesc(self.ty, self.sub_kinds,
                                  tied_input_index=tied_input_index,
                                  write_stage=self.write_stage)

    def with_fixed_loc(self, fixed_loc):
        # type: (Loc) -> Self
        return GenericOperandDesc(self.ty, self.sub_kinds, fixed_loc=fixed_loc,
                                  write_stage=self.write_stage)

    def with_write_stage(self, write_stage):
        # type: (OpStage) -> Self
        return GenericOperandDesc(self.ty, self.sub_kinds,
                                  fixed_loc=self.fixed_loc,
                                  tied_input_index=self.tied_input_index,
                                  spread=self.spread,
                                  write_stage=write_stage)

    def instantiate(self, maxvl):
        # type: (int) -> Iterable[OperandDesc]
        # assumes all spread operands have ty.reg_len = 1
        rep_count = 1
        if self.spread:
            rep_count = maxvl
        ty_before_spread = self.ty_before_spread.instantiate(maxvl=maxvl)

        def locs_before_spread():
            # type: () -> Iterable[Loc]
            if self.fixed_loc is not None:
                if ty_before_spread != self.fixed_loc.ty:
                    raise ValueError(
                        f"instantiation failed: type mismatch with fixed_loc: "
                        f"instantiated type: {ty_before_spread} "
                        f"fixed_loc: {self.fixed_loc}")
                yield self.fixed_loc
                return
            for sub_kind in self.sub_kinds:
                yield from sub_kind.allocatable_locs(ty_before_spread)
        loc_set_before_spread = LocSet(locs_before_spread())
        for idx in range(rep_count):
            if not self.spread:
                idx = None
            yield OperandDesc(loc_set_before_spread=loc_set_before_spread,
                              tied_input_index=self.tied_input_index,
                              spread_index=idx, write_stage=self.write_stage)


@plain_data(frozen=True, unsafe_hash=True)
@final
class OperandDesc(metaclass=InternedMeta):
    """Op operand descriptor"""
    __slots__ = ("loc_set_before_spread", "tied_input_index", "spread_index",
                 "write_stage")

    def __init__(self, loc_set_before_spread, tied_input_index, spread_index,
                 write_stage):
        # type: (LocSet, int | None, int | None, OpStage) -> None
        if len(loc_set_before_spread) == 0:
            raise ValueError("loc_set_before_spread must not be empty")
        self.loc_set_before_spread = loc_set_before_spread
        self.tied_input_index = tied_input_index
        if self.tied_input_index is not None and spread_index is not None:
            raise ValueError("operand can't be both spread and tied")
        self.spread_index = spread_index
        self.write_stage = write_stage

    @cached_property
    def ty_before_spread(self):
        # type: () -> Ty
        ty = self.loc_set_before_spread.ty
        assert ty is not None, (
            "__init__ checked that the LocSet isn't empty, "
            "non-empty LocSets should always have ty set")
        return ty

    @cached_property
    def ty(self):
        """ Ty after any spread is applied """
        if self.spread_index is not None:
            # assumes all spread operands have ty.reg_len = 1
            return Ty(base_ty=self.ty_before_spread.base_ty, reg_len=1)
        return self.ty_before_spread

    @property
    def reg_offset_in_unspread(self):
        """ the number of reg-sized slots in the unspread Loc before self's Loc

        e.g. if the unspread Loc containing self is:
        `Loc(kind=LocKind.GPR, start=8, reg_len=4)`
        and self's Loc is `Loc(kind=LocKind.GPR, start=10, reg_len=1)`
        then reg_offset_into_unspread == 2 == 10 - 8
        """
        if self.spread_index is None:
            return 0
        return self.spread_index * self.ty.reg_len


OD_BASE_SGPR = GenericOperandDesc(
    ty=GenericTy(base_ty=BaseTy.I64, is_vec=False),
    sub_kinds=[LocSubKind.BASE_GPR])
OD_EXTRA3_SGPR = GenericOperandDesc(
    ty=GenericTy(base_ty=BaseTy.I64, is_vec=False),
    sub_kinds=[LocSubKind.SV_EXTRA3_SGPR])
OD_EXTRA3_VGPR = GenericOperandDesc(
    ty=GenericTy(base_ty=BaseTy.I64, is_vec=True),
    sub_kinds=[LocSubKind.SV_EXTRA3_VGPR])
OD_EXTRA2_SGPR = GenericOperandDesc(
    ty=GenericTy(base_ty=BaseTy.I64, is_vec=False),
    sub_kinds=[LocSubKind.SV_EXTRA2_SGPR])
OD_EXTRA2_VGPR = GenericOperandDesc(
    ty=GenericTy(base_ty=BaseTy.I64, is_vec=True),
    sub_kinds=[LocSubKind.SV_EXTRA2_VGPR])
OD_CA = GenericOperandDesc(
    ty=GenericTy(base_ty=BaseTy.CA, is_vec=False),
    sub_kinds=[LocSubKind.CA])
OD_VL = GenericOperandDesc(
    ty=GenericTy(base_ty=BaseTy.VL_MAXVL, is_vec=False),
    sub_kinds=[LocSubKind.VL_MAXVL])


@plain_data(frozen=True, unsafe_hash=True)
@final
class GenericOpProperties(metaclass=InternedMeta):
    __slots__ = ("demo_asm", "inputs", "outputs", "immediates",
                 "is_copy", "is_load_immediate", "has_side_effects")

    def __init__(
        self, demo_asm,  # type: str
        inputs,  # type: Iterable[GenericOperandDesc]
        outputs,  # type: Iterable[GenericOperandDesc]
        immediates=(),  # type: Iterable[range]
        is_copy=False,  # type: bool
        is_load_immediate=False,  # type: bool
        has_side_effects=False,  # type: bool
    ):
        # type: (...) -> None
        self.demo_asm = demo_asm  # type: str
        self.inputs = tuple(inputs)  # type: tuple[GenericOperandDesc, ...]
        for inp in self.inputs:
            if inp.tied_input_index is not None:
                raise ValueError(
                    f"tied_input_index is not allowed on inputs: {inp}")
            if inp.write_stage is not OpStage.Early:
                raise ValueError(
                    f"write_stage is not allowed on inputs: {inp}")
        self.outputs = tuple(outputs)  # type: tuple[GenericOperandDesc, ...]
        fixed_locs = []  # type: list[tuple[Loc, int]]
        for idx, out in enumerate(self.outputs):
            if out.tied_input_index is not None:
                if out.tied_input_index >= len(self.inputs):
                    raise ValueError(f"tied_input_index out of range: {out}")
                tied_inp = self.inputs[out.tied_input_index]
                expected_out = tied_inp.tied_to_input(out.tied_input_index) \
                    .with_write_stage(out.write_stage)
                if expected_out != out:
                    raise ValueError(f"output can't be tied to non-equivalent "
                                     f"input: {out} tied to {tied_inp}")
            if out.fixed_loc is not None:
                for other_fixed_loc, other_idx in fixed_locs:
                    if not other_fixed_loc.conflicts(out.fixed_loc):
                        continue
                    raise ValueError(
                        f"conflicting fixed_locs: outputs[{idx}] and "
                        f"outputs[{other_idx}]: {out.fixed_loc} conflicts "
                        f"with {other_fixed_loc}")
                fixed_locs.append((out.fixed_loc, idx))
        self.immediates = tuple(immediates)  # type: tuple[range, ...]
        self.is_copy = is_copy  # type: bool
        self.is_load_immediate = is_load_immediate  # type: bool
        self.has_side_effects = has_side_effects  # type: bool


@plain_data(frozen=True, unsafe_hash=True)
@final
class OpProperties(metaclass=InternedMeta):
    __slots__ = "kind", "inputs", "outputs", "maxvl", "copy_reg_len"

    def __init__(self, kind, maxvl):
        # type: (OpKind, int) -> None
        self.kind = kind  # type: OpKind
        inputs = []  # type: list[OperandDesc]
        for inp in self.generic.inputs:
            inputs.extend(inp.instantiate(maxvl=maxvl))
        self.inputs = tuple(inputs)  # type: tuple[OperandDesc, ...]
        outputs = []  # type: list[OperandDesc]
        for out in self.generic.outputs:
            outputs.extend(out.instantiate(maxvl=maxvl))
        self.outputs = tuple(outputs)  # type: tuple[OperandDesc, ...]
        self.maxvl = maxvl  # type: int
        copy_input_reg_len = 0
        for inp in self.inputs[:self.copy_inputs_len]:
            copy_input_reg_len += inp.ty.reg_len
        copy_output_reg_len = 0
        for out in self.outputs[:self.copy_outputs_len]:
            copy_output_reg_len += out.ty.reg_len
        if copy_input_reg_len != copy_output_reg_len:
            raise ValueError(f"invalid copy: copy's input reg len must "
                             f"match its output reg len: "
                             f"{copy_input_reg_len} != {copy_output_reg_len}")
        self.copy_reg_len = copy_input_reg_len

    @property
    def generic(self):
        # type: () -> GenericOpProperties
        return self.kind.properties

    @property
    def immediates(self):
        # type: () -> tuple[range, ...]
        return self.generic.immediates

    @property
    def demo_asm(self):
        # type: () -> str
        return self.generic.demo_asm

    @property
    def is_copy(self):
        # type: () -> bool
        return self.generic.is_copy

    @property
    def is_load_immediate(self):
        # type: () -> bool
        return self.generic.is_load_immediate

    @property
    def has_side_effects(self):
        # type: () -> bool
        return self.generic.has_side_effects

    @cached_property
    def copy_inputs_len(self):
        # type: () -> int
        if not self.is_copy:
            return 0
        if self.inputs[0].spread_index is None:
            return 1
        retval = 0
        for i, inp in enumerate(self.inputs):
            if inp.spread_index != i:
                break
            retval += 1
        return retval

    @cached_property
    def copy_outputs_len(self):
        # type: () -> int
        if not self.is_copy:
            return 0
        if self.outputs[0].spread_index is None:
            return 1
        retval = 0
        for i, out in enumerate(self.outputs):
            if out.spread_index != i:
                break
            retval += 1
        return retval


IMM_S16 = range(-1 << 15, 1 << 15)

_SIM_FN = Callable[["Op", "BaseSimState"], None]
_SIM_FN2 = Callable[[], _SIM_FN]
_SIM_FNS = {}  # type: dict[GenericOpProperties | Any, _SIM_FN2]
_GEN_ASM_FN = Callable[["Op", "GenAsmState"], None]
_GEN_ASM_FN2 = Callable[[], _GEN_ASM_FN]
_GEN_ASMS = {}  # type: dict[GenericOpProperties | Any, _GEN_ASM_FN2]


@unique
@final
class OpKind(Enum):
    def __init__(self, properties):
        # type: (GenericOpProperties) -> None
        super().__init__()
        self.__properties = properties

    @property
    def properties(self):
        # type: () -> GenericOpProperties
        return self.__properties

    def instantiate(self, maxvl):
        # type: (int) -> OpProperties
        return OpProperties(self, maxvl=maxvl)

    def __repr__(self):
        # type: () -> str
        return "OpKind." + self._name_

    @cached_property
    def sim(self):
        # type: () -> _SIM_FN
        return _SIM_FNS[self.properties]()

    @cached_property
    def gen_asm(self):
        # type: () -> _GEN_ASM_FN
        return _GEN_ASMS[self.properties]()

    @staticmethod
    def __clearca_sim(op, state):
        # type: (Op, BaseSimState) -> None
        state[op.outputs[0]] = False,

    @staticmethod
    def __clearca_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        state.writeln("addic 0, 0, 0")
    ClearCA = GenericOpProperties(
        demo_asm="addic 0, 0, 0",
        inputs=[],
        outputs=[OD_CA.with_write_stage(OpStage.Late)],
    )
    _SIM_FNS[ClearCA] = lambda: OpKind.__clearca_sim
    _GEN_ASMS[ClearCA] = lambda: OpKind.__clearca_gen_asm

    @staticmethod
    def __setca_sim(op, state):
        # type: (Op, BaseSimState) -> None
        state[op.outputs[0]] = True,

    @staticmethod
    def __setca_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        state.writeln("subfc 0, 0, 0")
    SetCA = GenericOpProperties(
        demo_asm="subfc 0, 0, 0",
        inputs=[],
        outputs=[OD_CA.with_write_stage(OpStage.Late)],
    )
    _SIM_FNS[SetCA] = lambda: OpKind.__setca_sim
    _GEN_ASMS[SetCA] = lambda: OpKind.__setca_gen_asm

    @staticmethod
    def __svadde_sim(op, state):
        # type: (Op, BaseSimState) -> None
        RA = state[op.input_vals[0]]
        RB = state[op.input_vals[1]]
        carry, = state[op.input_vals[2]]
        VL, = state[op.input_vals[3]]
        RT = []  # type: list[int]
        for i in range(VL):
            v = RA[i] + RB[i] + carry
            RT.append(v & GPR_VALUE_MASK)
            carry = (v >> GPR_SIZE_IN_BITS) != 0
        state[op.outputs[0]] = tuple(RT)
        state[op.outputs[1]] = carry,

    @staticmethod
    def __svadde_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        RT = state.vgpr(op.outputs[0])
        RA = state.vgpr(op.input_vals[0])
        RB = state.vgpr(op.input_vals[1])
        state.writeln(f"sv.adde {RT}, {RA}, {RB}")
    SvAddE = GenericOpProperties(
        demo_asm="sv.adde *RT, *RA, *RB",
        inputs=[OD_EXTRA3_VGPR, OD_EXTRA3_VGPR, OD_CA, OD_VL],
        outputs=[OD_EXTRA3_VGPR, OD_CA.tied_to_input(2)],
    )
    _SIM_FNS[SvAddE] = lambda: OpKind.__svadde_sim
    _GEN_ASMS[SvAddE] = lambda: OpKind.__svadde_gen_asm

    @staticmethod
    def __addze_sim(op, state):
        # type: (Op, BaseSimState) -> None
        RA, = state[op.input_vals[0]]
        carry, = state[op.input_vals[1]]
        v = RA + carry
        RT = v & GPR_VALUE_MASK
        carry = (v >> GPR_SIZE_IN_BITS) != 0
        state[op.outputs[0]] = RT,
        state[op.outputs[1]] = carry,

    @staticmethod
    def __addze_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        RT = state.vgpr(op.outputs[0])
        RA = state.vgpr(op.input_vals[0])
        state.writeln(f"addze {RT}, {RA}")
    AddZE = GenericOpProperties(
        demo_asm="addze RT, RA",
        inputs=[OD_BASE_SGPR, OD_CA],
        outputs=[OD_BASE_SGPR, OD_CA.tied_to_input(1)],
    )
    _SIM_FNS[AddZE] = lambda: OpKind.__addze_sim
    _GEN_ASMS[AddZE] = lambda: OpKind.__addze_gen_asm

    @staticmethod
    def __svsubfe_sim(op, state):
        # type: (Op, BaseSimState) -> None
        RA = state[op.input_vals[0]]
        RB = state[op.input_vals[1]]
        carry, = state[op.input_vals[2]]
        VL, = state[op.input_vals[3]]
        RT = []  # type: list[int]
        for i in range(VL):
            v = (~RA[i] & GPR_VALUE_MASK) + RB[i] + carry
            RT.append(v & GPR_VALUE_MASK)
            carry = (v >> GPR_SIZE_IN_BITS) != 0
        state[op.outputs[0]] = tuple(RT)
        state[op.outputs[1]] = carry,

    @staticmethod
    def __svsubfe_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        RT = state.vgpr(op.outputs[0])
        RA = state.vgpr(op.input_vals[0])
        RB = state.vgpr(op.input_vals[1])
        state.writeln(f"sv.subfe {RT}, {RA}, {RB}")
    SvSubFE = GenericOpProperties(
        demo_asm="sv.subfe *RT, *RA, *RB",
        inputs=[OD_EXTRA3_VGPR, OD_EXTRA3_VGPR, OD_CA, OD_VL],
        outputs=[OD_EXTRA3_VGPR, OD_CA.tied_to_input(2)],
    )
    _SIM_FNS[SvSubFE] = lambda: OpKind.__svsubfe_sim
    _GEN_ASMS[SvSubFE] = lambda: OpKind.__svsubfe_gen_asm

    @staticmethod
    def __svandvs_sim(op, state):
        # type: (Op, BaseSimState) -> None
        RA = state[op.input_vals[0]]
        RB, = state[op.input_vals[1]]
        VL, = state[op.input_vals[2]]
        RT = []  # type: list[int]
        for i in range(VL):
            RT.append(RA[i] & RB & GPR_VALUE_MASK)
        state[op.outputs[0]] = tuple(RT)

    @staticmethod
    def __svandvs_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        RT = state.vgpr(op.outputs[0])
        RA = state.vgpr(op.input_vals[0])
        RB = state.sgpr(op.input_vals[1])
        state.writeln(f"sv.and {RT}, {RA}, {RB}")
    SvAndVS = GenericOpProperties(
        demo_asm="sv.and *RT, *RA, RB",
        inputs=[OD_EXTRA3_VGPR, OD_EXTRA3_SGPR, OD_VL],
        outputs=[OD_EXTRA3_VGPR],
    )
    _SIM_FNS[SvAndVS] = lambda: OpKind.__svandvs_sim
    _GEN_ASMS[SvAndVS] = lambda: OpKind.__svandvs_gen_asm

    @staticmethod
    def __svmaddedu_sim(op, state):
        # type: (Op, BaseSimState) -> None
        RA = state[op.input_vals[0]]
        RB, = state[op.input_vals[1]]
        carry, = state[op.input_vals[2]]
        VL, = state[op.input_vals[3]]
        RT = []  # type: list[int]
        for i in range(VL):
            v = RA[i] * RB + carry
            RT.append(v & GPR_VALUE_MASK)
            carry = v >> GPR_SIZE_IN_BITS
        state[op.outputs[0]] = tuple(RT)
        state[op.outputs[1]] = carry,

    @staticmethod
    def __svmaddedu_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        RT = state.vgpr(op.outputs[0])
        RA = state.vgpr(op.input_vals[0])
        RB = state.sgpr(op.input_vals[1])
        RC = state.sgpr(op.input_vals[2])
        state.writeln(f"sv.maddedu {RT}, {RA}, {RB}, {RC}")
    SvMAddEDU = GenericOpProperties(
        demo_asm="sv.maddedu *RT, *RA, RB, RC",
        inputs=[OD_EXTRA2_VGPR, OD_EXTRA2_SGPR, OD_EXTRA2_SGPR, OD_VL],
        outputs=[OD_EXTRA3_VGPR, OD_EXTRA2_SGPR.tied_to_input(2)],
    )
    _SIM_FNS[SvMAddEDU] = lambda: OpKind.__svmaddedu_sim
    _GEN_ASMS[SvMAddEDU] = lambda: OpKind.__svmaddedu_gen_asm

    @staticmethod
    def __sradi_sim(op, state):
        # type: (Op, BaseSimState) -> None
        rs, = state[op.input_vals[0]]
        imm = op.immediates[0]
        if rs >= 1 << (GPR_SIZE_IN_BITS - 1):
            rs -= 1 << GPR_SIZE_IN_BITS
        v = rs >> imm
        RA = v & GPR_VALUE_MASK
        CA = (RA << imm) != rs
        state[op.outputs[0]] = RA,
        state[op.outputs[1]] = CA,

    @staticmethod
    def __sradi_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        RA = state.sgpr(op.outputs[0])
        RS = state.sgpr(op.input_vals[0])
        imm = op.immediates[0]
        state.writeln(f"sradi {RA}, {RS}, {imm}")
    SRADI = GenericOpProperties(
        demo_asm="sradi RA, RS, imm",
        inputs=[OD_BASE_SGPR],
        outputs=[OD_BASE_SGPR.with_write_stage(OpStage.Late),
                 OD_CA.with_write_stage(OpStage.Late)],
        immediates=[range(GPR_SIZE_IN_BITS)],
    )
    _SIM_FNS[SRADI] = lambda: OpKind.__sradi_sim
    _GEN_ASMS[SRADI] = lambda: OpKind.__sradi_gen_asm

    @staticmethod
    def __setvli_sim(op, state):
        # type: (Op, BaseSimState) -> None
        state[op.outputs[0]] = op.immediates[0],

    @staticmethod
    def __setvli_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        imm = op.immediates[0]
        state.writeln(f"setvl 0, 0, {imm}, 0, 1, 1")
    SetVLI = GenericOpProperties(
        demo_asm="setvl 0, 0, imm, 0, 1, 1",
        inputs=(),
        outputs=[OD_VL.with_write_stage(OpStage.Late)],
        immediates=[range(1, 65)],
        is_load_immediate=True,
    )
    _SIM_FNS[SetVLI] = lambda: OpKind.__setvli_sim
    _GEN_ASMS[SetVLI] = lambda: OpKind.__setvli_gen_asm

    @staticmethod
    def __svli_sim(op, state):
        # type: (Op, BaseSimState) -> None
        VL, = state[op.input_vals[0]]
        imm = op.immediates[0] & GPR_VALUE_MASK
        state[op.outputs[0]] = (imm,) * VL

    @staticmethod
    def __svli_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        RT = state.vgpr(op.outputs[0])
        imm = op.immediates[0]
        state.writeln(f"sv.addi {RT}, 0, {imm}")
    SvLI = GenericOpProperties(
        demo_asm="sv.addi *RT, 0, imm",
        inputs=[OD_VL],
        outputs=[OD_EXTRA3_VGPR],
        immediates=[IMM_S16],
        is_load_immediate=True,
    )
    _SIM_FNS[SvLI] = lambda: OpKind.__svli_sim
    _GEN_ASMS[SvLI] = lambda: OpKind.__svli_gen_asm

    @staticmethod
    def __li_sim(op, state):
        # type: (Op, BaseSimState) -> None
        imm = op.immediates[0] & GPR_VALUE_MASK
        state[op.outputs[0]] = imm,

    @staticmethod
    def __li_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        RT = state.sgpr(op.outputs[0])
        imm = op.immediates[0]
        state.writeln(f"addi {RT}, 0, {imm}")
    LI = GenericOpProperties(
        demo_asm="addi RT, 0, imm",
        inputs=(),
        outputs=[OD_BASE_SGPR.with_write_stage(OpStage.Late)],
        immediates=[IMM_S16],
        is_load_immediate=True,
    )
    _SIM_FNS[LI] = lambda: OpKind.__li_sim
    _GEN_ASMS[LI] = lambda: OpKind.__li_gen_asm

    @staticmethod
    def __veccopytoreg_sim(op, state):
        # type: (Op, BaseSimState) -> None
        state[op.outputs[0]] = state[op.input_vals[0]]

    @staticmethod
    def __copy_to_from_reg_gen_asm(src_loc, dest_loc, is_vec, state):
        # type: (Loc, Loc, bool, GenAsmState) -> None
        sv = "sv." if is_vec else ""
        rev = ""
        if src_loc.conflicts(dest_loc) and src_loc.start < dest_loc.start:
            rev = "/mrr"
        if src_loc == dest_loc:
            return  # no-op
        if src_loc.kind not in (LocKind.GPR, LocKind.StackI64):
            raise ValueError(f"invalid src_loc.kind: {src_loc.kind}")
        if dest_loc.kind not in (LocKind.GPR, LocKind.StackI64):
            raise ValueError(f"invalid dest_loc.kind: {dest_loc.kind}")
        if src_loc.kind is LocKind.StackI64:
            if dest_loc.kind is LocKind.StackI64:
                raise ValueError(
                    f"can't copy from stack to stack: {src_loc} {dest_loc}")
            elif dest_loc.kind is not LocKind.GPR:
                assert_never(dest_loc.kind)
            src = state.stack(src_loc)
            dest = state.gpr(dest_loc, is_vec=is_vec)
            state.writeln(f"{sv}ld {dest}, {src}")
        elif dest_loc.kind is LocKind.StackI64:
            if src_loc.kind is not LocKind.GPR:
                assert_never(src_loc.kind)
            src = state.gpr(src_loc, is_vec=is_vec)
            dest = state.stack(dest_loc)
            state.writeln(f"{sv}std {src}, {dest}")
        elif src_loc.kind is LocKind.GPR:
            if dest_loc.kind is not LocKind.GPR:
                assert_never(dest_loc.kind)
            src = state.gpr(src_loc, is_vec=is_vec)
            dest = state.gpr(dest_loc, is_vec=is_vec)
            state.writeln(f"{sv}or{rev} {dest}, {src}, {src}")
        else:
            assert_never(src_loc.kind)

    @staticmethod
    def __veccopytoreg_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        OpKind.__copy_to_from_reg_gen_asm(
            src_loc=state.loc(
                op.input_vals[0], (LocKind.GPR, LocKind.StackI64)),
            dest_loc=state.loc(op.outputs[0], LocKind.GPR),
            is_vec=True, state=state)

    VecCopyToReg = GenericOpProperties(
        demo_asm="sv.mv dest, src",
        inputs=[GenericOperandDesc(
            ty=GenericTy(BaseTy.I64, is_vec=True),
            sub_kinds=[LocSubKind.SV_EXTRA3_VGPR, LocSubKind.StackI64],
        ), OD_VL],
        outputs=[OD_EXTRA3_VGPR.with_write_stage(OpStage.Late)],
        is_copy=True,
    )
    _SIM_FNS[VecCopyToReg] = lambda: OpKind.__veccopytoreg_sim
    _GEN_ASMS[VecCopyToReg] = lambda: OpKind.__veccopytoreg_gen_asm

    @staticmethod
    def __veccopyfromreg_sim(op, state):
        # type: (Op, BaseSimState) -> None
        state[op.outputs[0]] = state[op.input_vals[0]]

    @staticmethod
    def __veccopyfromreg_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        OpKind.__copy_to_from_reg_gen_asm(
            src_loc=state.loc(op.input_vals[0], LocKind.GPR),
            dest_loc=state.loc(
                op.outputs[0], (LocKind.GPR, LocKind.StackI64)),
            is_vec=True, state=state)
    VecCopyFromReg = GenericOpProperties(
        demo_asm="sv.mv dest, src",
        inputs=[OD_EXTRA3_VGPR, OD_VL],
        outputs=[GenericOperandDesc(
            ty=GenericTy(BaseTy.I64, is_vec=True),
            sub_kinds=[LocSubKind.SV_EXTRA3_VGPR, LocSubKind.StackI64],
            write_stage=OpStage.Late,
        )],
        is_copy=True,
    )
    _SIM_FNS[VecCopyFromReg] = lambda: OpKind.__veccopyfromreg_sim
    _GEN_ASMS[VecCopyFromReg] = lambda: OpKind.__veccopyfromreg_gen_asm

    @staticmethod
    def __copytoreg_sim(op, state):
        # type: (Op, BaseSimState) -> None
        state[op.outputs[0]] = state[op.input_vals[0]]

    @staticmethod
    def __copytoreg_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        OpKind.__copy_to_from_reg_gen_asm(
            src_loc=state.loc(
                op.input_vals[0], (LocKind.GPR, LocKind.StackI64)),
            dest_loc=state.loc(op.outputs[0], LocKind.GPR),
            is_vec=False, state=state)
    CopyToReg = GenericOpProperties(
        demo_asm="mv dest, src",
        inputs=[GenericOperandDesc(
            ty=GenericTy(BaseTy.I64, is_vec=False),
            sub_kinds=[LocSubKind.SV_EXTRA3_SGPR, LocSubKind.BASE_GPR,
                       LocSubKind.StackI64],
        )],
        outputs=[GenericOperandDesc(
            ty=GenericTy(BaseTy.I64, is_vec=False),
            sub_kinds=[LocSubKind.SV_EXTRA3_SGPR, LocSubKind.BASE_GPR],
            write_stage=OpStage.Late,
        )],
        is_copy=True,
    )
    _SIM_FNS[CopyToReg] = lambda: OpKind.__copytoreg_sim
    _GEN_ASMS[CopyToReg] = lambda: OpKind.__copytoreg_gen_asm

    @staticmethod
    def __copyfromreg_sim(op, state):
        # type: (Op, BaseSimState) -> None
        state[op.outputs[0]] = state[op.input_vals[0]]

    @staticmethod
    def __copyfromreg_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        OpKind.__copy_to_from_reg_gen_asm(
            src_loc=state.loc(op.input_vals[0], LocKind.GPR),
            dest_loc=state.loc(
                op.outputs[0], (LocKind.GPR, LocKind.StackI64)),
            is_vec=False, state=state)
    CopyFromReg = GenericOpProperties(
        demo_asm="mv dest, src",
        inputs=[GenericOperandDesc(
            ty=GenericTy(BaseTy.I64, is_vec=False),
            sub_kinds=[LocSubKind.SV_EXTRA3_SGPR, LocSubKind.BASE_GPR],
        )],
        outputs=[GenericOperandDesc(
            ty=GenericTy(BaseTy.I64, is_vec=False),
            sub_kinds=[LocSubKind.SV_EXTRA3_SGPR, LocSubKind.BASE_GPR,
                       LocSubKind.StackI64],
            write_stage=OpStage.Late,
        )],
        is_copy=True,
    )
    _SIM_FNS[CopyFromReg] = lambda: OpKind.__copyfromreg_sim
    _GEN_ASMS[CopyFromReg] = lambda: OpKind.__copyfromreg_gen_asm

    @staticmethod
    def __concat_sim(op, state):
        # type: (Op, BaseSimState) -> None
        state[op.outputs[0]] = tuple(
            state[i][0] for i in op.input_vals[:-1])

    @staticmethod
    def __concat_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        OpKind.__copy_to_from_reg_gen_asm(
            src_loc=state.loc(op.input_vals[0:-1], LocKind.GPR),
            dest_loc=state.loc(op.outputs[0], LocKind.GPR),
            is_vec=True, state=state)
    Concat = GenericOpProperties(
        demo_asm="sv.mv dest, src",
        inputs=[GenericOperandDesc(
            ty=GenericTy(BaseTy.I64, is_vec=False),
            sub_kinds=[LocSubKind.SV_EXTRA3_VGPR],
            spread=True,
        ), OD_VL],
        outputs=[OD_EXTRA3_VGPR.with_write_stage(OpStage.Late)],
        is_copy=True,
    )
    _SIM_FNS[Concat] = lambda: OpKind.__concat_sim
    _GEN_ASMS[Concat] = lambda: OpKind.__concat_gen_asm

    @staticmethod
    def __spread_sim(op, state):
        # type: (Op, BaseSimState) -> None
        for idx, inp in enumerate(state[op.input_vals[0]]):
            state[op.outputs[idx]] = inp,

    @staticmethod
    def __spread_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        OpKind.__copy_to_from_reg_gen_asm(
            src_loc=state.loc(op.input_vals[0], LocKind.GPR),
            dest_loc=state.loc(op.outputs, LocKind.GPR),
            is_vec=True, state=state)
    Spread = GenericOpProperties(
        demo_asm="sv.mv dest, src",
        inputs=[OD_EXTRA3_VGPR, OD_VL],
        outputs=[GenericOperandDesc(
            ty=GenericTy(BaseTy.I64, is_vec=False),
            sub_kinds=[LocSubKind.SV_EXTRA3_VGPR],
            spread=True,
            write_stage=OpStage.Late,
        )],
        is_copy=True,
    )
    _SIM_FNS[Spread] = lambda: OpKind.__spread_sim
    _GEN_ASMS[Spread] = lambda: OpKind.__spread_gen_asm

    @staticmethod
    def __svld_sim(op, state):
        # type: (Op, BaseSimState) -> None
        RA, = state[op.input_vals[0]]
        VL, = state[op.input_vals[1]]
        addr = RA + op.immediates[0]
        RT = []  # type: list[int]
        for i in range(VL):
            v = state.load(addr + GPR_SIZE_IN_BYTES * i)
            RT.append(v & GPR_VALUE_MASK)
        state[op.outputs[0]] = tuple(RT)

    @staticmethod
    def __svld_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        RA = state.sgpr(op.input_vals[0])
        RT = state.vgpr(op.outputs[0])
        imm = op.immediates[0]
        state.writeln(f"sv.ld {RT}, {imm}({RA})")
    SvLd = GenericOpProperties(
        demo_asm="sv.ld *RT, imm(RA)",
        inputs=[OD_EXTRA3_SGPR, OD_VL],
        outputs=[OD_EXTRA3_VGPR],
        immediates=[IMM_S16],
    )
    _SIM_FNS[SvLd] = lambda: OpKind.__svld_sim
    _GEN_ASMS[SvLd] = lambda: OpKind.__svld_gen_asm

    @staticmethod
    def __ld_sim(op, state):
        # type: (Op, BaseSimState) -> None
        RA, = state[op.input_vals[0]]
        addr = RA + op.immediates[0]
        v = state.load(addr)
        state[op.outputs[0]] = v & GPR_VALUE_MASK,

    @staticmethod
    def __ld_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        RA = state.sgpr(op.input_vals[0])
        RT = state.sgpr(op.outputs[0])
        imm = op.immediates[0]
        state.writeln(f"ld {RT}, {imm}({RA})")
    Ld = GenericOpProperties(
        demo_asm="ld RT, imm(RA)",
        inputs=[OD_BASE_SGPR],
        outputs=[OD_BASE_SGPR.with_write_stage(OpStage.Late)],
        immediates=[IMM_S16],
    )
    _SIM_FNS[Ld] = lambda: OpKind.__ld_sim
    _GEN_ASMS[Ld] = lambda: OpKind.__ld_gen_asm

    @staticmethod
    def __svstd_sim(op, state):
        # type: (Op, BaseSimState) -> None
        RS = state[op.input_vals[0]]
        RA, = state[op.input_vals[1]]
        VL, = state[op.input_vals[2]]
        addr = RA + op.immediates[0]
        for i in range(VL):
            state.store(addr + GPR_SIZE_IN_BYTES * i, value=RS[i])

    @staticmethod
    def __svstd_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        RS = state.vgpr(op.input_vals[0])
        RA = state.sgpr(op.input_vals[1])
        imm = op.immediates[0]
        state.writeln(f"sv.std {RS}, {imm}({RA})")
    SvStd = GenericOpProperties(
        demo_asm="sv.std *RS, imm(RA)",
        inputs=[OD_EXTRA3_VGPR, OD_EXTRA3_SGPR, OD_VL],
        outputs=[],
        immediates=[IMM_S16],
        has_side_effects=True,
    )
    _SIM_FNS[SvStd] = lambda: OpKind.__svstd_sim
    _GEN_ASMS[SvStd] = lambda: OpKind.__svstd_gen_asm

    @staticmethod
    def __std_sim(op, state):
        # type: (Op, BaseSimState) -> None
        RS, = state[op.input_vals[0]]
        RA, = state[op.input_vals[1]]
        addr = RA + op.immediates[0]
        state.store(addr, value=RS)

    @staticmethod
    def __std_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        RS = state.sgpr(op.input_vals[0])
        RA = state.sgpr(op.input_vals[1])
        imm = op.immediates[0]
        state.writeln(f"std {RS}, {imm}({RA})")
    Std = GenericOpProperties(
        demo_asm="std RS, imm(RA)",
        inputs=[OD_BASE_SGPR, OD_BASE_SGPR],
        outputs=[],
        immediates=[IMM_S16],
        has_side_effects=True,
    )
    _SIM_FNS[Std] = lambda: OpKind.__std_sim
    _GEN_ASMS[Std] = lambda: OpKind.__std_gen_asm

    @staticmethod
    def __funcargr3_sim(op, state):
        # type: (Op, BaseSimState) -> None
        pass  # return value set before simulation

    @staticmethod
    def __funcargr3_gen_asm(op, state):
        # type: (Op, GenAsmState) -> None
        pass  # no instructions needed
    FuncArgR3 = GenericOpProperties(
        demo_asm="",
        inputs=[],
        outputs=[OD_BASE_SGPR.with_fixed_loc(
            Loc(kind=LocKind.GPR, start=3, reg_len=1))],
    )
    _SIM_FNS[FuncArgR3] = lambda: OpKind.__funcargr3_sim
    _GEN_ASMS[FuncArgR3] = lambda: OpKind.__funcargr3_gen_asm


@plain_data(frozen=True, unsafe_hash=True, repr=False)
class SSAValOrUse(metaclass=InternedMeta):
    __slots__ = "op", "operand_idx"

    def __init__(self, op, operand_idx):
        # type: (Op, int) -> None
        super().__init__()
        self.op = op
        if operand_idx < 0 or operand_idx >= len(self.descriptor_array):
            raise ValueError("invalid operand_idx")
        self.operand_idx = operand_idx

    @abstractmethod
    def __repr__(self):
        # type: () -> str
        ...

    @property
    @abstractmethod
    def descriptor_array(self):
        # type: () -> tuple[OperandDesc, ...]
        ...

    @cached_property
    def defining_descriptor(self):
        # type: () -> OperandDesc
        return self.descriptor_array[self.operand_idx]

    @cached_property
    def ty(self):
        # type: () -> Ty
        return self.defining_descriptor.ty

    @cached_property
    def ty_before_spread(self):
        # type: () -> Ty
        return self.defining_descriptor.ty_before_spread

    @property
    def base_ty(self):
        # type: () -> BaseTy
        return self.ty_before_spread.base_ty

    @property
    def reg_offset_in_unspread(self):
        """ the number of reg-sized slots in the unspread Loc before self's Loc

        e.g. if the unspread Loc containing self is:
        `Loc(kind=LocKind.GPR, start=8, reg_len=4)`
        and self's Loc is `Loc(kind=LocKind.GPR, start=10, reg_len=1)`
        then reg_offset_into_unspread == 2 == 10 - 8
        """
        return self.defining_descriptor.reg_offset_in_unspread

    @property
    def unspread_start_idx(self):
        # type: () -> int
        return self.operand_idx - (self.defining_descriptor.spread_index or 0)

    @property
    def unspread_start(self):
        # type: () -> Self
        return self.__class__(op=self.op, operand_idx=self.unspread_start_idx)


@plain_data(frozen=True, unsafe_hash=True, repr=False)
@final
class SSAVal(SSAValOrUse):
    __slots__ = ()

    def __repr__(self):
        # type: () -> str
        return f"<{self.op.name}.outputs[{self.operand_idx}]: {self.ty}>"

    @cached_property
    def def_loc_set_before_spread(self):
        # type: () -> LocSet
        return self.defining_descriptor.loc_set_before_spread

    @cached_property
    def descriptor_array(self):
        # type: () -> tuple[OperandDesc, ...]
        return self.op.properties.outputs

    @cached_property
    def tied_input(self):
        # type: () -> None | SSAUse
        if self.defining_descriptor.tied_input_index is None:
            return None
        return SSAUse(op=self.op,
                      operand_idx=self.defining_descriptor.tied_input_index)

    @property
    def write_stage(self):
        # type: () -> OpStage
        return self.defining_descriptor.write_stage

    @property
    def current_debugging_value(self):
        # type: () -> tuple[int, ...]
        """ get the current value for debugging in pdb or similar.

        This is intended for use with
        `PreRASimState.set_current_debugging_state`.

        This is only intended for debugging, do not use in unit tests or
        production code.
        """
        return PreRASimState.get_current_debugging_state()[self]

    @cached_property
    def ssa_val_sub_regs(self):
        # type: () -> tuple[SSAValSubReg, ...]
        return tuple(SSAValSubReg(self, i) for i in range(self.ty.reg_len))


@plain_data(frozen=True, unsafe_hash=True, repr=False)
@final
class SSAUse(SSAValOrUse):
    __slots__ = ()

    @cached_property
    def use_loc_set_before_spread(self):
        # type: () -> LocSet
        return self.defining_descriptor.loc_set_before_spread

    @cached_property
    def descriptor_array(self):
        # type: () -> tuple[OperandDesc, ...]
        return self.op.properties.inputs

    def __repr__(self):
        # type: () -> str
        return f"<{self.op.name}.input_uses[{self.operand_idx}]: {self.ty}>"

    @property
    def ssa_val(self):
        # type: () -> SSAVal
        return self.op.input_vals[self.operand_idx]

    @ssa_val.setter
    def ssa_val(self, ssa_val):
        # type: (SSAVal) -> None
        self.op.input_vals[self.operand_idx] = ssa_val


_T = TypeVar("_T")
_Desc = TypeVar("_Desc")


class OpInputSeq(Sequence[_T], Generic[_T, _Desc]):
    @abstractmethod
    def _verify_write_with_desc(self, idx, item, desc):
        # type: (int, _T | Any, _Desc) -> None
        raise NotImplementedError

    @final
    def _verify_write(self, idx, item):
        # type: (int | Any, _T | Any) -> int
        if not isinstance(idx, int):
            if isinstance(idx, slice):
                raise TypeError(
                    f"can't write to slice of {self.__class__.__name__}")
            raise TypeError(f"can't write with index {idx!r}")
        # normalize idx, raising IndexError if it is out of range
        idx = range(len(self.descriptors))[idx]
        desc = self.descriptors[idx]
        self._verify_write_with_desc(idx, item, desc)
        return idx

    def _on_set(self, idx, new_item, old_item):
        # type: (int, _T, _T | None) -> None
        pass

    @abstractmethod
    def _get_descriptors(self):
        # type: () -> tuple[_Desc, ...]
        raise NotImplementedError

    @cached_property
    @final
    def descriptors(self):
        # type: () -> tuple[_Desc, ...]
        return self._get_descriptors()

    @property
    @final
    def op(self):
        return self.__op

    def __init__(self, items, op):
        # type: (Iterable[_T], Op) -> None
        super().__init__()
        self.__op = op
        self.__items = []  # type: list[_T]
        for idx, item in enumerate(items):
            if idx >= len(self.descriptors):
                raise ValueError("too many items")
            _ = self._verify_write(idx, item)
            self.__items.append(item)
        if len(self.__items) < len(self.descriptors):
            raise ValueError("not enough items")

    @final
    def __iter__(self):
        # type: () -> Iterator[_T]
        yield from self.__items

    @overload
    def __getitem__(self, idx):
        # type: (int) -> _T
        ...

    @overload
    def __getitem__(self, idx):
        # type: (slice) -> list[_T]
        ...

    @final
    def __getitem__(self, idx):
        # type: (int | slice) -> _T | list[_T]
        return self.__items[idx]

    @final
    def __setitem__(self, idx, item):
        # type: (int, _T) -> None
        idx = self._verify_write(idx, item)
        self.__items[idx] = item

    @final
    def __len__(self):
        # type: () -> int
        return len(self.__items)

    def __repr__(self):
        # type: () -> str
        return f"{self.__class__.__name__}({self.__items}, op=...)"


@final
class OpInputVals(OpInputSeq[SSAVal, OperandDesc]):
    def _get_descriptors(self):
        # type: () -> tuple[OperandDesc, ...]
        return self.op.properties.inputs

    def _verify_write_with_desc(self, idx, item, desc):
        # type: (int, SSAVal | Any, OperandDesc) -> None
        if not isinstance(item, SSAVal):
            raise TypeError("expected value of type SSAVal")
        if item.ty != desc.ty:
            raise ValueError(f"assigned item's type {item.ty!r} doesn't match "
                             f"corresponding input's type {desc.ty!r}")

    def _on_set(self, idx, new_item, old_item):
        # type: (int, SSAVal, SSAVal | None) -> None
        SSAUses._on_op_input_set(self, idx, new_item, old_item)  # type: ignore

    def __init__(self, items, op):
        # type: (Iterable[SSAVal], Op) -> None
        if hasattr(op, "inputs"):
            raise ValueError("Op.inputs already set")
        super().__init__(items, op)


@final
class OpImmediates(OpInputSeq[int, range]):
    def _get_descriptors(self):
        # type: () -> tuple[range, ...]
        return self.op.properties.immediates

    def _verify_write_with_desc(self, idx, item, desc):
        # type: (int, int | Any, range) -> None
        if not isinstance(item, int):
            raise TypeError("expected value of type int")
        if item not in desc:
            raise ValueError(f"immediate value {item!r} not in {desc!r}")

    def __init__(self, items, op):
        # type: (Iterable[int], Op) -> None
        if hasattr(op, "immediates"):
            raise ValueError("Op.immediates already set")
        super().__init__(items, op)


@plain_data(frozen=True, eq=False, repr=False)
@final
class Op:
    __slots__ = ("fn", "properties", "input_vals", "input_uses", "immediates",
                 "outputs", "name")

    def __init__(self, fn, properties, input_vals, immediates, name=""):
        # type: (Fn, OpProperties, Iterable[SSAVal], Iterable[int], str) -> None
        self.fn = fn
        self.properties = properties
        self.input_vals = OpInputVals(input_vals, op=self)
        inputs_len = len(self.properties.inputs)
        self.input_uses = tuple(SSAUse(self, i) for i in range(inputs_len))
        self.immediates = OpImmediates(immediates, op=self)
        outputs_len = len(self.properties.outputs)
        self.outputs = tuple(SSAVal(self, i) for i in range(outputs_len))
        self.name = fn._add_op_with_unused_name(self, name)  # type: ignore

    @property
    def kind(self):
        # type: () -> OpKind
        return self.properties.kind

    def __eq__(self, other):
        # type: (Op | Any) -> bool
        if isinstance(other, Op):
            return self is other
        return NotImplemented

    def __hash__(self):
        # type: () -> int
        return object.__hash__(self)

    def __repr__(self, wrap_width=63, indent="    "):
        # type: (int, str) -> str
        WRAP_POINT = "\u200B"  # zero-width space
        items = [f"{self.name}:\n"]
        for i, out in enumerate(self.outputs):
            item = f"<...outputs[{i}]: {out.ty}>"
            if i == 0:
                item = "(" + WRAP_POINT + item
            if i != len(self.outputs) - 1:
                item += ", " + WRAP_POINT
            else:
                item += WRAP_POINT + ") <= "
            items.append(item)
        items.append(self.kind._name_)
        if len(self.input_vals) + len(self.immediates) != 0:
            items[-1] += "("
        items[-1] += WRAP_POINT
        for i, inp in enumerate(self.input_vals):
            item = repr(inp)
            if i != len(self.input_vals) - 1 or len(self.immediates) != 0:
                item += ", " + WRAP_POINT
            else:
                item += ") " + WRAP_POINT
            items.append(item)
        for i, imm in enumerate(self.immediates):
            item = hex(imm)
            if i != len(self.immediates) - 1:
                item += ", " + WRAP_POINT
            else:
                item += ") " + WRAP_POINT
            items.append(item)
        lines = []  # type: list[str]
        for i, line_in in enumerate("".join(items).splitlines()):
            if i != 0:
                line_in = indent + line_in
            line_out = ""
            for part in line_in.split(WRAP_POINT):
                if line_out == "":
                    line_out = part
                    continue
                trial_line_out = line_out + part
                if len(trial_line_out.rstrip()) > wrap_width:
                    lines.append(line_out.rstrip())
                    line_out = indent + part
                else:
                    line_out = trial_line_out
            lines.append(line_out.rstrip())
        return "\n".join(lines)

    def sim(self, state):
        # type: (BaseSimState) -> None
        for inp in self.input_vals:
            try:
                val = state[inp]
            except KeyError:
                raise ValueError(f"SSAVal {inp} not yet assigned when "
                                 f"running {self}")
            except SimSkipOp:
                continue
            if len(val) != inp.ty.reg_len:
                raise ValueError(
                    f"value of SSAVal {inp} has wrong number of elements: "
                    f"expected {inp.ty.reg_len} found "
                    f"{len(val)}: {val!r}")
        if isinstance(state, PreRASimState):
            for out in self.outputs:
                if out in state.ssa_vals:
                    if self.kind is OpKind.FuncArgR3:
                        continue
                    raise ValueError(f"SSAVal {out} already assigned before "
                                     f"running {self}")
        try:
            self.kind.sim(self, state)
        except SimSkipOp:
            state.on_skip(self)
        for out in self.outputs:
            try:
                val = state[out]
            except KeyError:
                raise ValueError(f"running {self} failed to assign to {out}")
            except SimSkipOp:
                continue
            if len(val) != out.ty.reg_len:
                raise ValueError(
                    f"value of SSAVal {out} has wrong number of elements: "
                    f"expected {out.ty.reg_len} found "
                    f"{len(val)}: {val!r}")

    def gen_asm(self, state):
        # type: (GenAsmState) -> None
        all_loc_kinds = tuple(LocKind)
        for inp in self.input_vals:
            state.loc(inp, expected_kinds=all_loc_kinds)
        for out in self.outputs:
            state.loc(out, expected_kinds=all_loc_kinds)
        self.kind.gen_asm(self, state)


@plain_data(frozen=True, repr=False)
class BaseSimState(metaclass=ABCMeta):
    __slots__ = "memory",

    def __init__(self, memory):
        # type: (dict[int, int]) -> None
        super().__init__()
        self.memory = memory  # type: dict[int, int]

    def _default_memory_value(self):
        # type: () -> int
        return 0

    def on_skip(self, op):
        # type: (Op) -> None
        raise ValueError("skipping instructions not supported")

    def load_byte(self, addr):
        # type: (int) -> int
        addr &= GPR_VALUE_MASK
        try:
            return self.memory[addr] & 0xFF
        except KeyError:
            return self._default_memory_value()

    def store_byte(self, addr, value):
        # type: (int, int) -> None
        addr &= GPR_VALUE_MASK
        value &= 0xFF
        self.memory[addr] = value

    def load(self, addr, size_in_bytes=GPR_SIZE_IN_BYTES, signed=False):
        # type: (int, int, bool) -> int
        if addr % size_in_bytes != 0:
            raise ValueError(f"address not aligned: {hex(addr)} "
                             f"required alignment: {size_in_bytes}")
        retval = 0
        for i in range(size_in_bytes):
            retval |= self.load_byte(addr + i) << i * BITS_IN_BYTE
        if signed and retval >> (size_in_bytes * BITS_IN_BYTE - 1) != 0:
            retval -= 1 << size_in_bytes * BITS_IN_BYTE
        return retval

    def store(self, addr, value, size_in_bytes=GPR_SIZE_IN_BYTES):
        # type: (int, int, int) -> None
        if addr % size_in_bytes != 0:
            raise ValueError(f"address not aligned: {hex(addr)} "
                             f"required alignment: {size_in_bytes}")
        for i in range(size_in_bytes):
            self.store_byte(addr + i, (value >> i * BITS_IN_BYTE) & 0xFF)

    def _memory__repr(self):
        # type: () -> str
        if len(self.memory) == 0:
            return "{}"
        keys = sorted(self.memory.keys(), reverse=True)
        CHUNK_SIZE = GPR_SIZE_IN_BYTES
        items = []  # type: list[str]
        while len(keys) != 0:
            addr = keys[-1]
            if (len(keys) >= CHUNK_SIZE
                    and addr % CHUNK_SIZE == 0
                    and keys[-CHUNK_SIZE:]
                    == list(reversed(range(addr, addr + CHUNK_SIZE)))):
                value = self.load(addr, size_in_bytes=CHUNK_SIZE)
                items.append(f"0x{addr:05x}: <0x{value:0{CHUNK_SIZE * 2}x}>")
                keys[-CHUNK_SIZE:] = ()
            else:
                items.append(f"0x{addr:05x}: 0x{self.memory[keys.pop()]:02x}")
        if len(items) == 1:
            return f"{{{items[0]}}}"
        items_str = ",\n".join(items)
        return f"{{\n{items_str}}}"

    def __repr__(self):
        # type: () -> str
        field_vals = []  # type: list[str]
        for name in fields(self):
            try:
                value = getattr(self, name)
            except AttributeError:
                field_vals.append(f"{name}=<not set>")
                continue
            repr_fn = getattr(self, f"_{name}__repr", None)
            if callable(repr_fn):
                field_vals.append(f"{name}={repr_fn()}")
            else:
                field_vals.append(f"{name}={value!r}")
        field_vals_str = ", ".join(field_vals)
        return f"{self.__class__.__name__}({field_vals_str})"

    @abstractmethod
    def __getitem__(self, ssa_val):
        # type: (SSAVal) -> tuple[int, ...]
        ...

    @abstractmethod
    def __setitem__(self, ssa_val, value):
        # type: (SSAVal, Iterable[int]) -> None
        ...


@plain_data(frozen=True, repr=False)
class PreRABaseSimState(BaseSimState):
    __slots__ = "ssa_vals",

    def __init__(self, ssa_vals, memory):
        # type: (dict[SSAVal, tuple[int, ...]], dict[int, int]) -> None
        super().__init__(memory)
        self.ssa_vals = ssa_vals  # type: dict[SSAVal, tuple[int, ...]]

    def _ssa_vals__repr(self):
        # type: () -> str
        if len(self.ssa_vals) == 0:
            return "{}"
        items = []  # type: list[str]
        CHUNK_SIZE = 4
        for k, v in self.ssa_vals.items():
            element_strs = []  # type: list[str]
            for i, el in enumerate(v):
                if i % CHUNK_SIZE != 0:
                    element_strs.append(" " + hex(el))
                else:
                    element_strs.append("\n    " + hex(el))
            if len(element_strs) <= CHUNK_SIZE:
                element_strs[0] = element_strs[0].lstrip()
            if len(element_strs) == 1:
                element_strs.append("")
            v_str = ",".join(element_strs)
            items.append(f"{k!r}: ({v_str})")
        if len(items) == 1 and "\n" not in items[0]:
            return f"{{{items[0]}}}"
        items_str = ",\n".join(items)
        return f"{{\n{items_str},\n}}"

    def __getitem__(self, ssa_val):
        # type: (SSAVal) -> tuple[int, ...]
        try:
            return self.ssa_vals[ssa_val]
        except KeyError:
            return self._handle_undefined_ssa_val(ssa_val)

    def _handle_undefined_ssa_val(self, ssa_val):
        # type: (SSAVal) -> tuple[int, ...]
        raise KeyError("SSAVal has no value set", ssa_val)

    def __setitem__(self, ssa_val, value):
        # type: (SSAVal, Iterable[int]) -> None
        value = tuple(map(int, value))
        if len(value) != ssa_val.ty.reg_len:
            raise ValueError("value has wrong len")
        self.ssa_vals[ssa_val] = value


class SimSkipOp(Exception):
    pass


@plain_data(frozen=True, repr=False)
@final
class ConstPropagationState(PreRABaseSimState):
    __slots__ = "skipped_ops",

    def __init__(self, ssa_vals, memory, skipped_ops):
        # type: (dict[SSAVal, tuple[int, ...]], dict[int, int], OSet[Op]) -> None
        super().__init__(ssa_vals, memory)
        self.skipped_ops = skipped_ops

    def _default_memory_value(self):
        # type: () -> int
        raise SimSkipOp

    def _handle_undefined_ssa_val(self, ssa_val):
        # type: (SSAVal) -> tuple[int, ...]
        raise SimSkipOp

    def on_skip(self, op):
        # type: (Op) -> None
        self.skipped_ops.add(op)


@plain_data(frozen=True, repr=False)
class PreRASimState(PreRABaseSimState):
    __slots__ = ()

    __CURRENT_DEBUGGING_STATE = []  # type: list[PreRASimState]

    @contextmanager
    def set_as_current_debugging_state(self):
        """ return a context manager that sets self as the current state for
        debugging in pdb or similar. This is intended only for use with
        `get_current_debugging_state` which should not be used in unit tests
        or production code.
        """
        try:
            PreRASimState.__CURRENT_DEBUGGING_STATE.append(self)
            yield
        finally:
            assert self is PreRASimState.__CURRENT_DEBUGGING_STATE.pop(), \
                "inconsistent __CURRENT_DEBUGGING_STATE"

    @staticmethod
    def get_current_debugging_state():
        # type: () -> PreRASimState
        """ get the current state for debugging in pdb or similar.

        This is intended for use with `set_current_debugging_state`.

        This is only intended for debugging, do not use in unit tests or
        production code.
        """
        if len(PreRASimState.__CURRENT_DEBUGGING_STATE) == 0:
            raise ValueError("no current debugging state")
        return PreRASimState.__CURRENT_DEBUGGING_STATE[-1]


@plain_data(frozen=True, repr=False)
@final
class PostRASimState(BaseSimState):
    __slots__ = "ssa_val_to_loc_map", "loc_values"

    def __init__(self, ssa_val_to_loc_map, memory, loc_values):
        # type: (dict[SSAVal, Loc], dict[int, int], dict[Loc, int]) -> None
        super().__init__(memory)
        self.ssa_val_to_loc_map = FMap(ssa_val_to_loc_map)
        for ssa_val, loc in self.ssa_val_to_loc_map.items():
            if ssa_val.ty != loc.ty:
                raise ValueError(
                    f"type mismatch for SSAVal and Loc: {ssa_val} {loc}")
        self.loc_values = loc_values
        for loc in self.loc_values.keys():
            if loc.reg_len != 1:
                raise ValueError(
                    "loc_values must only contain Locs with reg_len=1, all "
                    "larger Locs will be split into reg_len=1 sub-Locs")

    def _loc_values__repr(self):
        # type: () -> str
        locs = sorted(self.loc_values.keys(),
                      key=lambda v: (v.kind.name, v.start))
        items = []  # type: list[str]
        for loc in locs:
            items.append(f"{loc}: 0x{self.loc_values[loc]:x}")
        items_str = ",\n".join(items)
        return f"{{\n{items_str},\n}}"

    def __getitem__(self, ssa_val):
        # type: (SSAVal) -> tuple[int, ...]
        loc = self.ssa_val_to_loc_map[ssa_val]
        subloc_ty = Ty(base_ty=loc.ty.base_ty, reg_len=1)
        retval = []  # type: list[int]
        for i in range(loc.reg_len):
            subloc = loc.get_subloc_at_offset(subloc_ty=subloc_ty, offset=i)
            retval.append(self.loc_values.get(subloc, 0))
        return tuple(retval)

    def __setitem__(self, ssa_val, value):
        # type: (SSAVal, Iterable[int]) -> None
        value = tuple(map(int, value))
        if len(value) != ssa_val.ty.reg_len:
            raise ValueError("value has wrong len")
        loc = self.ssa_val_to_loc_map[ssa_val]
        subloc_ty = Ty(base_ty=loc.ty.base_ty, reg_len=1)
        for i in range(loc.reg_len):
            subloc = loc.get_subloc_at_offset(subloc_ty=subloc_ty, offset=i)
            self.loc_values[subloc] = value[i]


@plain_data(frozen=True)
class GenAsmState:
    __slots__ = "allocated_locs", "output"

    def __init__(self, allocated_locs, output=None):
        # type: (Mapping[SSAVal, Loc], StringIO | list[str] | None) -> None
        super().__init__()
        self.allocated_locs = FMap(allocated_locs)
        for ssa_val, loc in self.allocated_locs.items():
            if ssa_val.ty != loc.ty:
                raise ValueError(
                    f"Ty mismatch: ssa_val.ty:{ssa_val.ty} != loc.ty:{loc.ty}")
        if output is None:
            output = []
        self.output = output

    __SSA_VAL_OR_LOCS = Union[SSAVal, Loc, Sequence["SSAVal | Loc"]]

    def loc(self, ssa_val_or_locs, expected_kinds):
        # type: (__SSA_VAL_OR_LOCS, LocKind | tuple[LocKind, ...]) -> Loc
        if isinstance(ssa_val_or_locs, (SSAVal, Loc)):
            ssa_val_or_locs = [ssa_val_or_locs]
        locs = []  # type: list[Loc]
        for i in ssa_val_or_locs:
            if isinstance(i, SSAVal):
                locs.append(self.allocated_locs[i])
            else:
                locs.append(i)
        if len(locs) == 0:
            raise ValueError("invalid Loc sequence: must not be empty")
        retval = locs[0].try_concat(*locs[1:])
        if retval is None:
            raise ValueError("invalid Loc sequence: try_concat failed")
        if isinstance(expected_kinds, LocKind):
            expected_kinds = expected_kinds,
        if retval.kind not in expected_kinds:
            if len(expected_kinds) == 1:
                expected_kinds = expected_kinds[0]
            raise ValueError(f"LocKind mismatch: {ssa_val_or_locs}: found "
                             f"{retval.kind} expected {expected_kinds}")
        return retval

    def gpr(self, ssa_val_or_locs, is_vec):
        # type: (__SSA_VAL_OR_LOCS, bool) -> str
        loc = self.loc(ssa_val_or_locs, LocKind.GPR)
        vec_str = "*" if is_vec else ""
        return vec_str + str(loc.start)

    def sgpr(self, ssa_val_or_locs):
        # type: (__SSA_VAL_OR_LOCS) -> str
        return self.gpr(ssa_val_or_locs, is_vec=False)

    def vgpr(self, ssa_val_or_locs):
        # type: (__SSA_VAL_OR_LOCS) -> str
        return self.gpr(ssa_val_or_locs, is_vec=True)

    def stack(self, ssa_val_or_locs):
        # type: (__SSA_VAL_OR_LOCS) -> str
        loc = self.loc(ssa_val_or_locs, LocKind.StackI64)
        return f"{loc.start}(1)"

    def writeln(self, *line_segments):
        # type: (*str) -> None
        line = " ".join(line_segments)
        if isinstance(self.output, list):
            self.output.append(line)
        else:
            self.output.write(line + "\n")