Interpreter

class ReturnedValues(NamedTuple):
    """
    If the terminator exits the region being interpreted, such as for function returns,
    return these values to yield control back to the parent operation of the block.
    """

    values: PythonValues

class Successor(NamedTuple):
    """
    If the terminator jumps to another block within the region being interpreted, such
    as for jumps, return a block along with the values to pass as arguments.
    """

    block: Block
    args: PythonValues

class OpImplResult(NamedTuple):
    """
    The result of interpreting an Operation. If and only if the Operation is a terminator,
    it must set the terminator_value.
    """

    values: PythonValues
    terminator_value: TerminatorValue | None

@dataclass
class InterpreterFunctions:
    """
    Holds the Python implementations for Operations. Users should
    subclass this class, and define the functions to run during interpretation.
    For example:

    ``` python
    @register_impls
    class ArithFunctions(InterpreterFunctions):

        @impl(arith.Addi)
        def run_addi(self, interpreter: Interpreter, op: arith.Addi,
                     args: tuple[Any, ...]) -> tuple[Any, ...]:
            lhs, rhs = args
            return lhs + rhs,
    ```

    The interpreter will take care of fetching the Python values associated
    with the operand SSAValues, and setting the return values to the
    appropriate OpResults.

    To override the definition of an operation implementation, subclass the
    class to override, and redefine the functions, annotating them with
    `@impl`.

    ``` python
    @register_impls
    class DebugArithFunctions(ArithFunctions):

        @impl(arith.Addi)
        def run_addi(self, interpreter: Interpreter, op: arith.Addi,
                     args: tuple[Any, ...]) -> tuple[Any, ...]:
            lhs, rhs = args
            print(lhs, rhs, lhs + rhs)
            return lhs + rhs,
    ```

    To register an implementation of a cast for UnrealizedConversionCastOp, use
    `impl_cast`, like so:

    ``` python
    @register_impls
    class ArithFunctions(InterpreterFunctions):
        @impl_cast(IntegerType, IndexType)
        def cast_integer_to_index(
            self,
            input_type: IntegerType,
            output_type: IndexType,
            value: Any,
        ) -> Any:
            # Both input and output represented by a Python `int`
            return value
    ```
    """

    @classmethod
    def _impls(
        cls,
    ) -> Iterable[tuple[type[Operation], OpImpl[InterpreterFunctions, Operation]]]:
        try:
            impl_dict = getattr(cls, _IMPL_DICT_KEY)
            return impl_dict.items()
        except AttributeError as e:
            raise ValueError(f"Use `@register_impls` on class {cls.__name__}") from e

    @classmethod
    def _cast_impls(
        cls,
    ) -> Iterable[
        tuple[
            tuple[type[Attribute], type[Attribute]],
            CastImpl[InterpreterFunctions, Attribute, Attribute],
        ]
    ]:
        try:
            impl_dict = getattr(cls, _CAST_IMPL_DICT_KEY)
            return impl_dict.items()
        except AttributeError as e:
            raise ValueError(f"Use `@register_impls` on class {cls.__name__}") from e

    @classmethod
    def _attr_impls(
        cls,
    ) -> Iterable[
        tuple[
            type[Attribute],
            AttrImpl[InterpreterFunctions, Attribute],
        ]
    ]:
        try:
            impl_dict = getattr(cls, _ATTR_IMPL_DICT_KEY)
            return impl_dict.items()
        except AttributeError as e:
            raise ValueError(f"Use `@register_impls` on class {cls.__name__}") from e

    @classmethod
    def _ext_impls(
        cls,
    ) -> Iterable[tuple[str, ExtFuncImpl[InterpreterFunctions]]]:
        try:
            impl_dict = getattr(cls, _EXT_FUNC_DICT_KEY)
            return impl_dict.items()
        except AttributeError as e:
            raise ValueError(f"Use `@register_impls` on class {cls.__name__}") from e

    @classmethod
    def _callable_impls(
        cls,
    ) -> Iterable[tuple[type[Operation], OpImpl[InterpreterFunctions, Operation]]]:
        try:
            impl_dict = getattr(cls, _CALLABLE_IMPL_DICT_KEY)
            return impl_dict.items()
        except AttributeError as e:
            raise ValueError(f"Use `@register_impls` on class {cls.__name__}") from e

