import typing as t

import attrs

from puya.ir import models
from puya.ir.visitor import IRVisitor

@attrs.define
class IRMutator(IRVisitor[t.Any]):
    def visit_block(self, block: models.BasicBlock) -> None:
        new_phis = []
        for phi in block.phis:
            new_phi = self.visit_phi(phi)
            if new_phi is not None:
                new_phis.append(new_phi)
        block.phis = new_phis
        new_ops = []
        for op in block.ops:
            new_op = op.accept(self)
            if new_op is not None:
                new_ops.append(new_op)
        block.ops = new_ops
        if block.terminator is not None:
            block.terminator = block.terminator.accept(self)

    def visit_assignment(self, ass: models.Assignment) -> models.Assignment | None:
        ass.targets = [self.visit_register(r) for r in ass.targets]
        ass.source = ass.source.accept(self)
        return ass

    def visit_register(self, reg: models.Register) -> models.Register:
        return reg

    def visit_undefined(self, val: models.Undefined) -> models.Undefined:
        return val

    def visit_template_var(self, deploy_var: models.TemplateVar) -> models.TemplateVar:
        return deploy_var

    def visit_uint64_constant(self, const: models.UInt64Constant) -> models.UInt64Constant:
        return const

    def visit_biguint_constant(self, const: models.BigUIntConstant) -> models.BigUIntConstant:
        return const

    def visit_bytes_constant(self, const: models.BytesConstant) -> models.BytesConstant:
        return const

    def visit_address_constant(self, const: models.AddressConstant) -> models.AddressConstant:
        return const

    def visit_method_constant(self, const: models.MethodConstant) -> models.MethodConstant:
        return const

    def visit_itxn_constant(self, const: models.ITxnConstant) -> models.ITxnConstant:
        return const

    def visit_compiled_contract_reference(
        self, const: models.CompiledContractReference
    ) -> models.CompiledContractReference:
        return attrs.evolve(
            const,
            template_variables={
                var: value.accept(self) for var, value in const.template_variables.items()
            },
        )

    def visit_compiled_logicsig_reference(
        self, const: models.CompiledLogicSigReference
    ) -> models.CompiledLogicSigReference:
        return attrs.evolve(
            const,
            template_variables={
                var: value.accept(self) for var, value in const.template_variables.items()
            },
        )

    def visit_phi(self, phi: models.Phi) -> models.Phi | None:
        phi.register = self.visit_register(phi.register)
        phi.args = [self.visit_phi_argument(a) for a in phi.args]
        return phi

    def visit_phi_argument(self, arg: models.PhiArgument) -> models.PhiArgument:
        arg.value = arg.value.accept(self)
        return arg

    def visit_intrinsic_op(self, intrinsic: models.Intrinsic) -> models.Intrinsic | None:
        intrinsic.args = [a.accept(self) for a in intrinsic.args]
        return intrinsic

    def visit_inner_transaction_field(
        self, field: models.InnerTransactionField
    ) -> models.InnerTransactionField | models.Intrinsic:
        field.group_index = field.group_index.accept(self)
        field.is_last_in_group = field.is_last_in_group.accept(self)
        if field.array_index:
            field.array_index = field.array_index.accept(self)
        return field

    def visit_invoke_subroutine(self, callsub: models.InvokeSubroutine) -> models.InvokeSubroutine:
        callsub.args = [a.accept(self) for a in callsub.args]
        return callsub

    def visit_conditional_branch(self, branch: models.ConditionalBranch) -> models.ControlOp:
        branch.condition = branch.condition.accept(self)
        return branch

    def visit_goto_nth(self, goto_nth: models.GotoNth) -> models.ControlOp:
        goto_nth.value = goto_nth.value.accept(self)
        return goto_nth

    def visit_goto(self, goto: models.Goto) -> models.ControlOp:
        return goto

    def visit_switch(self, switch: models.Switch) -> models.ControlOp:
        switch.value = switch.value.accept(self)
        switch.cases = {case.accept(self): target for case, target in switch.cases.items()}
        return switch

    def visit_subroutine_return(self, retsub: models.SubroutineReturn) -> models.ControlOp:
        retsub.result = [v.accept(self) for v in retsub.result]
        return retsub

    def visit_program_exit(self, exit_: models.ProgramExit) -> models.ControlOp:
        exit_.result = exit_.result.accept(self)
        return exit_

    def visit_fail(self, fail: models.Fail) -> models.ControlOp:
        return fail

    def visit_value_tuple(self, tup: models.ValueTuple) -> models.ValueTuple:
        tup.values = [v.accept(self) for v in tup.values]
        return tup


import itertools
import typing
from collections.abc import Sequence, Set

import attrs

from puya.errors import InternalError
from puya.ir import models as ops
from puya.ir.visitor import IRTraverser
from puya.utils import StableSet

IrOp: typing.TypeAlias = ops.Op | ops.ControlOp | ops.Phi

@attrs.define(kw_only=True)
class _OpLifetime:
    block: ops.BasicBlock
    used: StableSet[ops.Register] = attrs.field(on_setattr=attrs.setters.frozen)
    defined: StableSet[ops.Register] = attrs.field(on_setattr=attrs.setters.frozen)
    successors: Sequence[IrOp] = attrs.field(on_setattr=attrs.setters.frozen)

    live_in: StableSet[ops.Register] = attrs.field(factory=StableSet)
    live_out: StableSet[ops.Register] = attrs.field(factory=StableSet)

@attrs.define
class _VlaTraverser(IRTraverser):
    used: StableSet[ops.Register] = attrs.field(factory=StableSet)
    defined: StableSet[ops.Register] = attrs.field(factory=StableSet)

    @classmethod
    def apply(cls, op: IrOp) -> tuple[StableSet[ops.Register], StableSet[ops.Register]]:
        traverser = cls()
        op.accept(traverser)
        return traverser.used, traverser.defined

    def visit_register(self, reg: ops.Register) -> None:
        self.used.add(reg)

    def visit_assignment(self, ass: ops.Assignment) -> None:
        ass.source.accept(self)
        self.defined = StableSet.from_iter(ass.targets)

    def visit_phi(self, _phi: ops.Phi) -> None:
        # WARNING: this is slightly trickier than it might seem for SSA,
        #          consider how this translates when coming out of SSA -
        #          the target register isn't defined here, but defined at
        #          the end of each predecessor block.
        #          Similarly, the arguments aren't live-in at this location necessarily.
        raise InternalError("IR VLA not capable of handling SSA yet")

@attrs.define
class VariableLifetimeAnalysis:
    """Performs VLA analysis for a subroutine, providing a mapping of ops to sets of live local_ids
    see https://www.classes.cs.uchicago.edu/archive/2004/spring/22620-1/docs/liveness.pdf"""

    subroutine: ops.Subroutine
    _op_lifetimes: dict[IrOp, _OpLifetime] = attrs.field(init=False)

    @_op_lifetimes.default
    def _op_lifetimes_factory(self) -> dict[IrOp, _OpLifetime]:
        result = dict[IrOp, _OpLifetime]()
        block_map = {b.id: next(b.all_ops) for b in self.subroutine.body}
        for block in self.subroutine.body:
            assert not block.phis
            all_ops = list(block.all_ops)
            for op, next_op in itertools.zip_longest(all_ops, all_ops[1:]):
                used, defined = _VlaTraverser.apply(op)
                if next_op is None:
                    # for last op, add first op of each successor block
                    successors = [block_map[s.id] for s in block.successors]
                else:
                    successors = [next_op]
                result[op] = _OpLifetime(
                    block=block,
                    used=used,
                    defined=defined,
                    successors=successors,
                )
        return result

    def get_live_out_variables(self, op: IrOp) -> Set[ops.Register]:
        return self._op_lifetimes[op].live_out

    def get_live_in_variables(self, op: IrOp) -> Set[ops.Register]:
        return self._op_lifetimes[op].live_in

    @classmethod
    def analyze(cls, subroutine: ops.Subroutine) -> typing.Self:
        analysis = cls(subroutine)
        analysis._analyze()  # noqa: SLF001
        return analysis

    def _analyze(self) -> None:
        changes = True
        while changes:
            changes = False
            for n in self._op_lifetimes.values():
                # For OUT, find out the union of previous variables
                # in the IN set for each succeeding node of n.

                # out[n] = U s ∈ succ[n] in[s]
                live_out = StableSet[ops.Register]()
                for s in n.successors:
                    live_out |= self._op_lifetimes[s].live_in

                # in[n] = use[n] U (out[n] - def [n])
                live_in = n.used | (live_out - n.defined)

                if not (live_in == n.live_in and live_out == n.live_out):
                    n.live_in = live_in
                    n.live_out = live_out
                    changes = True


import contextlib
import json
import logging
import os.path
import sys
import typing
from collections import Counter
from collections.abc import Iterator, Mapping, Sequence
from contextvars import ContextVar
from enum import IntEnum, StrEnum, auto
from io import StringIO
from pathlib import Path

import attrs
import structlog

from puya.parse import SourceLocation

class LogFormat(StrEnum):
    default = auto()
    json = auto()

    @staticmethod
    def from_string(s: str) -> "LogFormat":
        try:
            return LogFormat[s]
        except KeyError as err:
            raise ValueError from err

class LogLevel(IntEnum):
    notset = logging.NOTSET
    debug = logging.DEBUG
    info = logging.INFO
    warning = logging.WARNING
    error = logging.ERROR
    critical = logging.CRITICAL

    def __str__(self) -> str:
        return self.name

    @staticmethod
    def from_string(s: str) -> "LogLevel":
        try:
            return LogLevel[s]
        except KeyError as err:
            raise ValueError from err

@attrs.frozen
class Log:
    level: LogLevel
    message: str
    location: SourceLocation | None

    @property
    def file(self) -> Path | None:
        return None if self.location is None else self.location.file

    @property
    def line(self) -> int | None:
        return None if self.location is None else self.location.line

@attrs.define
class LoggingContext:
    logs: list[Log] = attrs.field(factory=list)
    sources_by_path: Mapping[Path, Sequence[str] | None] | None = None

    def _log_level_counts(self) -> Mapping[LogLevel, int]:
        return Counter(log.level for log in self.logs)

    @property
    def num_errors(self) -> int:
        level_counts = self._log_level_counts()
        return sum(count for level, count in level_counts.items() if level >= LogLevel.error)

    def exit_if_errors(self) -> None:
        level_counts = self._log_level_counts()
        if level_counts[LogLevel.critical]:
            sys.exit(2)
        elif level_counts[LogLevel.error]:
            sys.exit(1)

_current_ctx: ContextVar[LoggingContext] = ContextVar("current_ctx")

class PuyaJsonRender(structlog.processors.JSONRenderer):
    def __init__(self, *, base_path: str) -> None:
        super().__init__()
        self.base_path = base_path + os.path.sep

    def _location_json(
        self, location: SourceLocation | None
    ) -> Mapping[str, str | int | None] | None:
        if not location or not location.file:
            return None

        file = str(location.file)
        if file.startswith(self.base_path):
            file = file[len(self.base_path) :]

        return {
            "file": file,
            "line": location.line,
            "end_line": location.end_line,
            "column": location.column,
            "end_column": location.end_column,
        }

    def __call__(
        self,
        _logger: structlog.typing.WrappedLogger,
        _name: str,
        event_dict: structlog.typing.EventDict,
    ) -> str:
        # force event to str for compatibility with standard library
        event = event_dict.pop("event", None)
        if not isinstance(event, str):
            event = str(event)

        important: bool = event_dict.pop("important", False)
        location: SourceLocation | None = event_dict.pop("location", None)
        level = event_dict.pop("level", "info")

        return json.dumps(
            {
                "level": level,
                "location": self._location_json(location),
                "message": event,
                "important": important,
            }
        )

class PuyaConsoleRender(structlog.dev.ConsoleRenderer):
    def __init__(self, *, colors: bool, base_path: str):
        super().__init__(colors=colors)
        self.level_to_color = self.get_default_level_styles(colors)
        self.base_path = base_path
        if not self.base_path.endswith(
            os.path.sep
        ):  # TODO: can we always append the path seperator?
            self.base_path += os.path.sep

    def _location_as_link(self, location: SourceLocation | None) -> str:
        if not location or not location.file:
            return ""

        file = str(location.file)
        if file.startswith(self.base_path):
            file = file[len(self.base_path) :]

        line = str(location.line) if location.line else "1"
        col = f":{location.column + 1}" if location.column else ""
        return f"{file}:{line}{col}"

    def __call__(
        self,
        _logger: structlog.typing.WrappedLogger,
        _name: str,
        event_dict: structlog.typing.EventDict,
    ) -> str:
        # force event to str for compatibility with standard library
        event = event_dict.pop("event", None)
        if not isinstance(event, str):
            event = str(event)
        lines = [event]
        related_errors = event_dict.pop("related_lines", None)
        if related_errors:
            assert isinstance(related_errors, list)
            lines.extend(related_errors)
        important: bool = event_dict.pop("important", False)
        location: SourceLocation | None = event_dict.pop("location", None)
        if location and location.file is None:
            location = None
        location_as_link = self._location_as_link(location) if location else ""
        level = event_dict.pop("level", "info")

        align_related_lines = " " * (len(location_as_link) + 1 + len(level) + 1)
        sio = StringIO()
        reset_colour = self._styles.reset
        if important:
            sio.write(self._styles.bright)
            reset_colour += self._styles.bright
        for idx, line in enumerate(lines):
            if idx:
                sio.write("\n")
                sio.write(align_related_lines)
            else:
                if location:
                    sio.write(self._styles.logger_name)
                    location_link = self._location_as_link(location)
                    sio.write(location_link)
                    sio.write(" ")
                    sio.write(reset_colour)

                sio.write(self.level_to_color.get(level, ""))
                sio.write(level)
                sio.write(": ")
                sio.write(reset_colour)
            sio.write(line)
        sio.write(self._styles.reset)

        stack = event_dict.pop("stack", None)
        exc = event_dict.pop("exception", None)
        exc_info = event_dict.pop("exc_info", None)

        event_dict_keys: typing.Iterable[str] = event_dict.keys()
        if self._sort_keys:
            event_dict_keys = sorted(event_dict_keys)

        sio.write(
            " ".join(
                self._styles.kv_key
                + key
                + self._styles.reset
                + "="
                + self._styles.kv_value
                + self._repr(event_dict[key])
                + self._styles.reset
                for key in event_dict_keys
            )
        )

        if stack is not None:
            sio.write("\n" + stack)
            if exc_info or exc is not None:
                sio.write("\n\n" + "=" * 79 + "\n")

        if exc_info:
            exc_info = figure_out_exc_info(exc_info)

            self._exception_formatter(sio, exc_info)
        elif exc is not None:
            sio.write("\n" + exc)

        return sio.getvalue()

# copied from structlog.dev._figure_out_exc_info
def figure_out_exc_info(
    v: BaseException | structlog.typing.ExcInfo | bool,
) -> structlog.typing.ExcInfo | bool:
    """
    Depending on the Python version will try to do the smartest thing possible
    to transform *v* into an ``exc_info`` tuple.
    """
    if isinstance(v, BaseException):
        return type(v), v, v.__traceback__

    if isinstance(v, tuple):
        return v

    if v:
        return sys.exc_info()  # type: ignore[return-value]

    return v

@attrs.define
class FilterByLogLevel:
    min_log_level: LogLevel

    def __call__(
        self,
        _logger: structlog.typing.WrappedLogger,
        method: str,
        event_dict: structlog.typing.EventDict,
    ) -> structlog.typing.EventDict:
        if LogLevel[method] < self.min_log_level:
            raise structlog.DropEvent
        return event_dict

def configure_logging(
    *,
    min_log_level: LogLevel = LogLevel.notset,
    cache_logger: bool = True,
    log_format: LogFormat = LogFormat.default,
) -> None:
    if cache_logger and structlog.is_configured():
        raise ValueError(
            "Logging can not be configured more than once if using cache_logger = True"
        )
    base_path = str(Path.cwd())  # TODO: don't assume this?
    match log_format:
        case LogFormat.json:
            log_renderer: structlog.typing.Processor = PuyaJsonRender(base_path=base_path)
        case LogFormat.default:
            log_renderer = PuyaConsoleRender(
                colors="NO_COLOR" not in os.environ, base_path=base_path
            )
        case never:
            typing.assert_never(never)

    processors: list[structlog.typing.Processor] = [
        structlog.contextvars.merge_contextvars,
        structlog.processors.add_log_level,
        structlog.processors.StackInfoRenderer(),
        log_renderer,
    ]
    if min_log_level != LogLevel.notset:
        # filtering via a processor instead of via the logger like
        # structlog.make_filtering_bound_logger(min_log_level.value)
        # so that structlog.testing.capture_logs() still works in test cases
        processors.insert(0, FilterByLogLevel(min_log_level))
    structlog.configure(
        processors=processors,
        context_class=dict,
        logger_factory=structlog.PrintLoggerFactory(),
        cache_logger_on_first_use=cache_logger,
    )

class _Logger:
    def __init__(self, name: str, initial_values: dict[str, typing.Any]):
        self._logger = structlog.get_logger(name, **initial_values)

    def debug(
        self,
        event: object,
        *args: typing.Any,
        location: SourceLocation | None = None,
        **kwargs: typing.Any,
    ) -> None:
        self._report(LogLevel.debug, event, *args, location=location, **kwargs)

    def info(
        self,
        event: object,
        *args: typing.Any,
        location: SourceLocation | None = None,
        **kwargs: typing.Any,
    ) -> None:
        self._report(LogLevel.info, event, *args, location=location, **kwargs)

    def warning(
        self,
        event: object,
        *args: typing.Any,
        location: SourceLocation | None = None,
        **kwargs: typing.Any,
    ) -> None:
        self._report(LogLevel.warning, event, *args, location=location, **kwargs)

    def error(
        self,
        event: object,
        *args: typing.Any,
        location: SourceLocation | None = None,
        **kwargs: typing.Any,
    ) -> None:
        self._report(LogLevel.error, event, *args, location=location, **kwargs)

    def exception(
        self,
        event: object,
        *args: typing.Any,
        location: SourceLocation | None = None,
        **kwargs: typing.Any,
    ) -> None:
        kwargs.setdefault("exc_info", True)
        self._report(LogLevel.critical, event, *args, location=location, **kwargs)

    def critical(
        self,
        event: object,
        *args: typing.Any,
        location: SourceLocation | None = None,
        **kwargs: typing.Any,
    ) -> None:
        self._report(LogLevel.critical, event, *args, location=location, **kwargs)

    def log(
        self,
        level: LogLevel,
        event: object,
        *args: typing.Any,
        location: SourceLocation | None = None,
        **kwargs: typing.Any,
    ) -> None:
        self._report(level, event, *args, location=location, **kwargs)

    def _report(
        self,
        level: LogLevel,
        event: object,
        *args: typing.Any,
        location: SourceLocation | None = None,
        **kwargs: typing.Any,
    ) -> None:
        log_ctx = _current_ctx.get(None)
        if (
            level >= LogLevel.error
            and location
            and log_ctx
            and log_ctx.sources_by_path
            and location.file
        ):
            file_source = log_ctx.sources_by_path.get(location.file)
            if file_source is not None:
                kwargs["related_lines"] = _get_pretty_source(file_source, location)
        self._logger.log(level, event, *args, location=location, **kwargs)
        if log_ctx:
            if isinstance(event, str) and args:
                message = event % args
            else:
                message = str(event)
            log_ctx.logs.append(Log(level, message, location))

def _get_pretty_source(
    file_source: Sequence[str], location: SourceLocation
) -> Sequence[str] | None:
    lines = file_source[location.line - 1 : location.end_line]
    if len(lines) != location.line_count:
        logger = get_logger(__name__)
        logger.warning(f"source length mismatch for {location}")
        return None
    try:
        (source_line,) = lines
    except ValueError:
        # source line is followed by additional lines, don't bother annotating columns
        return lines
    # Shifts column after tab expansion
    column = len(source_line[: location.column].expandtabs())
    end_column = len(source_line[: location.end_column].expandtabs())

    return [
        source_line.expandtabs(),
        " " * column + f"^{'~' * max(end_column - column - 1, 0)}",
    ]

def get_logger(name: str, **initial_values: typing.Any) -> _Logger:
    return _Logger(name, initial_values)

@contextlib.contextmanager
def logging_context() -> Iterator[LoggingContext]:
    ctx = LoggingContext()
    restore = _current_ctx.set(ctx)
    try:
        yield ctx
    finally:
        _current_ctx.reset(restore)


from pathlib import Path

import attrs
import cattrs.preconf.json

from puya import log
from puya.awst import (
    nodes as awst_nodes,
    serialize,
)
from puya.compile import awst_to_teal
from puya.errors import log_exceptions
from puya.options import PuyaOptions
from puya.program_refs import ContractReference, LogicSigReference

logger = log.get_logger(__name__)

@attrs.frozen(kw_only=True)
class _PuyaOptionsWithCompilationSet(PuyaOptions):
    compilation_set: dict[str, Path]