@dataclass
class Interpreter:
    """
    An extensible interpreter, initialised with a Module to interpret. The
    implementation for each Operation subclass should be provided via a
    `InterpretationFunctions` instance. Interpretations can be overridden, and
    the override must be specified explicitly, by passing `override=True` to
    the `register_functions` method.
    """

    class Listener:
        """
        Base class for observing the operations that are interpreted during a run.
        """

        def will_interpret_op(self, op: Operation, args: PythonValues) -> None: ...

        def did_interpret_op(self, op: Operation, results: PythonValues) -> None: ...

    SYSTEM_BITWIDTH: ClassVar[Literal[32, 64] | None] = _get_system_bitwidth()
    DEFAULT_BITWIDTH: ClassVar[Literal[32, 64]] = (
        32 if SYSTEM_BITWIDTH is None else SYSTEM_BITWIDTH
    )

    module: ModuleOp
    index_bitwidth: Literal[32, 64] = field(default=DEFAULT_BITWIDTH)
    """
    Number of bits in the binary representation of the index
    """
    _impls: _InterpreterFunctionImpls = field(default_factory=_InterpreterFunctionImpls)
    _ctx: ScopedDict[SSAValue, Any] = field(
        default_factory=lambda: ScopedDict[SSAValue, Any](name="root")
    )
    """
    Object holding the Python values associated with SSAValues during an
    interpretation context. An environment is a stack of scopes, values are
    assigned to the current scope, but can be fetched from a parent scope.
    """
    file: IO[str] | None = field(default=None)
    _symbol_table: dict[str, Operation] | None = None
    _impl_data: _IMPL_DATA = field(default_factory=_IMPL_DATA)
    """
    Runtime data associated with an interpreter functions implementation.
    """
    listeners: tuple[Listener, ...] = field(default=())

    def get_values(self, values: Iterable[SSAValue]) -> tuple[Any, ...]:
        """
        Get values from current environment.
        """
        return tuple(self._ctx[value] for value in values)

    def set_values(self, pairs: Iterable[tuple[SSAValue, Any]]):
        """
        Set values to current scope.
        Raises InterpretationError if len(ssa_values) != len(result_values), or
        if SSA value already has a Python value in the current scope.
        """
        for ssa_value, result_value in pairs:
            self._ctx[ssa_value] = result_value

    def push_scope(self, name: str | None = None) -> None:
        """
        Create new scope in current environment, with optional custom `name`.
        """
        self._ctx = ScopedDict(name=name, parent=self._ctx)

    def pop_scope(self) -> None:
        """
        Discard the current scope, and all the values registered in it. Sets
        parent scope of current scope to new current scope.
        Raises InterpretationError if current scope is root scope.
        """
        if self._ctx.parent is None:
            raise InterpretationError("Attempting to pop root env")

        self._ctx = self._ctx.parent

    def register_implementations(
        self, impls: InterpreterFunctions, /, override: bool = False
    ) -> None:
        """
        Register implementations for operations defined in given
        `InterpreterFunctions` object. Raise InterpretationError if an
        operation already has an implementation registered, unless override is
        set to True.
        """
        self._impls.register_from(impls, override=override)

    def _run_op(self, op: Operation, inputs: PythonValues) -> OpImplResult:
        if (operands_count := len(op.operands)) != (inputs_count := len(inputs)):
            raise InterpretationError(
                f"Number of operands ({operands_count}) doesn't match the number of inputs ({inputs_count})."
            )
        for listener in self.listeners:
            listener.will_interpret_op(op, inputs)
        result = self._impls.run(self, op, inputs)
        if (results_count := len(op.results)) != (
            actual_result_count := len(result.values)
        ):
            raise InterpretationError(
                f"Number of operation results ({results_count}) doesn't match the "
                f"number of implementation results ({actual_result_count})."
            )
        for listener in self.listeners:
            listener.did_interpret_op(op, result.values)
        return result

    def run_op(
        self, op: Operation | str | SymbolRefAttr, inputs: PythonValues = ()
    ) -> PythonValues:
        """
        Calls the implementation for the given operation.
        """
        if not isinstance(op, Operation):
            op = self.get_op_for_symbol(op)

        return self._run_op(op, inputs).values

    def call_op(
        self, op: Operation | str | SymbolRefAttr, inputs: PythonValues = ()
    ) -> PythonValues:
        """
        Calls the implementation for the given operation.
        """
        if not isinstance(op, Operation):
            op = self.get_op_for_symbol(op)
        results = self._impls.call(self, op, inputs)
        return results

    def call_external(
        self, sym_name: str, op: Operation, inputs: PythonValues = ()
    ) -> PythonValues:
        return self._impls.call_external(self, sym_name, op, inputs)

    def run_ssacfg_region(
        self, region: Region, args: PythonValues, name: str = "unknown"
    ) -> PythonValues:
        """
        Interpret an SSACFG-semantic Region.
        Creates a new scope, then executes the first block in the region. The first block
        is expected to return the results of the region directly.
        """
        results = ()
        if not region.blocks:
            return results

        initial_scope = self._ctx
        block = region.blocks.first

        while block is not None:
            self.push_scope(name)
            self.set_values(zip(block.args, args))

            op: Operation | None = block.first_op
            block = None

            while op is not None:
                inputs = self.get_values(op.operands)
                result = self._run_op(op, inputs)
                self.interpreter_assert(
                    len(op.results) == len(result.values),
                    f"Incorrect number of results for op {op.name}, expected {len(op.results)} but got {len(result.values)}",
                )
                self.set_values(zip(op.results, result.values))

                if result.terminator_value is not None:
                    match result.terminator_value:
                        case ReturnedValues():
                            # update results and break out of outer loop
                            results = result.terminator_value.values
                            break
                        case Successor():
                            # block won't be None, so only break out of inner loop
                            block, args = result.terminator_value
                            break

                # Set up next iteration
                op = op.next_op

        self._ctx = initial_scope
        return results

    def cast_value(self, o: Attribute, r: Attribute, value: Any) -> Any:
        """
        If the type of the operand and result are not the same, then look up the
        user-provided conversion function.
        """
        if o == r:
            return value

        return self._impls.cast(o, r, value)

    def value_for_attribute(self, attr: Attribute, type_attr: Attribute) -> Any:
        return self._impls.attr_value(self, attr, type_attr)

    def get_op_for_symbol(self, symbol: str | SymbolRefAttr) -> Operation:
        op = SymbolTable.lookup_symbol(self.module, symbol)
        if op is not None:
            return op
        raise InterpretationError(f"Could not find symbol {symbol}")

    def get_data(
        self,
        functions: type[InterpreterFunctions],
        key: str,
        factory: Callable[[], Any],
    ) -> Any:
        """
        Get data associated with a specific interpreter functions class, with a given key.
        If the data is missing, the `factory` argument will be executed and stored on the
        interpreter.
        """
        if functions not in self._impl_data:
            functions_data: dict[str, Any] = {}
            self._impl_data[functions] = functions_data
        else:
            functions_data = self._impl_data[functions]

        if key not in functions_data:
            data = factory()
            functions_data[key] = data
        else:
            data = functions_data[key]

        return data

    def set_data(
        self,
        functions: type[InterpreterFunctions],
        key: str,
        value: Any,
    ):
        if functions not in self._impl_data:
            functions_data: dict[str, Any] = {}
            self._impl_data[functions] = functions_data
        else:
            functions_data = self._impl_data[functions]
        functions_data[key] = value

    def print(self, *args: Any, **kwargs: Any):
        """Print to current file."""
        print(*args, **kwargs, file=self.file)

    def interpreter_assert(self, condition: bool, message: str | None = None):
        """Raise InterpretationError if condition is not satisfied."""
        if not condition:
            self.raise_error(message)

    def scope_names(self):
        ctx = self._ctx

        while ctx is not None:
            yield ctx.name or "unknown"
            ctx = ctx.parent

    def raise_error(self, message: str | None = None):
        scope_description = "/".join(self.scope_names())
        raise InterpretationError(f"AssertionError: ({scope_description})({message})")

@dataclass
class OpCounter(Interpreter.Listener):
    """
    Counts the number of times that an op has been run by the interpreter.
    """

    ops: Counter[str] = field(default_factory=Counter[str])

    def will_interpret_op(self, op: Operation, args: PythonValues) -> None:
        self.ops[op.name] += 1