def main(*, options_json: str, awst_json: str, source_annotations_json: str | None) -> None:
    with log.logging_context() as log_ctx, log_exceptions():
        json_converter = cattrs.preconf.json.make_converter()
        sources_by_path = {}
        if source_annotations_json:
            sources_by_path = json_converter.loads(
                source_annotations_json, dict[Path, list[str] | None]
            )
        log_ctx.sources_by_path = sources_by_path
        awst = serialize.awst_from_json(awst_json)
        options = json_converter.loads(options_json, _PuyaOptionsWithCompilationSet)
        compilation_set = dict[ContractReference | LogicSigReference, Path]()
        awst_lookup = {n.id: n for n in awst}
        for target_id, path in options.compilation_set.items():
            match awst_lookup.get(target_id):
                case awst_nodes.Contract(id=contract_id):
                    compilation_set[contract_id] = path
                case awst_nodes.LogicSignature(id=lsig_id):
                    compilation_set[lsig_id] = path
                case None:
                    logger.error(f"compilation target {target_id!r} not found in AWST")
                case other:
                    logger.error(f"unexpected compilation target type: {type(other).__name__}")
        awst_to_teal(log_ctx, options, compilation_set, sources_by_path, awst)
    # note: needs to be outside the with block
    log_ctx.exit_if_errors()


import contextlib
from collections.abc import Iterator

import attrs

from puya.errors import InternalError

@attrs.define
class AlignedWriter:
    _headers: list[str] = attrs.field(factory=list)
    _lines: list[list[str] | str] = attrs.field(factory=list)
    _current_line: list[str] = attrs.field(factory=list)
    _ignore_current_line: bool = False
    _padding: dict[int, int] = attrs.field(factory=dict)
    _indent: str = ""
    omit_empty_columns: bool = True

    @property
    def ignore_current_line(self) -> bool:
        return self._ignore_current_line

    @property
    def current_indent(self) -> str:
        return self._indent

    @contextlib.contextmanager
    def indent(self) -> Iterator[None]:
        indent_width = 4
        original_indent = self._indent
        self._indent += " " * indent_width
        try:
            yield
        finally:
            self._indent = original_indent

    def add_header(self, header: str, padding: int = 1) -> None:
        self._padding[len(self._headers)] = padding
        self._headers.append(header)

    def append(self, part: str) -> None:
        self._current_line.append(part)

    def new_line(self) -> None:
        parts = self._current_line
        if not self._ignore_current_line:
            if parts and self._indent:
                parts[0] = f"{self._indent}{parts[0]}"
            self._lines.append(parts)
        self._ignore_current_line = False
        self._current_line = []

    def append_line(self, line: str) -> None:
        if self._current_line or self._ignore_current_line:
            raise InternalError(
                "Cannot append a new line while a current line is in progress, missing new_line()?"
            )
        self._lines.append(line)

    def ignore_line(self) -> None:
        self._ignore_current_line = True

    def _join_columns(self, line: list[str], widths: dict[int, int]) -> str:
        return "".join(
            part.ljust(widths.get(column, 0))
            for column, part in enumerate(line)
            if widths.get(column, 0) or not self.omit_empty_columns
        ).rstrip()

    def write(self) -> list[str]:
        widths = dict[int, int]()
        all_lines = self._lines.copy()
        for parts in all_lines:
            if isinstance(parts, list):
                for column, part in enumerate(parts):
                    widths[column] = max(widths.get(column, 0), len(part))

        for column, width in widths.items():
            if width == 0 and self.omit_empty_columns:
                continue
            if column < len(self._headers):
                width = max(width, len(self._headers[column]))
            widths[column] = width + self._padding.get(column, 1)

        if self._headers:
            all_lines.insert(0, self._headers)

        return [
            line if isinstance(line, str) else self._join_columns(line, widths)
            for line in all_lines
        ]


import typing
from collections.abc import Sequence

import attrs

from puya import log
from puya.errors import CodeError, InternalError
from puya.ir import models as ir
from puya.ir.types_ import AVMBytesEncoding
from puya.ir.visitor import IRVisitor
from puya.mir import models
from puya.mir.context import ProgramMIRContext
from puya.utils import biguint_bytes_eval

logger = log.get_logger(__name__)

@attrs.define
class MemoryIRBuilder(IRVisitor[None]):
    context: ProgramMIRContext = attrs.field(on_setattr=attrs.setters.frozen)
    current_subroutine: ir.Subroutine
    is_main: bool
    current_ops: list[models.Op] = attrs.field(factory=list)
    terminator: models.ControlOp | None = None
    active_op: ir.Op | ir.ControlOp | None = None

    def _add_op(self, op: models.Op) -> None:
        self.current_ops.append(op)

    def _terminate(self, op: models.ControlOp) -> None:
        assert self.terminator is None
        self.terminator = op

    def _get_block_name(self, block: ir.BasicBlock) -> str:
        if block is self.current_subroutine.entry:
            return self.context.subroutine_names[self.current_subroutine]
        assert block in self.current_subroutine.body
        comment = (block.comment or "block").replace(" ", "_")
        subroutine_name = self.context.subroutine_names[self.current_subroutine]
        return f"{subroutine_name}_{comment}@{block.id}"

    def visit_assignment(self, ass: ir.Assignment) -> None:
        ass.source.accept(self)
        # right most target is top of stack
        for target in reversed(ass.targets):
            try:
                param_idx = self.current_subroutine.parameters.index(target)
            except ValueError:
                self._add_op(
                    models.AbstractStore(
                        local_id=target.local_id,
                        source_location=ass.source_location,
                        atype=target.atype,
                    )
                )
            else:
                index = param_idx - len(self.current_subroutine.parameters)
                self._add_op(
                    models.StoreParam(
                        local_id=target.local_id,
                        index=index,
                        source_location=ass.source_location,
                        atype=target.atype,
                    )
                )

    def visit_register(self, reg: ir.Register) -> None:
        produces = (reg.local_id,)
        if isinstance(self.active_op, ir.Assignment):
            if reg is self.active_op.source:
                (target,) = self.active_op.targets
                produces = (target.local_id,)
            elif (
                isinstance(self.active_op.source, ir.ValueTuple)
                and reg in self.active_op.source.values
            ):
                index = self.active_op.source.values.index(reg)
                target = self.active_op.targets[index]
                produces = (target.local_id,)
        try:
            param_idx = self.current_subroutine.parameters.index(reg)
        except ValueError:
            self._add_op(
                models.AbstractLoad(
                    local_id=reg.local_id,
                    produces=produces,
                    source_location=(self.active_op or reg).source_location,
                    atype=reg.atype,
                )
            )
        else:
            index = param_idx - len(self.current_subroutine.parameters)
            if produces[0] == reg.local_id:
                produces = (f"{produces[0]} (copy)",)
            self._add_op(
                models.LoadParam(
                    local_id=reg.local_id,
                    produces=produces,
                    index=index,
                    source_location=(self.active_op or reg).source_location,
                    atype=reg.atype,
                )
            )

    def visit_undefined(self, val: ir.Undefined) -> None:
        self._add_op(
            models.Undefined(
                atype=val.atype,
                source_location=val.source_location,
            )
        )

    def visit_template_var(self, deploy_var: ir.TemplateVar) -> None:
        self._add_op(
            models.TemplateVar(
                name=deploy_var.name,
                atype=deploy_var.atype,
                source_location=deploy_var.source_location,
            )
        )

    def visit_uint64_constant(self, const: ir.UInt64Constant) -> None:
        self._add_op(
            models.Int(
                const.value if not const.teal_alias else const.teal_alias,
                source_location=const.source_location,
            )
        )

    def visit_biguint_constant(self, const: ir.BigUIntConstant) -> None:
        big_uint_bytes = biguint_bytes_eval(const.value)
        self._add_op(
            models.Byte(
                big_uint_bytes,
                source_location=const.source_location,
                encoding=AVMBytesEncoding.base16,
            )
        )

    def visit_bytes_constant(self, const: ir.BytesConstant) -> None:
        self._add_op(
            models.Byte(
                const.value, encoding=const.encoding, source_location=const.source_location
            )
        )

    def visit_address_constant(self, const: ir.AddressConstant) -> None:
        self._add_op(
            models.Address(
                const.value,
                source_location=const.source_location,
            )
        )

    def visit_method_constant(self, const: ir.MethodConstant) -> None:
        self._add_op(
            models.Method(
                const.value,
                source_location=const.source_location,
            )
        )

    def visit_intrinsic_op(self, intrinsic: ir.Intrinsic) -> None:
        discard_results = intrinsic is self.active_op

        if intrinsic.op.code.startswith("box_"):
            try:
                box_key = intrinsic.args[0]
            except ValueError:
                raise InternalError("box key arg not found", intrinsic.source_location) from None
            if isinstance(box_key, ir.BytesConstant) and not box_key.value:
                raise CodeError("AVM does not support empty box keys", intrinsic.source_location)

        for arg in intrinsic.args:
            arg.accept(self)
        produces = len(intrinsic.op_signature.returns)
        self._add_op(
            models.IntrinsicOp(
                op_code=intrinsic.op.code,
                immediates=intrinsic.immediates,
                source_location=intrinsic.source_location,
                consumes=len(intrinsic.op_signature.args),
                produces=_produces_from_op(intrinsic.op.code, produces, self.active_op),
                error_message=intrinsic.error_message,
            )
        )
        if discard_results and produces:
            self._add_op(models.Pop(produces))

    def visit_invoke_subroutine(self, callsub: ir.InvokeSubroutine) -> None:
        discard_results = callsub is self.active_op
        target = callsub.target

        callsub_op = models.CallSub(
            target=self.context.subroutine_names[target],
            parameters=len(target.parameters),
            returns=len(target.returns),
            produces=_produces_from_op(
                self.context.subroutine_names[target], len(target.returns), self.active_op
            ),
            source_location=callsub.source_location,
        )

        # prepare args
        for arg in callsub.args:
            arg.accept(self)

        # call sub
        self._add_op(callsub_op)

        if discard_results and target.returns:
            num_returns = len(target.returns)
            self._add_op(models.Pop(num_returns))

    def visit_conditional_branch(self, branch: ir.ConditionalBranch) -> None:
        branch.condition.accept(self)
        self._terminate(
            models.ConditionalBranch(
                nonzero_target=self._get_block_name(branch.non_zero),
                zero_target=self._get_block_name(branch.zero),
                source_location=branch.source_location,
            )
        )

    def visit_goto(self, goto: ir.Goto) -> None:
        self._terminate(
            models.Goto(
                target=self._get_block_name(goto.target),
                source_location=goto.source_location,
            )
        )

    def visit_goto_nth(self, goto_nth: ir.GotoNth) -> None:
        block_labels = [self._get_block_name(block) for block in goto_nth.blocks]
        goto_nth.value.accept(self)
        self._terminate(
            models.Switch(
                switch_targets=block_labels,
                default_target=self._get_block_name(goto_nth.default),
                source_location=goto_nth.source_location,
            )
        )

    def visit_switch(self, switch: ir.Switch) -> None:
        blocks = list[str]()
        for case, block in switch.cases.items():
            case.accept(self)
            block_name = self._get_block_name(block)
            blocks.append(block_name)
        switch.value.accept(self)
        self._terminate(
            models.Match(
                match_targets=blocks,
                default_target=self._get_block_name(switch.default),
                source_location=switch.source_location,
            )
        )

    def visit_subroutine_return(self, retsub: ir.SubroutineReturn) -> None:
        for r in retsub.result:
            r.accept(self)
        if self.is_main:
            assert len(retsub.result) == 1
            self._terminate(models.ProgramExit(source_location=retsub.source_location))
        else:
            self._terminate(
                models.RetSub(returns=len(retsub.result), source_location=retsub.source_location)
            )

    def visit_program_exit(self, exit_: ir.ProgramExit) -> None:
        exit_.result.accept(self)
        self._terminate(models.ProgramExit(source_location=exit_.source_location))

    def visit_fail(self, fail: ir.Fail) -> None:
        self._terminate(
            models.Err(error_message=fail.error_message, source_location=fail.source_location)
        )

    def lower_block_to_mir(self, block: ir.BasicBlock) -> models.MemoryBasicBlock:
        self.current_ops = list[models.Op]()
        self.terminator = None
        for op in block.all_ops:
            assert not isinstance(op, ir.Phi)
            self.active_op = op
            op.accept(self)

        assert self.terminator is not None
        block_name = self._get_block_name(block)
        predecessors = [self._get_block_name(b) for b in block.predecessors]
        assert block.id is not None
        return models.MemoryBasicBlock(
            id=block.id,
            block_name=block_name,
            mem_ops=self.current_ops,
            terminator=self.terminator,
            predecessors=predecessors,
            source_location=block.source_location,
        )

    def visit_compiled_contract_reference(self, const: ir.CompiledContractReference) -> None:
        _unexpected_node(const)

    def visit_compiled_logicsig_reference(self, const: ir.CompiledLogicSigReference) -> None:
        _unexpected_node(const)

    def visit_value_tuple(self, tup: ir.ValueTuple) -> None:
        _unexpected_node(tup)

    def visit_itxn_constant(self, const: ir.ITxnConstant) -> None:
        _unexpected_node(const)

    def visit_inner_transaction_field(self, field: ir.InnerTransactionField) -> None:
        _unexpected_node(field)

    def visit_phi(self, phi: ir.Phi) -> None:
        _unexpected_node(phi)

    def visit_phi_argument(self, arg: ir.PhiArgument) -> None:
        _unexpected_node(arg)

def _unexpected_node(node: ir.IRVisitable) -> typing.Never:
    raise InternalError(
        f"Encountered node of type {type(node).__name__!r} during codegen"
        f" - should have been eliminated in prior stages",
        node.source_location,
    )

def _produces_from_op(
    prefix: str, size: int, maybe_assignment: ir.IRVisitable | None
) -> Sequence[str]:
    produces = models.produces_from_desc(prefix, size)
    if isinstance(maybe_assignment, ir.Assignment):
        produces = [r.local_id for r in maybe_assignment.targets]
    return produces


from collections.abc import Mapping

import attrs

from puya.context import ArtifactCompileContext
from puya.ir import models as ir
from puya.mir import models
from puya.mir.vla import VariableLifetimeAnalysis
from puya.utils import attrs_extend

@attrs.define(kw_only=True)
class ProgramMIRContext(ArtifactCompileContext):
    program: ir.Program
    subroutine_names: Mapping[ir.Subroutine, str] = attrs.field(init=False)

    @subroutine_names.default
    def _get_short_subroutine_names(self) -> dict[ir.Subroutine, str]:
        """Return a mapping of unique TEAL names for all subroutines in program, while attempting
        to use the shortest name possible"""
        names = {"main": self.program.main}
        for subroutine in self.program.subroutines:
            if subroutine.short_name and subroutine.short_name not in names:
                name = subroutine.short_name
            else:
                assert subroutine.id not in names
                name = subroutine.id
            names[name] = subroutine

        return {v: k for k, v in names.items()}

    def for_subroutine(self, subroutine: models.MemorySubroutine) -> "SubroutineCodeGenContext":
        return attrs_extend(SubroutineCodeGenContext, self, subroutine=subroutine)

@attrs.define(frozen=False)
class SubroutineCodeGenContext(ProgramMIRContext):
    subroutine: models.MemorySubroutine
    _vla: VariableLifetimeAnalysis | None = None

    @property
    def vla(self) -> VariableLifetimeAnalysis:
        if self._vla is None:
            self._vla = VariableLifetimeAnalysis.analyze(self.subroutine)
        return self._vla

    def invalidate_vla(self) -> None:
        self._vla = None


from puya.context import ArtifactCompileContext
from puya.ir import models as ir
from puya.mir import models
from puya.mir.builder import MemoryIRBuilder
from puya.mir.context import ProgramMIRContext
from puya.mir.stack_allocation import global_stack_allocation
from puya.utils import attrs_extend

def program_ir_to_mir(context: ArtifactCompileContext, program_ir: ir.Program) -> models.Program:
    ctx = attrs_extend(ProgramMIRContext, context, program=program_ir)

    result = models.Program(
        kind=program_ir.kind,
        main=_lower_subroutine_to_mir(ctx, program_ir.main, is_main=True),
        subroutines=[
            _lower_subroutine_to_mir(ctx, ir_sub, is_main=False)
            for ir_sub in program_ir.subroutines
        ],
        avm_version=program_ir.avm_version,
    )
    global_stack_allocation(ctx, result)
    return result

def _lower_subroutine_to_mir(
    context: ProgramMIRContext, subroutine: ir.Subroutine, *, is_main: bool
) -> models.MemorySubroutine:
    builder = MemoryIRBuilder(context=context, current_subroutine=subroutine, is_main=is_main)
    body = [builder.lower_block_to_mir(block) for block in subroutine.body]
    signature = models.Signature(
        name=subroutine.id,
        parameters=[
            models.Parameter(name=p.name, local_id=p.local_id, atype=p.atype)
            for p in subroutine.parameters
        ],
        returns=[r.avm_type for r in subroutine.returns],
    )
    return models.MemorySubroutine(
        id=context.subroutine_names[subroutine],
        signature=signature,
        is_main=is_main,
        body=body,
        source_location=subroutine.source_location,
    )


from __future__ import annotations

import abc
import typing
import typing as t
from functools import cached_property

import attrs

from puya.avm import AVMType
from puya.errors import InternalError
from puya.ir.utils import format_bytes, format_error_comment
from puya.program_refs import ProgramKind

if t.TYPE_CHECKING:
    from collections.abc import Iterator, Mapping, Sequence

    from puya.ir.types_ import AVMBytesEncoding
    from puya.mir.visitor import MIRVisitor
    from puya.parse import SourceLocation

_T = t.TypeVar("_T")

@attrs.frozen(kw_only=True, eq=False, str=False)
class BaseOp(abc.ABC):
    error_message: str | None = None
    source_location: SourceLocation | None = None
    consumes: int
    """How many values are removed from the top of l-stack
    Does not take into account any manipulations lower in the stack e.g. from Load*Stack"""
    produces: Sequence[str]
    """The local ids that are appended to the l-stack
    Does not take into account any manipulations lower in the stack e.g. from Store*Stack"""

    @abc.abstractmethod
    def accept(self, visitor: MIRVisitor[_T]) -> _T: ...

def _is_single_item(_: object, __: object, value: Sequence[str]) -> None:
    assert len(value) == 1, "expected single item"

@attrs.frozen(eq=False)
class Op(BaseOp, abc.ABC):
    @abc.abstractmethod
    def __str__(self) -> str: ...

@attrs.frozen(eq=False)
class Int(Op):
    value: int | str
    consumes: int = attrs.field(default=0, init=False)
    produces: Sequence[str] = attrs.field(validator=_is_single_item)

    @produces.default
    def _produces(self) -> Sequence[str]:
        return (str(self.value),)

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_int(self)

    def __str__(self) -> str:
        return f"int {self.value}"

@attrs.frozen(eq=False)
class Byte(Op):
    value: bytes
    encoding: AVMBytesEncoding
    consumes: int = attrs.field(default=0, init=False)
    produces: Sequence[str] = attrs.field(validator=_is_single_item)

    @produces.default
    def _produces(self) -> Sequence[str]:
        return (format_bytes(self.value, self.encoding),)

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_byte(self)

    def __str__(self) -> str:
        return f"byte {format_bytes(self.value, self.encoding)}"

@attrs.frozen(eq=False)
class Undefined(Op):
    atype: typing.Literal[AVMType.bytes, AVMType.uint64]
    consumes: int = attrs.field(default=0, init=False)
    produces: Sequence[str] = attrs.field(validator=_is_single_item)

    @produces.default
    def _produces(self) -> Sequence[str]:
        return ("undefined",)

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_undefined(self)

    def __str__(self) -> str:
        return "undefined"

@attrs.frozen(eq=False)
class TemplateVar(Op):
    name: str
    atype: AVMType
    op_code: typing.Literal["int", "byte"] = attrs.field(init=False)
    consumes: int = attrs.field(default=0, init=False)
    produces: Sequence[str] = attrs.field(validator=_is_single_item)

    @produces.default
    def _produces(self) -> Sequence[str]:
        return (self.name,)

    @op_code.default
    def _default_opcode(self) -> typing.Literal["int", "byte"]:
        match self.atype:
            case AVMType.bytes:
                return "byte"
            case AVMType.uint64:
                return "int"
            case _:
                raise InternalError(
                    f"Unsupported atype for TemplateVar: {self.atype}", self.source_location
                )

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_template_var(self)

    def __str__(self) -> str:
        return f"{self.op_code} {self.name}"

@attrs.frozen(eq=False)
class Address(Op):
    value: str
    consumes: int = attrs.field(default=0, init=False)
    produces: Sequence[str] = attrs.field(validator=_is_single_item)

    @produces.default
    def _produces(self) -> Sequence[str]:
        return (f"Address({self.value})",)

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_address(self)

    def __str__(self) -> str:
        return f'addr "{self.value}"'

@attrs.frozen(eq=False)
class Method(Op):
    value: str
    consumes: int = attrs.field(default=0, init=False)
    produces: Sequence[str] = attrs.field(validator=_is_single_item)

    @produces.default
    def _produces(self) -> Sequence[str]:
        return (f"Method({self.value})",)

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_method(self)

    def __str__(self) -> str:
        return f"method {self.value}"

@attrs.frozen(eq=False)
class Comment(Op):
    comment: str
    consumes: int = attrs.field(default=0, init=False)
    produces: Sequence[str] = attrs.field(default=(), init=False)

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_comment(self)

    def __str__(self) -> str:
        return f"// {self.comment}"

@attrs.frozen(kw_only=True, eq=False)
class StoreOp(Op, abc.ABC):
    """An op for storing values"""

    local_id: str
    atype: AVMType = attrs.field()

    @atype.validator
    def _validate_not_any(self, _attribute: object, atype: AVMType) -> None:
        if atype is AVMType.any:
            raise InternalError(f"Register has type any: {self}", self.source_location)

@attrs.frozen(kw_only=True, eq=False)
class LoadOp(Op, abc.ABC):
    """An op for loading values"""

    local_id: str
    atype: AVMType = attrs.field()

    @atype.validator
    def _validate_not_any(self, _attribute: object, atype: AVMType) -> None:
        if atype is AVMType.any:
            raise InternalError(f"Register has type any: {self}", self.source_location)

@attrs.frozen(eq=False)
class AbstractStore(StoreOp):
    consumes: int = attrs.field(default=1, init=False)
    produces: Sequence[str] = attrs.field(default=(), init=False)

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_abstract_store(self)

    def __str__(self) -> str:
        return f"v-store {self.local_id}"

@attrs.frozen(eq=False)
class AbstractLoad(LoadOp):
    consumes: int = attrs.field(default=0, init=False)
    produces: Sequence[str] = attrs.field(validator=_is_single_item)

    @produces.default
    def _produces(self) -> Sequence[str]:
        return (self.local_id,)

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_abstract_load(self)

    def __str__(self) -> str:
        return f"v-load {self.local_id}"

@attrs.frozen(eq=False, kw_only=True)
class StoreLStack(StoreOp):
    depth: int = attrs.field(validator=attrs.validators.ge(0))
    copy: bool
    consumes: int = attrs.field(default=0, init=False)
    produces: Sequence[str] = attrs.field()

    @produces.default
    def _produces(self) -> Sequence[str]:
        return (f"{self.local_id} (copy)",) if self.copy else ()

    @produces.validator
    def _validate_produces(self, _: object, value: Sequence[str]) -> None:
        assert len(value) == (1 if self.copy else 0), "invalid produces size"

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_store_l_stack(self)

    def __str__(self) -> str:
        op = "l-store-copy" if self.copy else "l-store"
        return f"{op} {self.local_id} {self.depth}"

@attrs.frozen(eq=False)
class LoadLStack(LoadOp):
    copy: bool
    consumes: int = attrs.field(init=False, default=0)
    produces: Sequence[str] = attrs.field(validator=_is_single_item)
    # depth can only be defined after koopmans pass and dead store removal
    depth: int | None = None

    @produces.default
    def _produces(self) -> Sequence[str]:
        produces = self.local_id
        if self.copy:
            produces += " (copy)"
        return (produces,)

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_load_l_stack(self)

    def __str__(self) -> str:
        depth = "" if self.depth is None else f" {self.depth}"
        op = "l-load-copy" if self.copy else "l-load"
        return f"{op} {self.local_id}{depth}"

@attrs.frozen(eq=False)
class StoreXStack(StoreOp):
    depth: int = attrs.field(validator=attrs.validators.ge(0))
    consumes: int = attrs.field(default=1, init=False)
    produces: Sequence[str] = attrs.field(default=(), init=False)

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_store_x_stack(self)

    def __str__(self) -> str:
        return f"x-store {self.local_id}"

@attrs.frozen(eq=False)
class LoadXStack(LoadOp):
    depth: int = attrs.field(validator=attrs.validators.ge(0))
    consumes: int = attrs.field(default=0, init=False)
    produces: Sequence[str] = attrs.field(validator=_is_single_item)

    @produces.default
    def _produces(self) -> Sequence[str]:
        return (self.local_id,)

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_load_x_stack(self)

    def __str__(self) -> str:
        return f"x-load {self.local_id}"

@attrs.frozen(eq=False)
class StoreFStack(StoreOp):
    depth: int = attrs.field(validator=attrs.validators.ge(0))
    frame_index: int = attrs.field(validator=attrs.validators.ge(0))
    insert: bool = False
    consumes: int = attrs.field(default=1, init=False)
    produces: Sequence[str] = attrs.field(default=(), init=False)

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_store_f_stack(self)

    def __str__(self) -> str:
        return f"f-store {self.local_id}"

@attrs.frozen(eq=False)
class LoadFStack(LoadOp):
    depth: int = attrs.field(validator=attrs.validators.ge(0))
    frame_index: int = attrs.field(validator=attrs.validators.ge(0))
    consumes: int = attrs.field(default=0, init=False)
    produces: Sequence[str] = attrs.field(validator=_is_single_item)

    @produces.default
    def _produces(self) -> Sequence[str]:
        return (f"{self.local_id} (copy)",)

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_load_f_stack(self)

    def __str__(self) -> str:
        return f"f-load {self.local_id}"

@attrs.frozen(eq=False)
class LoadParam(LoadOp):
    index: int
    consumes: int = attrs.field(default=0, init=False)
    produces: Sequence[str] = attrs.field(validator=_is_single_item)

    @produces.default
    def _produces(self) -> Sequence[str]:
        return (f"{self.local_id} (copy)",)

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_load_param(self)

    def __str__(self) -> str:
        return f"p-load {self.local_id}"

@attrs.frozen(eq=False)
class StoreParam(StoreOp):
    index: int
    consumes: int = attrs.field(default=1, init=False)
    produces: Sequence[str] = attrs.field(default=(), init=False)

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_store_param(self)

    def __str__(self) -> str:
        return f"p-store {self.local_id}"

@attrs.frozen(eq=False)
class Pop(Op):
    n: int
    consumes: int = attrs.field(init=False)
    produces: Sequence[str] = attrs.field(default=(), init=False)

    @consumes.default
    def _consumes(self) -> int:
        return self.n

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_pop(self)

    def __str__(self) -> str:
        return f"pop {self.n}"

@attrs.frozen(eq=False)
class Allocate(Op):
    bytes_vars: Sequence[str]
    uint64_vars: Sequence[str]
    consumes: int = attrs.field(default=0, init=False)
    produces: Sequence[str] = attrs.field(default=(), init=False)

    @property
    def allocate_on_entry(self) -> Sequence[str]:
        return [*self.bytes_vars, *self.uint64_vars]

    @property
    def num_bytes(self) -> int:
        return len(self.bytes_vars)

    @property
    def num_uints(self) -> int:
        return len(self.uint64_vars)

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_allocate(self)

    def __str__(self) -> str:
        return f"allocate {len(self.allocate_on_entry)} to stack"

@attrs.frozen(eq=False)
class CallSub(Op):
    target: str
    parameters: int
    returns: int
    consumes: int = attrs.field(init=False)
    produces: Sequence[str] = attrs.field()

    @consumes.default
    def _consumes(self) -> int:
        return self.parameters

    @produces.validator
    def _validate_produces(self, _: object, value: Sequence[str]) -> None:
        assert len(value) == self.returns, "invalid produces size"

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_callsub(self)

    def __str__(self) -> str:
        return f"callsub {self.target}"

@attrs.frozen(eq=False)
class IntrinsicOp(Op):
    """An Op that does something other than just manipulating memory"""

    # TODO: use enum values for these ops
    op_code: str
    immediates: Sequence[str | int] = attrs.field(default=(), converter=tuple[str | int, ...])

    def __attrs_post_init__(self) -> None:
        if self.op_code in ("b", "bz", "bnz", "switch", "match", "retsub", "err", "return"):
            raise InternalError(
                f"Branching op {self.op_code} should map to explicit MIR ControlOp",
                self.source_location,
            )

    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_intrinsic(self)

    def __str__(self) -> str:
        result = [self.op_code, *map(str, self.immediates)]
        if self.error_message:
            result.append("//")
            result.append(format_error_comment(self.op_code, self.error_message))
        return " ".join(result)

@attrs.frozen(eq=False)
class ControlOp(BaseOp, abc.ABC):
    @abc.abstractmethod
    def targets(self) -> Sequence[str]: ...

    @abc.abstractmethod
    def _str_components(self) -> tuple[str, ...]: ...

    @typing.final
    def __str__(self) -> str:
        result = tuple(self._str_components())
        if self.error_message:
            result += ("//", self.error_message)
        return " ".join(result)

@attrs.frozen(eq=False)
class RetSub(ControlOp):
    returns: int
    fx_height: int = 0
    # l-stack is discarded after this op
    consumes: int = attrs.field(default=0, init=False)
    produces: Sequence[str] = attrs.field(default=(), init=False)

    @typing.override
    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_retsub(self)

    @typing.override
    def targets(self) -> Sequence[str]:
        return ()

    @typing.override
    def _str_components(self) -> tuple[str, ...]:
        return ("retsub",)

@attrs.frozen(eq=False)
class ProgramExit(ControlOp):
    consumes: int = attrs.field(default=1, init=False)
    produces: Sequence[str] = attrs.field(default=(), init=False)

    @typing.override
    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_program_exit(self)

    @typing.override
    def targets(self) -> Sequence[str]:
        return ()

    @typing.override
    def _str_components(self) -> tuple[str, ...]:
        return ("return",)

@attrs.frozen(eq=False)
class Err(ControlOp):
    consumes: int = attrs.field(default=0, init=False)
    produces: Sequence[str] = attrs.field(default=(), init=False)

    @typing.override
    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_err(self)

    @typing.override
    def targets(self) -> Sequence[str]:
        return ()

    @typing.override
    def _str_components(self) -> tuple[str, ...]:
        return ("err",)

@attrs.frozen(eq=False)
class Goto(ControlOp):
    consumes: int = attrs.field(default=0, init=False)
    produces: Sequence[str] = attrs.field(default=(), init=False)
    target: str

    @typing.override
    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_goto(self)

    @typing.override
    def targets(self) -> Sequence[str]:
        return (self.target,)

    @typing.override
    def _str_components(self) -> tuple[str, ...]:
        return "b", self.target

@attrs.frozen(eq=False)
class ConditionalBranch(ControlOp):
    consumes: int = attrs.field(default=1, init=False)
    produces: Sequence[str] = attrs.field(default=(), init=False)
    zero_target: str
    nonzero_target: str

    @typing.override
    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_conditional_branch(self)

    @typing.override
    def targets(self) -> Sequence[str]:
        return self.zero_target, self.nonzero_target

    @typing.override
    def _str_components(self) -> tuple[str, ...]:
        return "bz", self.zero_target, ";", "b", self.nonzero_target

@attrs.frozen(eq=False)
class Switch(ControlOp):
    consumes: int = attrs.field(default=1, init=False)
    produces: Sequence[str] = attrs.field(default=(), init=False)
    switch_targets: Sequence[str] = attrs.field(converter=tuple[str, ...])
    default_target: str

    @typing.override
    def targets(self) -> Sequence[str]:
        return *self.switch_targets, self.default_target

    @typing.override
    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_switch(self)

    @typing.override
    def _str_components(self) -> tuple[str, ...]:
        return "switch", *self.switch_targets, ";", "b", self.default_target

@attrs.frozen(eq=False)
class Match(ControlOp):
    produces: Sequence[str] = attrs.field(default=(), init=False)
    match_targets: Sequence[str] = attrs.field(converter=tuple[str, ...])
    consumes: int = attrs.field(init=False)
    default_target: str

    @consumes.default
    def _consumes(self) -> int:
        return len(self.match_targets) + 1

    @typing.override
    def targets(self) -> Sequence[str]:
        return *self.match_targets, self.default_target

    @typing.override
    def accept(self, visitor: MIRVisitor[_T]) -> _T:
        return visitor.visit_match(self)

    @typing.override
    def _str_components(self) -> tuple[str, ...]:
        return "match", *self.match_targets, ";", "b", self.default_target

@attrs.define(eq=False, repr=False, kw_only=True)
class MemoryBasicBlock:
    id: int
    block_name: str
    mem_ops: list[Op]
    terminator: ControlOp
    predecessors: list[str]
    source_location: SourceLocation
    # the ordering of values on the stack is used by debug maps
    # the assumption is lower levels won't change the order of variables in the stack
    # however they can introduce changes that do that ordering more efficiently
    x_stack_in: Sequence[str] | None = None
    """local_ids on x-stack on entry to a block"""
    x_stack_out: Sequence[str] | None = None
    """local_ids on x-stack on exit from a block"""
    f_stack_in: Sequence[str] = attrs.field(factory=list)
    """local_ids on f-stack on entry to a block"""
    f_stack_out: Sequence[str] = attrs.field(factory=list)
    """local_ids on f-stack on exit from a block"""

    @property
    def ops(self) -> Sequence[BaseOp]:
        return *self.mem_ops, self.terminator

    def __repr__(self) -> str:
        return self.block_name

    @property
    def entry_stack_height(self) -> int:
        return len(self.f_stack_in) + len(self.x_stack_in or ())

    @property
    def exit_stack_height(self) -> int:
        return len(self.f_stack_out) + len(self.x_stack_out or ())

    @property
    def successors(self) -> Sequence[str]:
        return self.terminator.targets()

@attrs.frozen(kw_only=True)
class Parameter:
    name: str
    local_id: str
    atype: AVMType

@attrs.frozen(str=False)
class Signature:
    name: str
    parameters: Sequence[Parameter]
    returns: Sequence[AVMType]

    def __str__(self) -> str:
        params = ", ".join(f"{p.name}: {p.atype.name}" for p in self.parameters)
        returns = ", ".join(str(r.name) for r in self.returns)
        return f"{self.name}({params}) -> {returns or 'void'}:"

@attrs.define
class MemorySubroutine:
    """A lower form of IR that is concerned with memory assignment (both stack and scratch)"""

    id: str
    is_main: bool
    signature: Signature
    body: Sequence[MemoryBasicBlock]
    source_location: SourceLocation | None

    @cached_property
    def block_map(self) -> Mapping[str, MemoryBasicBlock]:
        return {b.block_name: b for b in self.body}

    def get_block(self, block_name: str) -> MemoryBasicBlock:
        return self.block_map[block_name]

@attrs.define
class Program:
    kind: ProgramKind
    main: MemorySubroutine
    subroutines: list[MemorySubroutine]
    avm_version: int

    @property
    def all_subroutines(self) -> Iterator[MemorySubroutine]:
        yield self.main
        yield from self.subroutines

def produces_from_desc(desc: str, size: int) -> Sequence[str]:
    desc = f"{{{desc}}}"
    if size > 1:
        produces = [f"{desc}.{n}" for n in range(size)]
    elif size == 1:
        produces = [desc]
    else:
        produces = []
    return produces


import textwrap

import attrs

from puya import log
from puya.context import ArtifactCompileContext, CompileContext
from puya.mir import models
from puya.mir.aligned_writer import AlignedWriter
from puya.mir.stack import Stack
from puya.parse import SourceLocation

logger = log.get_logger(__name__)

def output_memory_ir(
    ctx: ArtifactCompileContext, program: models.Program, *, qualifier: str
) -> None:
    output_path = ctx.build_output_path(program.kind, qualifier, "mir")
    if output_path is None:
        return
    writer = AlignedWriter()
    writer.add_header("// Op")
    writer.add_header("Stack (out)", 4)
    for subroutine in program.all_subroutines:
        writer.append_line(f"// {subroutine.signature}")
        for block in subroutine.body:
            stack = Stack.begin_block(subroutine, block)
            last_location = None
            writer.append(f"{block.block_name}:")
            writer.append(stack.full_stack_desc)
            writer.new_line()
            with writer.indent():
                for op in block.ops:
                    last_location = _output_src_comment(
                        ctx, writer, last_location, op.source_location
                    )
                    op_str = str(op)
                    op.accept(stack)
                    # some ops can be very long (generally due to labels)
                    # in those (rare?) cases bypass the column alignment
                    if len(op_str) > 80:
                        writer.append_line(
                            writer.current_indent + op_str + " " + stack.full_stack_desc
                        )
                    else:
                        writer.append(op_str)
                        writer.append(stack.full_stack_desc)
                        writer.new_line()
                writer.new_line()
        writer.new_line()
    writer.new_line()
    output_path.write_text("\n".join(writer.write()), "utf8")

def _output_src_comment(
    ctx: CompileContext,
    writer: AlignedWriter,
    last_loc: SourceLocation | None,
    op_loc: SourceLocation | None,
) -> SourceLocation | None:
    if op_loc:
        whole_lines_location = attrs.evolve(op_loc, column=None, end_column=None)
        if whole_lines_location != last_loc:
            last_loc = whole_lines_location
            src = ctx.try_get_source(whole_lines_location)
            if src is not None:
                writer.append(f"// {whole_lines_location}")
                writer.new_line()
                lines = textwrap.dedent("\n".join(src)).splitlines()
                for line in lines:
                    writer.append(f"// {line.rstrip()}")
                    writer.new_line()
    return last_loc


import typing
from collections.abc import Sequence

import attrs

from puya.mir import models
from puya.mir.visitor import MIRVisitor

@attrs.define
class Stack(MIRVisitor[None]):
    parameters: Sequence[str]
    _f_stack: list[str]
    """f-stack holds variables above the current frame"""
    _x_stack: list[str]
    """x-stack holds variable that are carried between blocks"""
    _l_stack: list[str] = attrs.field(factory=list)
    """l-stack holds variables that are used within a block"""

    @classmethod
    def begin_block(
        cls, subroutine: models.MemorySubroutine, block: models.MemoryBasicBlock
    ) -> typing.Self:
        return cls(
            parameters=[p.local_id for p in subroutine.signature.parameters],
            f_stack=list(block.f_stack_in),
            x_stack=list(block.x_stack_in or ()),  # x-stack might not be assigned yet
        )

    @property
    def f_stack(self) -> Sequence[str]:
        return self._f_stack

    @property
    def x_stack(self) -> Sequence[str]:
        return self._x_stack

    @property
    def l_stack(self) -> Sequence[str]:
        return self._l_stack

    @property
    def xl_height(self) -> int:
        return len(self._l_stack) + len(self._x_stack)

    @property
    def fxl_height(self) -> int:
        return len(self._f_stack) + len(self._l_stack) + len(self._x_stack)

    @property
    def full_stack_desc(self) -> str:
        stack_descs = []
        if self.parameters:
            stack_descs.append("(𝕡) " + ",".join(self.parameters))  # noqa: RUF001
        if self._f_stack:
            stack_descs.append("(𝕗) " + ",".join(self._f_stack))  # noqa: RUF001
        if self._x_stack:
            stack_descs.append("(𝕏) " + ",".join(self._x_stack))  # noqa: RUF001
        stack_descs.append(",".join(self._l_stack))
        return " | ".join(stack_descs)

    def _get_f_stack_dig_bury(self, value: str) -> int:
        return (
            len(self._f_stack)
            + len(self._x_stack)
            + len(self._l_stack)
            - self._f_stack.index(value)
            - 1
        )

    def visit_int(self, const: models.Int) -> None:
        self._apply_lstack_effects(const)

    def visit_byte(self, const: models.Byte) -> None:
        self._apply_lstack_effects(const)

    def visit_undefined(self, const: models.Undefined) -> None:
        self._apply_lstack_effects(const)

    def visit_template_var(self, const: models.TemplateVar) -> None:
        self._apply_lstack_effects(const)

    def visit_address(self, const: models.Address) -> None:
        self._apply_lstack_effects(const)

    def visit_method(self, const: models.Method) -> None:
        self._apply_lstack_effects(const)

    def visit_comment(self, _: models.Comment) -> None:
        pass

    def visit_abstract_store(self, store: models.AbstractStore) -> None:
        self._apply_lstack_effects(store)

    def visit_abstract_load(self, load: models.AbstractLoad) -> None:
        self._apply_lstack_effects(load)

    def _store_f_stack(self, store: models.StoreFStack) -> None:
        local_id = store.local_id
        # must calculate bury offsets BEFORE modifying l-stack
        assert local_id in self._f_stack, f"{local_id} not in f-stack"
        bury = self._get_f_stack_dig_bury(local_id)
        assert bury == store.depth, f"expected {bury=} == {store.depth=}"
        self._apply_lstack_effects(store)

    def _insert_f_stack(self, store: models.StoreFStack) -> None:
        local_id = store.local_id
        assert local_id not in self._f_stack, f"{local_id} already in f-stack"
        # inserting something at the top of the f-stack
        # is equivalent to inserting at the bottom of the x-stack
        cover = len(self._x_stack) + len(self._l_stack) - 1
        assert cover == store.depth, f"expected {cover=} == {store.depth=}"
        self._f_stack.append(local_id)
        self._apply_lstack_effects(store)

    def visit_store_f_stack(self, store: models.StoreFStack) -> None:
        assert self._l_stack, f"l-stack is empty, can not store {store.local_id} to f-stack"
        if store.insert:
            self._insert_f_stack(store)
        else:
            self._store_f_stack(store)

    def visit_load_f_stack(self, load: models.LoadFStack) -> None:
        local_id = load.local_id
        assert local_id in self._f_stack, f"{local_id} not found in f-stack"
        dig = self._get_f_stack_dig_bury(local_id)
        assert dig == load.depth, f"expected {dig=} == {load.depth=}"
        self._apply_lstack_effects(load)

    def visit_store_x_stack(self, store: models.StoreXStack) -> None:
        local_id = store.local_id
        assert self._l_stack, f"l-stack too small to store {local_id} to x-stack"

        cover = len(self._x_stack) + len(self._l_stack) - 1
        assert cover == store.depth, f"expected {cover=} == {store.depth=}"
        self._x_stack.insert(0, local_id)
        self._apply_lstack_effects(store)

    def visit_load_x_stack(self, load: models.LoadXStack) -> None:
        local_id = load.local_id
        assert local_id in self._x_stack, f"{local_id} not found in x-stack"
        index = self._x_stack.index(local_id)
        uncover = len(self._l_stack) + len(self._x_stack) - index - 1
        assert uncover == load.depth, f"expected {uncover=} == {load.depth=}"
        self._x_stack.pop(index)
        self._apply_lstack_effects(load)

    def visit_store_l_stack(self, store: models.StoreLStack) -> None:
        cover = store.depth
        local_id = store.local_id
        assert cover < len(
            self._l_stack
        ), f"l-stack too small to store (cover {cover}) {store.local_id} to l-stack"
        index = len(self._l_stack) - cover - 1
        self._l_stack.pop()
        self._l_stack.insert(index, local_id)
        self._apply_lstack_effects(store)

    def visit_load_l_stack(self, load: models.LoadLStack) -> None:
        local_id = load.local_id
        uncover = load.depth
        if uncover is None:  # during l-stack construction depth is not fixed
            index = self._l_stack.index(local_id)
        else:
            index = len(self._l_stack) - uncover - 1
        if not load.copy:
            self._l_stack.pop(index)
        self._apply_lstack_effects(load)

    def visit_load_param(self, load: models.LoadParam) -> None:
        assert load.local_id in self.parameters, f"{load.local_id} is not a parameter"
        self._apply_lstack_effects(load)

    def visit_store_param(self, store: models.StoreParam) -> None:
        assert store.local_id in self.parameters, f"{store.local_id} is not a parameter"
        self._apply_lstack_effects(store)

    def visit_allocate(self, allocate: models.Allocate) -> None:
        self._f_stack.extend(allocate.allocate_on_entry)

    def visit_pop(self, pop: models.Pop) -> None:
        self._apply_lstack_effects(pop)

    def visit_callsub(self, callsub: models.CallSub) -> None:
        self._apply_lstack_effects(callsub)

    def visit_retsub(self, retsub: models.RetSub) -> None:
        assert len(self._l_stack) == retsub.returns, (
            f"Inconsistent l-stack height for retsub,"
            f" expected {retsub.returns}, actual: {len(self._l_stack)}"
        )

        # retsub moves return values down below the frame to the stack height before the sub was
        # called and discards anything above.
        # represent this in the virtual stack with a new stack state with only the current
        # l-stack (i.e. discard all values in parameters, f-stack and x-stack)
        self.parameters = []
        self._f_stack = []
        self._x_stack = []

    def visit_conditional_branch(self, op: models.ConditionalBranch) -> None:
        self._apply_lstack_effects(op)

    def visit_err(self, op: models.Err) -> None:
        self._apply_lstack_effects(op)

    def visit_goto(self, op: models.Goto) -> None:
        self._apply_lstack_effects(op)

    def visit_match(self, op: models.Match) -> None:
        self._apply_lstack_effects(op)

    def visit_program_exit(self, op: models.ProgramExit) -> None:
        self._apply_lstack_effects(op)

    def visit_switch(self, op: models.Switch) -> None:
        self._apply_lstack_effects(op)

    def visit_intrinsic(self, intrinsic: models.IntrinsicOp) -> None:
        self._apply_lstack_effects(intrinsic)

    def _apply_lstack_effects(self, op: models.BaseOp) -> None:
        assert len(self._l_stack) >= op.consumes, f"l-stack too small for {op}"
        start = len(self._l_stack) - op.consumes
        self._l_stack[start:] = op.produces


from puya.mir import models
from puya.mir.context import ProgramMIRContext, SubroutineCodeGenContext
from puya.mir.output import output_memory_ir
from puya.mir.stack_allocation.f_stack import f_stack_allocation
from puya.mir.stack_allocation.l_stack import l_stack_allocation
from puya.mir.stack_allocation.peephole import peephole_optimization_single_pass
from puya.mir.stack_allocation.x_stack import x_stack_allocation

# Note: implementation of http://www.euroforth.org/ef06/shannon-bailey06.pdf

def global_stack_allocation(ctx: ProgramMIRContext, program: models.Program) -> None:
    for desc, method in {
        "lstack": l_stack_allocation,
        "lstack.opt": _peephole_optimization,
        "xstack": x_stack_allocation,
        "xstack.opt": _peephole_optimization,
        "fstack": f_stack_allocation,
        "fstack.opt": _peephole_optimization,
    }.items():
        for mir_sub in program.all_subroutines:
            sub_ctx = ctx.for_subroutine(mir_sub)
            method(sub_ctx)
        if ctx.options.output_memory_ir:
            output_memory_ir(ctx, program, qualifier=desc)
    if ctx.options.output_memory_ir:
        output_memory_ir(ctx, program, qualifier="")

def _peephole_optimization(ctx: SubroutineCodeGenContext) -> None:
    # replace sequences of stack manipulations with shorter ones
    vla_modified = False
    for block in ctx.subroutine.body:
        result = peephole_optimization_single_pass(ctx, block)
        vla_modified = vla_modified or result.vla_modified
    if vla_modified:
        ctx.invalidate_vla()


import typing
from collections.abc import Sequence

import attrs

from puya import log
from puya.avm import AVMType
from puya.errors import InternalError
from puya.mir import models as mir
from puya.mir.context import SubroutineCodeGenContext
from puya.mir.stack import Stack
from puya.utils import attrs_extend

logger = log.get_logger(__name__)

def _get_lazy_fstack(subroutine: mir.MemorySubroutine) -> dict[str, mir.AbstractStore]:
    # TODO: consider more than the entry block
    entry = subroutine.body[0]
    # if entry is re-entrant then can't lazy allocate anything
    if entry.predecessors:
        return {}
    result = dict[str, mir.AbstractStore]()
    for op in entry.ops:
        if isinstance(op, mir.AbstractStore):
            result.setdefault(op.local_id, op)
    return result

def _get_local_id_types(subroutine: mir.MemorySubroutine) -> dict[str, AVMType]:
    variable_mapping = dict[str, AVMType]()
    for block in subroutine.body:
        for op in block.ops:
            if isinstance(op, mir.AbstractStore):
                try:
                    existing_type = variable_mapping[op.local_id]
                except KeyError:
                    existing_type = op.atype
                variable_mapping[op.local_id] = existing_type | op.atype
    return variable_mapping

def _get_allocate_op(
    subroutine: mir.MemorySubroutine, all_variables: Sequence[str]
) -> mir.Allocate:
    # determine variables to allocate at beginning of frame,
    # and order them so bytes are listed first, followed by uints
    byte_vars = []
    uint64_vars = []
    variable_type_mapping = _get_local_id_types(subroutine)
    for variable in all_variables:
        match variable_type_mapping.get(variable):
            case AVMType.uint64:
                uint64_vars.append(variable)
            case AVMType.bytes:
                byte_vars.append(variable)
            case AVMType.any:
                raise InternalError(
                    "Encountered AVM type any on preamble construction",
                    subroutine.source_location,
                )
            case None:
                # shouldn't occur, undefined variables should still have an Undefined entry
                # with a type
                raise InternalError(f"Undefined register: {variable}", subroutine.source_location)
            case unexpected:
                typing.assert_never(unexpected)
    return mir.Allocate(bytes_vars=byte_vars, uint64_vars=uint64_vars)

def f_stack_allocation(ctx: SubroutineCodeGenContext) -> None:
    all_variables = ctx.vla.all_variables
    if not all_variables:
        return

    subroutine = ctx.subroutine
    first_store_ops = _get_lazy_fstack(subroutine)
    allocate_on_first_store = list(first_store_ops)

    unsorted_pre_allocate = [x for x in all_variables if x not in first_store_ops]
    if unsorted_pre_allocate:
        allocate = _get_allocate_op(subroutine, unsorted_pre_allocate)
        allocate_at_entry = allocate.allocate_on_entry
        subroutine.body[0].mem_ops.insert(0, allocate)
    else:
        allocate_at_entry = []
    logger.debug(f"{subroutine.signature.name} f-stack entry: {allocate_at_entry}")
    logger.debug(f"{subroutine.signature.name} f-stack on first store: {allocate_on_first_store}")

    all_f_stack = [*allocate_at_entry, *first_store_ops.keys()]
    subroutine.body[0].f_stack_out = all_f_stack
    for block in subroutine.body[1:]:
        block.f_stack_in = block.f_stack_out = all_f_stack

    removed_virtual = False
    for block in subroutine.body:
        stack = Stack.begin_block(subroutine, block)
        for index, op in enumerate(block.mem_ops):
            match op:
                case mir.AbstractStore() as store:
                    insert = op in first_store_ops.values()
                    if insert:
                        depth = stack.xl_height - 1
                    else:
                        depth = stack.fxl_height - stack.f_stack.index(store.local_id) - 1

                    block.mem_ops[index] = op = attrs_extend(
                        mir.StoreFStack,
                        store,
                        depth=depth,
                        frame_index=stack.fxl_height - depth - 1,
                        insert=insert,
                    )
                    removed_virtual = True
                case mir.AbstractLoad() as load:
                    depth = stack.fxl_height - stack.f_stack.index(load.local_id) - 1
                    block.mem_ops[index] = op = attrs_extend(
                        mir.LoadFStack,
                        load,
                        depth=depth,
                        frame_index=stack.fxl_height - depth - 1,
                    )
                    removed_virtual = True
            op.accept(stack)
        match block.terminator:
            case mir.RetSub() as retsub:
                block.terminator = attrs.evolve(
                    retsub, fx_height=len(stack.f_stack) + len(stack.x_stack)
                )
    if removed_virtual:
        ctx.invalidate_vla()


import itertools

import attrs

from puya import log
from puya.mir import models as mir
from puya.mir.context import SubroutineCodeGenContext
from puya.mir.stack import Stack

logger = log.get_logger(__name__)

@attrs.define
class UsagePair:
    a: mir.AbstractLoad | mir.AbstractStore
    b: mir.AbstractLoad
    a_index: int
    b_index: int

    @staticmethod
    def by_distance(pair: "UsagePair") -> tuple[int, int, int]:
        return pair.b_index - pair.a_index, pair.a_index, pair.b_index

def l_stack_allocation(ctx: SubroutineCodeGenContext) -> None:
    # the following is basically koopmans algorithm
    # done as part of http://www.euroforth.org/ef06/shannon-bailey06.pdf
    # see also https://users.ece.cmu.edu/~koopman/stack_compiler/stack_co.html#appendix
    for block in ctx.subroutine.body:
        usage_pairs = _find_usage_pairs(block)
        _copy_usage_pairs(ctx, block, usage_pairs)
    for block in ctx.subroutine.body:
        _dead_store_removal(ctx, block)
        if ctx.options.optimization_level:
            _implicit_store_removal(block)
    # update vla after dead store removal
    ctx.invalidate_vla()
    # calculate load depths now that l-stack allocations are done
    for block in ctx.subroutine.body:
        _calculate_load_depths(ctx, block)

def _find_usage_pairs(block: mir.MemoryBasicBlock) -> list[UsagePair]:
    # find usage pairs of variables within the block
    # the first element of the pair is an op that defines or uses a variable
    # the second element of the pair is an op that uses the variable
    # return pairs in ascending order, based on the number of instruction between each pair
    variables = dict[str, list[tuple[int, mir.AbstractStore | mir.AbstractLoad]]]()
    for index, op in enumerate(block.ops):
        match op:
            case mir.AbstractStore(local_id=local_id) | mir.AbstractLoad(local_id=local_id):
                variables.setdefault(local_id, []).append((index, op))

    pairs = list[UsagePair]()
    for uses in variables.values():
        # pairwise iteration means an op can only be in at most 2 pairs
        for (a_index, a), (b_index, b) in itertools.pairwise(uses):
            if isinstance(b, mir.AbstractStore):
                continue  # skip redefines, if they are used they will be picked up in next pair
            pairs.append(UsagePair(a=a, b=b, a_index=a_index, b_index=b_index))

    return sorted(pairs, key=UsagePair.by_distance)

def _copy_usage_pairs(
    ctx: SubroutineCodeGenContext, block: mir.MemoryBasicBlock, pairs: list[UsagePair]
) -> None:
    # 1. copy define or use to bottom of l-stack
    # 2. replace usage with instruction to rotate the value from the bottom of l-stack to the top

    # to copy: dup, cover {stack_height}
    # to rotate: uncover {stack_height} - 1
    replaced_ops = dict[mir.StoreOp | mir.LoadOp, mir.LoadOp]()
    for pair in pairs:
        # note: pairs may refer to ops that have been replaced by an earlier iteration
        a = replaced_ops.get(pair.a, pair.a)
        b = replaced_ops.get(pair.b, pair.b)
        local_id = a.local_id
        # step 1. copy define or use to bottom of stack

        # redetermine index as block ops may have changed
        a_index = block.mem_ops.index(a)

        # insert replacement before store, or after load
        insert_index = a_index if isinstance(a, mir.AbstractStore) else a_index + 1
        stack = Stack.begin_block(ctx.subroutine, block)
        for op in block.mem_ops[:insert_index]:
            op.accept(stack)
        dup = mir.StoreLStack(
            depth=len(stack.l_stack) - 1,
            local_id=local_id,
            # leave a copy for the original consumer of this value which is either:
            #   a.) the virtual store we are inserting before
            #   b.) whatever came after the virtual load we are inserting after
            # The copy will be eliminated during dead store removal if no longer required
            copy=True,
            source_location=a.source_location,
            atype=a.atype,
        )
        block.mem_ops.insert(insert_index, dup)
        logger.debug(f"Inserted {block.block_name}.ops[{insert_index}]: '{dup}'")

        # step 2. replace b usage with instruction to rotate the value from the bottom of the stack
        # determine index of b, as inserts may have shifted its location
        b_index = block.mem_ops.index(b)

        uncover = mir.LoadLStack(
            # can not determine depth yet
            # as it depends on any other l-stack operations between the store and this load
            # which could change until after dead store removal is complete
            depth=None,
            local_id=local_id,
            copy=False,
            source_location=b.source_location,
            atype=b.atype,
        )
        # replace op
        block.mem_ops[b_index] = uncover

        # remember replacement in case it is part of another pair
        # an op can only be at most in 2 pairs, so don't need to do this recursively
        replaced_ops[b] = uncover
        logger.debug(f"Replaced {block.block_name}.ops[{b_index}]: '{b}' with '{uncover}'")

def _dead_store_removal(ctx: SubroutineCodeGenContext, block: mir.MemoryBasicBlock) -> None:
    ops = block.mem_ops
    op_idx = 0
    while op_idx < len(ops) - 1:
        window = slice(op_idx, op_idx + 2)
        a, b = ops[window]
        if (
            isinstance(a, mir.StoreLStack)
            and a.copy
            and isinstance(b, mir.AbstractStore)
            and b.local_id not in ctx.vla.get_live_out_variables(b)
        ):
            # StoreLStack is used to:
            #   1.) create copy of the value to be immediately stored via virtual store
            #   2.) rotate the value to the bottom of the stack for use in a later op in this block
            # If it is a dead store, then the 1st scenario is no longer needed
            # and instead just need to ensure the value is moved to the bottom of the stack
            a = attrs.evolve(a, copy=False, produces=())
            ops[window] = [a]
        elif (
            isinstance(a, mir.LoadLStack)
            and not a.copy
            and isinstance(b, mir.StoreLStack)
            and b.copy
            and a.local_id == b.local_id
        ):
            a = attrs.evolve(
                a,
                copy=True,
                produces=(f"{a.local_id} (copy)",),
            )
            ops[window] = [a]
        op_idx += 1

def _implicit_store_removal(block: mir.MemoryBasicBlock) -> None:
    ops = block.mem_ops
    op_idx = 0
    while op_idx < len(ops):
        op = ops[op_idx]
        # see if ops immediately after this op are all storing to the l-stack what this op produces
        next_op_idx = op_idx + 1
        maybe_remove_window = slice(next_op_idx, next_op_idx + len(op.produces))
        maybe_remove = [
            maybe_store
            for maybe_store in ops[maybe_remove_window]
            if isinstance(maybe_store, mir.StoreLStack)
            and not maybe_store.copy
            and maybe_store.local_id in op.produces
        ]
        # if they all match then this means all values are implicitly on the l-stack
        # and we can safely remove the store ops
        if len(maybe_remove) == len(op.produces):
            ops[maybe_remove_window] = []
        op_idx = next_op_idx

def _calculate_load_depths(ctx: SubroutineCodeGenContext, block: mir.MemoryBasicBlock) -> None:
    stack = Stack.begin_block(ctx.subroutine, block)
    for idx, op in enumerate(block.mem_ops):
        if isinstance(op, mir.LoadLStack):
            local_id_index = stack.l_stack.index(op.local_id)
            block.mem_ops[idx] = attrs.evolve(op, depth=len(stack.l_stack) - local_id_index - 1)
        op.accept(stack)


from collections.abc import Sequence

import attrs

from puya.mir import models as mir
from puya.mir.context import SubroutineCodeGenContext

def optimize_pair(
    ctx: SubroutineCodeGenContext,
    a: mir.Op,
    b: mir.Op,
) -> Sequence[mir.Op] | None:
    """Given a pair of ops, returns which ops should be kept including replacements"""
    # this function has been optimized to reduce the number of isinstance checks,
    # consider this when making any modifications

    # move local_ids to produces of previous op where possible
    if (
        isinstance(b, mir.StoreLStack | mir.StoreXStack | mir.StoreFStack)
        and len(a.produces)
        and a.produces[-1] != b.local_id
    ):
        a = attrs.evolve(a, produces=(*a.produces[:-1], b.local_id))
        return a, b

    # remove redundant stores and loads
    if a.produces and a.produces[-1] == _get_local_id_alias(b):
        return (a,)

    if isinstance(b, mir.AbstractStore) and b.local_id not in ctx.vla.get_live_out_variables(b):
        # note l-stack dead store removal occurs during l-stack allocation
        # this handles any other cases
        return a, mir.Pop(n=1, source_location=b.source_location)

    if isinstance(a, mir.LoadOp) and isinstance(b, mir.StoreOp) and a.local_id == b.local_id:
        match a, b:
            case mir.LoadXStack(), mir.StoreXStack():
                return ()
            case mir.LoadFStack(), mir.StoreFStack():
                return ()
            case mir.AbstractLoad(), mir.AbstractStore():
                # this is used see test_cases/bug_load_store_load_store
                return ()
    return None

@attrs.define(kw_only=True)
class PeepholeResult:
    modified: bool
    vla_modified: bool

def peephole_optimization_single_pass(
    ctx: SubroutineCodeGenContext, block: mir.MemoryBasicBlock
) -> PeepholeResult:
    result = block.mem_ops
    op_idx = 0
    modified = False
    vla_modified = False
    while op_idx < len(result) - 1:
        window = slice(op_idx, op_idx + 2)
        curr_op, next_op = result[window]
        pair_result = optimize_pair(ctx, curr_op, next_op)
        if pair_result is not None:
            modified = True
            result[window] = pair_result
            # check if VLA needs updating
            vla_modified = (
                vla_modified
                or (
                    curr_op not in pair_result
                    and isinstance(curr_op, mir.AbstractStore | mir.AbstractLoad)
                )
                or (
                    next_op not in pair_result
                    and isinstance(next_op, mir.AbstractStore | mir.AbstractLoad)
                )
            )
        else:  # if nothing optimized, then advance
            op_idx += 1
    return PeepholeResult(modified=modified, vla_modified=vla_modified)

def _get_local_id_alias(op: mir.BaseOp) -> str | None:
    """Returns the local_id of a memory op if it has no effect
    apart from renaming the top variable on the stack"""
    if isinstance(op, mir.StoreLStack | mir.LoadLStack) and not op.copy and not op.depth:
        return op.local_id
    # TODO: the following can only be done if the movement between l-stack and the other stack
    #       is captured somehow (also check assumption that it needs to be captured...)
    # if isinstance(op, mir.StoreXStack | mir.LoadXStack) and not op.depth:
    #    return op.local_id
    # if isinstance(op, mir.StoreFStack) and op.insert and not op.depth:
    #    return op.local_id
    return None