import base64
import contextlib
import math
import operator
import typing
from collections import defaultdict, deque
from collections.abc import Callable, Container, Generator, Iterable, Mapping, Sequence, Set
from itertools import zip_longest

import attrs

from puya import algo_constants, log
from puya.avm import AVMType
from puya.ir import models
from puya.ir.avm_ops import AVMOp
from puya.ir.models import Intrinsic, UInt64Constant
from puya.ir.optimize.context import IROptimizationContext
from puya.ir.optimize.dead_code_elimination import SIDE_EFFECT_FREE_AVM_OPS
from puya.ir.register_read_collector import RegisterReadCollector
from puya.ir.types_ import AVMBytesEncoding, IRType
from puya.ir.visitor_mutator import IRMutator
from puya.parse import SourceLocation
from puya.utils import biguint_bytes_eval, method_selector_hash, set_add

logger = log.get_logger(__name__)

_AnyOp = models.Op | models.ControlOp | models.Phi

_RegisterAssignments = Mapping[models.Register, models.Assignment]

COMPILE_TIME_CONSTANT_OPS = frozenset(
    [
        # "generic" comparison ops
        "==",
        "!=",
        # uint64 comparison ops
        "<",
        "<=",
        ">",
        ">=",
        # boolean ops
        "!",
        "&&",
        "||",
        # uint64 bitwise ops
        "&",
        "|",
        "^",
        "~",
        "shl",
        "shr",
        # uint64 math
        "+",
        "-",
        "*",
        "/",
        "%",
        "exp",
        "sqrt",
        # wide math
        "addw",
        "mulw",
        "divw",
        "expw",
        "divmodw",
        # bit/byte ops
        "concat",
        "extract",
        "extract3",
        "getbit",
        "getbyte",
        "len",
        "replace2",
        "replace3",
        "setbit",
        "setbyte",
        "substring",
        "substring3",
        # conversion
        "itob",
        "btoi",
        "extract_uint16",
        "extract_uint32",
        "extract_uint64",
        # byte math
        "b+",
        "b-",
        "b*",
        "b/",
        "b%",
        "bsqrt",
        # byte comaprison ops
        "b==",
        "b!=",
        "b<",
        "b<=",
        "b>",
        "b>=",
        # byte bitwise ops
        "b&",
        "b|",
        "b^",
        "b~",
        # misc
        "assert",
        "bzero",
        "select",
        "bitlen",
        # ! unimplemented for constant arg evaluation
        "base64_decode",
        "ec_add",
        "ec_map_to",
        "ec_multi_scalar_mul",
        "ec_pairing_check",
        "ec_scalar_mul",
        "ec_subgroup_check",
        "ecdsa_pk_decompress",
        "ecdsa_pk_recover",
        "ecdsa_verify",
        "ed25519verify",
        "ed25519verify_bare",
        "json_ref",
        "keccak256",
        "sha256",
        "sha3_256",
        "sha512_256",
        "vrf_verify",
        "sumhash512",  # AVM 11
    ]
)
_CONSTANT_EVALUABLE: typing.Final[frozenset[str]] = COMPILE_TIME_CONSTANT_OPS - {
    AVMOp.itob.code,
    AVMOp.bzero.code,
}

def intrinsic_simplifier(context: IROptimizationContext, subroutine: models.Subroutine) -> bool:
    if context.expand_all_bytes:
        work_list = _AssignmentWorkQueue(COMPILE_TIME_CONSTANT_OPS)
    else:
        work_list = _AssignmentWorkQueue(_CONSTANT_EVALUABLE)
    ssa_reads = _SSAReadTracker()

    register_assignments = dict[models.Register, models.Assignment]()
    for block in subroutine.body:
        for op in block.all_ops:
            ssa_reads.add(op)
            if isinstance(op, models.Assignment):
                work_list.enqueue(op)
                if len(op.targets) == 1:
                    (target,) = op.targets
                    register_assignments[target] = op

    modified = 0
    while work_list:
        ass, source = work_list.dequeue()
        simplified = _try_fold_intrinsic(register_assignments, source)
        if simplified is not None:
            logger.debug(f"Simplified {source} to {simplified}")
            with ssa_reads.update(ass):
                ass.source = simplified
            modified += 1
            # if it became a Value, propagate to any assignment-readers and add to work list
            if isinstance(simplified, models.Value):
                (target,) = ass.targets
                replacer = _RegisterValueReplacer(register=target, replacement=simplified)
                for target_read in ssa_reads.get(target, copy=True):
                    if isinstance(target_read, models.Assignment):
                        work_list.enqueue(target_read)
                        # special case for indirection of otherwise non-inlined constants
                        match target_read:
                            case models.Assignment(
                                targets=[target_read_target],
                                source=models.Intrinsic(op=(AVMOp.bzero | AVMOp.itob)),
                            ) if not context.expand_all_bytes:
                                for indirect_target_read in ssa_reads.get(target_read_target):
                                    if isinstance(indirect_target_read, models.Assignment):
                                        work_list.enqueue(indirect_target_read)
                    with ssa_reads.update(target_read):
                        target_read.accept(replacer)
                modified += replacer.modified
            else:
                typing.assert_type(simplified, models.Intrinsic)
                # source is still an intrinsic, add it back to the work list
                work_list.enqueue(ass)
                # add any assignment-readers to the work list
                for target in ass.targets:
                    for target_read in ssa_reads.get(target):
                        if isinstance(target_read, models.Assignment):
                            work_list.enqueue(target_read)

    for block in subroutine.body:
        for op in block.ops:
            if (
                isinstance(op, models.Assignment)
                and isinstance(op.source, models.Intrinsic)
                and op.source.args
            ):
                with_immediates = _try_convert_stack_args_to_immediates(op.source)
                if with_immediates is not None:
                    logger.debug(f"Simplified {op.source} to {with_immediates}")
                    op.source = with_immediates
                    modified += 1

    register_intrinsics = {
        target: ass.source
        for target, ass in register_assignments.items()
        if isinstance(ass.source, models.Intrinsic)
    }
    modified += _simplify_conditional_branches(subroutine, register_intrinsics)
    modified += _simplify_non_returning_intrinsics(subroutine, register_intrinsics)
    return modified > 0

class _AssignmentWorkQueue:
    def __init__(self, constant_evaluable: Container[str]) -> None:
        self._constant_evaluable = constant_evaluable
        self._dq = deque[tuple[models.Assignment, models.Intrinsic]]()
        self._set = set[Sequence[models.Register]]()

    def enqueue(self, op: models.Assignment) -> bool:
        if (
            isinstance(op.source, models.Intrinsic)
            and op.source.op.code in self._constant_evaluable
            and set_add(self._set, op.targets)
        ):
            self._dq.append((op, op.source))
            return True
        return False

    def dequeue(self) -> tuple[models.Assignment, models.Intrinsic]:
        op, source = self._dq.popleft()
        assert source is op.source
        self._set.remove(op.targets)
        return op, source

    def __bool__(self) -> int:
        return bool(self._dq)

class _SSAReadTracker:
    def __init__(self) -> None:
        self._data = defaultdict[models.Register, set[_AnyOp]](set)

    def add(self, op: _AnyOp) -> None:
        for read_reg in self._register_reads(op):
            self._data[read_reg].add(op)

    def get(self, reg: models.Register, *, copy: bool = False) -> Iterable[_AnyOp]:
        reads = self._data.get(reg)
        if reads is None:
            return ()
        if copy:
            return reads.copy()
        return reads

    @contextlib.contextmanager
    def update(self, op: _AnyOp) -> Generator[None, None, None]:
        old_reads = self._register_reads(op)
        yield
        new_reads = self._register_reads(op)
        for removed_read in old_reads - new_reads:
            self._data[removed_read].remove(op)
        for added_read in new_reads - old_reads:
            self._data[added_read].add(op)

    @staticmethod
    def _register_reads(visitable: models.IRVisitable) -> Set[models.Register]:
        collector = RegisterReadCollector()
        visitable.accept(collector)
        return collector.used_registers

@attrs.define(kw_only=True)
class _RegisterValueReplacer(IRMutator):
    register: models.Register
    replacement: models.Value
    modified: int = 0

    @typing.override
    def visit_assignment(self, ass: models.Assignment) -> models.Assignment:
        # don't visit target(s), needs to stay as Register
        ass.source = ass.source.accept(self)
        return ass

    @typing.override
    def visit_phi(self, phi: models.Phi) -> models.Phi:
        # don't visit phi nodes, needs to stay as Register
        return phi

    @typing.override
    def visit_register(self, reg: models.Register) -> models.Value:  # type: ignore[override]
        if reg != self.register:
            return reg
        self.modified += 1
        return self.replacement

def _simplify_conditional_branches(
    subroutine: models.Subroutine, register_intrinsics: Mapping[models.Register, models.Intrinsic]
) -> int:
    modified = 0
    branch_registers = dict[
        models.Register, list[tuple[models.ConditionalBranch, models.BasicBlock]]
    ]()
    for block in subroutine.body:
        match block.terminator:
            case (
                models.ConditionalBranch(condition=models.Register() as cond) as branch
            ) if cond in register_intrinsics:
                branch_registers.setdefault(cond, []).append((branch, block))
    for target, usages in branch_registers.items():
        intrinsic = register_intrinsics[target]
        cond_maybe_simplified = _try_simplify_bool_intrinsic(intrinsic)
        if cond_maybe_simplified is not None:
            for branch, used_block in usages:
                used_block.terminator = attrs.evolve(branch, condition=cond_maybe_simplified)
                modified += 1
    return modified

def _simplify_non_returning_intrinsics(
    subroutine: models.Subroutine, register_intrinsics: Mapping[models.Register, models.Intrinsic]
) -> int:
    modified = 0
    for block in subroutine.body:
        ops = []
        for op in block.ops:
            if not isinstance(op, models.Intrinsic):
                ops.append(op)
            else:
                result = _visit_intrinsic_op(op, register_intrinsics)
                if result is not op:
                    modified += 1
                if result is not None:
                    ops.append(result)
        block.ops[:] = ops
    return modified

def _visit_intrinsic_op(
    intrinsic: Intrinsic, register_intrinsics: Mapping[models.Register, models.Intrinsic]
) -> Intrinsic | None:
    # if we get here, it means either the intrinsic doesn't have a return or it's ignored,
    # in either case, the result has to be either an Op or None (ie delete),
    # so we don't invoke _try_fold_intrinsic here
    if intrinsic.op == AVMOp.assert_:
        (cond,) = intrinsic.args
        if isinstance(cond, models.UInt64Constant):
            value = cond.value
            if value:
                return None
            else:
                # an assert 0 could be simplified to an err, but
                # this would make it a ControlOp, so the block would
                # need to be restructured
                pass
        elif cond in register_intrinsics:
            cond_op = register_intrinsics[cond]  # type: ignore[index]
            assert_cond_maybe_simplified = _try_simplify_bool_intrinsic(cond_op)
            if assert_cond_maybe_simplified is not None:
                return attrs.evolve(intrinsic, args=[assert_cond_maybe_simplified])
        return intrinsic
    elif intrinsic.op == AVMOp.itxn_field:
        (field_im,) = intrinsic.immediates
        if field_im in ("ApprovalProgramPages", "ClearStateProgramPages"):
            (page_value,) = intrinsic.args
            if isinstance(page_value, models.BytesConstant) and page_value.value == b"":
                return None
        return intrinsic
    elif intrinsic.op.code in SIDE_EFFECT_FREE_AVM_OPS:
        logger.debug(f"Removing unused pure op {intrinsic}")
        return None
    elif intrinsic.args:
        simplified = _try_convert_stack_args_to_immediates(intrinsic)
        if simplified is not None:
            return simplified
        else:
            return intrinsic
    else:
        return intrinsic

def _try_simplify_bool_condition(
    register_assignments: _RegisterAssignments, cond: models.Value
) -> models.Value | None:
    if cond in register_assignments:
        cond_defn = register_assignments[cond]  # type: ignore[index]
        return _try_simplify_bool_intrinsic(cond_defn.source)
    return None

def _try_simplify_bool_intrinsic(cond_op: models.ValueProvider) -> models.Value | None:
    match cond_op:
        case (
            models.Intrinsic(
                args=[
                    models.Value(atype=AVMType.uint64) as a,
                    models.Value(atype=AVMType.uint64) as b,
                ]
            ) as intrinsic
        ):
            cond_maybe_simplified = _try_simplify_uint64_binary_op(
                intrinsic, a, b, bool_context=True
            )
            if isinstance(cond_maybe_simplified, models.Value):
                return cond_maybe_simplified
    return None

def _try_convert_stack_args_to_immediates(intrinsic: Intrinsic) -> Intrinsic | None:
    match intrinsic:
        case Intrinsic(
            op=AVMOp.gitxnas,
            args=[models.UInt64Constant(value=array_index)],
            immediates=[group_index, field],
        ):
            return attrs.evolve(
                intrinsic,
                op=AVMOp.gitxna,
                args=[],
                immediates=[group_index, field, array_index],
            )
        case Intrinsic(
            op=AVMOp.itxnas,
            args=[models.UInt64Constant(value=array_index)],
            immediates=[field],
        ):
            return attrs.evolve(
                intrinsic,
                op=AVMOp.itxna,
                args=[],
                immediates=[field, array_index],
            )
        case Intrinsic(
            op=(AVMOp.loads | AVMOp.stores as op),
            args=[models.UInt64Constant(value=slot), *rest],
        ):
            return attrs.evolve(
                intrinsic,
                immediates=[slot],
                args=rest,
                op=AVMOp.load if op == AVMOp.loads else AVMOp.store,
            )
        case Intrinsic(
            op=(AVMOp.extract3 | AVMOp.extract),
            args=[
                models.Value(atype=AVMType.bytes),
                models.UInt64Constant(value=S),
                models.UInt64Constant(value=L),
            ],
        ) if S <= 255 and 1 <= L <= 255:
            # note the lower bound of 1 on length, extract with immediates vs extract3
            # have *very* different behaviour if the length is 0
            return attrs.evolve(
                intrinsic, immediates=[S, L], args=intrinsic.args[:1], op=AVMOp.extract
            )
        case Intrinsic(
            op=AVMOp.substring3,
            args=[
                models.Value(atype=AVMType.bytes),
                models.UInt64Constant(value=S),
                models.UInt64Constant(value=E),
            ],
        ) if S <= 255 and E <= 255:
            return attrs.evolve(
                intrinsic, immediates=[S, E], args=intrinsic.args[:1], op=AVMOp.substring
            )
        case Intrinsic(
            op=AVMOp.replace3,
            args=[a, models.UInt64Constant(value=S), b],
        ) if S <= 255:
            return attrs.evolve(intrinsic, immediates=[S], args=[a, b], op=AVMOp.replace2)
        case Intrinsic(
            op=AVMOp.args,
            args=[models.UInt64Constant(value=idx)],
        ) if idx <= 255:
            return attrs.evolve(intrinsic, op=AVMOp.arg, immediates=[idx], args=[])
    return None

def _try_fold_intrinsic(
    register_assignments: _RegisterAssignments, intrinsic: models.Intrinsic
) -> models.Value | models.Intrinsic | None:
    op_loc = intrinsic.source_location
    if intrinsic.op is AVMOp.select:
        false, true, selector = intrinsic.args
        selector_const = _get_int_constant(selector)
        if selector_const is not None:
            return true if selector_const else false
        maybe_simplified_select_cond = _try_simplify_bool_condition(register_assignments, selector)
        if maybe_simplified_select_cond is not None:
            return attrs.evolve(intrinsic, args=[false, true, maybe_simplified_select_cond])
        if false == true:
            return true
        match (
            _get_byte_constant(register_assignments, false),
            _get_byte_constant(register_assignments, true),
        ):
            case (None, _) | (_, None):
                pass
            case (
                models.BytesConstant(value=false_value),
                models.BytesConstant(value=true_value) as true_bytes_const,
            ) if false_value == true_value:
                return true_bytes_const
        match _get_int_constant(false), _get_int_constant(true):
            case (None, _) | (_, None):
                pass
            case 1, 0:
                return attrs.evolve(intrinsic, op=AVMOp.not_, args=[selector])
            case 0, int(true_int_value) if selector.ir_type == IRType.bool:
                if true_int_value == 1:
                    return selector
                return attrs.evolve(intrinsic, op=AVMOp.mul, args=[selector, true])
            case 0, 1:
                zero_const = UInt64Constant(value=0, source_location=intrinsic.source_location)
                return attrs.evolve(intrinsic, op=AVMOp.neq, args=[selector, zero_const])
    elif intrinsic.op is AVMOp.replace2:
        (start,) = intrinsic.immediates
        assert isinstance(start, int)
        byte_arg_a, byte_arg_b = intrinsic.args
        if (byte_const_a := _get_byte_constant(register_assignments, byte_arg_a)) is not None and (
            byte_const_b := _get_byte_constant(register_assignments, byte_arg_b)
        ) is not None:
            replaced = bytearray(byte_const_a.value)
            replaced[start : start + len(byte_const_b.value)] = byte_const_b.value
            return models.BytesConstant(
                value=bytes(replaced),
                encoding=_choose_encoding(byte_const_a.encoding, byte_const_b.encoding),
                source_location=op_loc,
            )
    elif intrinsic.op is AVMOp.replace3:
        byte_arg_a, start_s, byte_arg_b = intrinsic.args
        if (
            (start2 := _get_int_constant(start_s)) is not None
            and (byte_const_a := _get_byte_constant(register_assignments, byte_arg_a)) is not None
            and (byte_const_b := _get_byte_constant(register_assignments, byte_arg_b)) is not None
        ):
            replaced = bytearray(byte_const_a.value)
            replaced[start2 : start2 + len(byte_const_b.value)] = byte_const_b.value
            return models.BytesConstant(
                value=bytes(replaced),
                encoding=_choose_encoding(byte_const_a.encoding, byte_const_b.encoding),
                source_location=op_loc,
            )
    elif intrinsic.op is AVMOp.getbit:
        match intrinsic.args:
            case [
                models.UInt64Constant(value=source, ir_type=IRType.uint64),
                models.UInt64Constant(value=index),
            ]:
                getbit_result = 1 if (source & (1 << index)) else 0
                return models.UInt64Constant(value=getbit_result, source_location=op_loc)
            case [
                models.Value(atype=AVMType.bytes) as byte_arg,
                models.UInt64Constant(value=index),
            ] if (byte_const := _get_byte_constant(register_assignments, byte_arg)) is not None:
                binary_array = [
                    x for xs in [bin(bb)[2:].zfill(8) for bb in byte_const.value] for x in xs
                ]
                the_bit = binary_array[index]
                return models.UInt64Constant(source_location=op_loc, value=int(the_bit))
    elif intrinsic.op is AVMOp.setbit:
        match intrinsic.args:
            case [
                models.UInt64Constant(value=source, ir_type=IRType.uint64),
                models.UInt64Constant(value=index),
                models.UInt64Constant(value=value),
            ]:
                if value:
                    setbit_result = source | (1 << index)
                else:
                    setbit_result = source & ~(1 << index)
                return models.UInt64Constant(value=setbit_result, source_location=op_loc)
            case [
                models.Value(atype=AVMType.bytes) as byte_arg,
                models.UInt64Constant(value=index),
                models.UInt64Constant(value=value),
            ] if (byte_const := _get_byte_constant(register_assignments, byte_arg)) is not None:
                binary_array = [
                    x for xs in [bin(bb)[2:].zfill(8) for bb in byte_const.value] for x in xs
                ]
                try:
                    binary_array[index] = "1" if value else "0"
                except IndexError:
                    return None  # would fail at runtime
                binary_string = "".join(binary_array)
                adjusted_const_value = int(binary_string, 2).to_bytes(
                    len(byte_const.value), byteorder="big"
                )
                return models.BytesConstant(
                    source_location=op_loc,
                    encoding=byte_const.encoding,
                    value=adjusted_const_value,
                )
            case [a, b, c]:
                bool_arg_maybe_simplified = _try_simplify_bool_condition(register_assignments, c)
                if bool_arg_maybe_simplified is not None:
                    return attrs.evolve(intrinsic, args=[a, b, bool_arg_maybe_simplified])
    elif intrinsic.op.code.startswith("extract_uint"):
        match intrinsic.args:
            case [
                models.BytesConstant(value=bytes_value),
                models.UInt64Constant(value=offset),
            ]:
                bit_size = int(intrinsic.op.code.removeprefix("extract_uint"))
                byte_size = bit_size // 8
                extracted = bytes_value[offset : offset + byte_size]
                if len(extracted) != byte_size:
                    return None  # would fail at runtime, lets hope this is unreachable 😬
                uint64_result = int.from_bytes(extracted, byteorder="big", signed=False)
                return models.UInt64Constant(
                    value=uint64_result,
                    source_location=op_loc,
                )
    elif intrinsic.op is AVMOp.concat:
        left_arg, right_arg = intrinsic.args
        left_const = _get_byte_constant(register_assignments, left_arg)
        right_const = _get_byte_constant(register_assignments, right_arg)
        if left_const is not None:
            if left_const.value == b"":
                return right_arg
            if right_const is not None:
                # two constants, just fold
                target_encoding = _choose_encoding(left_const.encoding, right_const.encoding)
                result_value = left_const.value + right_const.value
                result = models.BytesConstant(
                    value=result_value,
                    encoding=target_encoding,
                    source_location=op_loc,
                )
                return result
        elif right_const is not None:
            if right_const.value == b"":
                return left_arg
    elif intrinsic.op.code.startswith("extract"):
        match intrinsic:
            case (
                models.Intrinsic(
                    immediates=[int(S), int(L)],
                    args=[byte_arg],
                )
                | models.Intrinsic(
                    immediates=[],
                    args=[
                        byte_arg,
                        models.UInt64Constant(value=S),
                        models.UInt64Constant(value=L),
                    ],
                )
            ) if (byte_const := _get_byte_constant(register_assignments, byte_arg)) is not None:
                # note there is a difference of behaviour between extract with stack args
                # and with immediates - zero is to the end with immediates,
                # and zero length with stacks
                if intrinsic.immediates and L == 0:
                    extracted = byte_const.value[S:]
                else:
                    extracted = byte_const.value[S : S + L]
                return models.BytesConstant(
                    source_location=op_loc, encoding=byte_const.encoding, value=extracted
                )
    elif intrinsic.op.code.startswith("substring"):
        match intrinsic:
            case (
                models.Intrinsic(
                    immediates=[int(S), int(E)],
                    args=[byte_arg],
                )
                | models.Intrinsic(
                    immediates=[],
                    args=[
                        byte_arg,
                        models.UInt64Constant(value=S),
                        models.UInt64Constant(value=E),
                    ],
                )
            ) if (byte_const := _get_byte_constant(register_assignments, byte_arg)) is not None:
                if E < S:
                    return None  # would fail at runtime, lets hope this is unreachable 😬
                extracted = byte_const.value[S:E]
                return models.BytesConstant(
                    source_location=op_loc, encoding=byte_const.encoding, value=extracted
                )
    elif not intrinsic.immediates:
        match intrinsic.args:
            case [models.Value(atype=AVMType.uint64) as x]:
                return _try_simplify_uint64_unary_op(intrinsic, x)
            case [
                models.Value(atype=AVMType.uint64) as a,
                models.Value(atype=AVMType.uint64) as b,
            ]:
                return _try_simplify_uint64_binary_op(intrinsic, a, b)
            case [models.Value(atype=AVMType.bytes) as x]:
                return _try_simplify_bytes_unary_op(register_assignments, intrinsic, x)
            case [
                models.Value(atype=AVMType.bytes) as a,
                models.Value(atype=AVMType.bytes) as b,
            ]:
                return _try_simplify_bytes_binary_op(register_assignments, intrinsic, a, b)

    return None

def _get_int_constant(value: models.Value) -> int | None:
    if isinstance(value, models.UInt64Constant):
        return value.value
    return None

def _get_biguint_constant(
    register_assignments: _RegisterAssignments, value: models.Value
) -> tuple[int, bytes, AVMBytesEncoding] | tuple[None, None, None]:
    if isinstance(value, models.BigUIntConstant):
        return value.value, biguint_bytes_eval(value.value), AVMBytesEncoding.base16
    byte_const = _get_byte_constant(register_assignments, value)
    if byte_const is not None and len(byte_const.value) <= 64:
        return (
            int.from_bytes(byte_const.value, byteorder="big", signed=False),
            byte_const.value,
            byte_const.encoding,
        )
    return None, None, None

def _byte_wise(op: Callable[[int, int], int], lhs: bytes, rhs: bytes) -> bytes:
    return bytes([op(a, b) for a, b in zip_longest(lhs[::-1], rhs[::-1], fillvalue=0)][::-1])

def _choose_encoding(a: AVMBytesEncoding, b: AVMBytesEncoding) -> AVMBytesEncoding:
    if a == b:
        # preserve encoding if both equal
        return a
    # exclude utf8 from known choices, we don't preserve that encoding choice unless
    # they're both utf8 strings, which is covered by the first check
    known_binary_choices = {a, b} - {AVMBytesEncoding.utf8, AVMBytesEncoding.unknown}
    if not known_binary_choices:
        return AVMBytesEncoding.unknown

    # pick the most compact encoding of the known binary encodings
    if AVMBytesEncoding.base64 in known_binary_choices:
        return AVMBytesEncoding.base64
    if AVMBytesEncoding.base32 in known_binary_choices:
        return AVMBytesEncoding.base32
    return AVMBytesEncoding.base16

def _decode_address(address: str) -> bytes:
    # Pad address so it's a valid b32 string
    padded_address = address + (6 * "=")
    address_bytes = base64.b32decode(padded_address)
    public_key_hash = address_bytes[: algo_constants.PUBLIC_KEY_HASH_LENGTH]
    return public_key_hash

def _get_byte_constant(
    register_assignments: _RegisterAssignments, byte_arg: models.Value
) -> models.BytesConstant | None:
    if byte_arg in register_assignments:
        byte_arg_defn = register_assignments[byte_arg]  # type: ignore[index]
        match byte_arg_defn.source:
            case models.Intrinsic(op=AVMOp.itob, args=[models.UInt64Constant(value=itob_arg)]):
                return _eval_itob(itob_arg, byte_arg_defn.source_location)
            case models.Intrinsic(op=AVMOp.bzero, args=[models.UInt64Constant(value=bzero_arg)]):
                return _eval_bzero(bzero_arg, byte_arg_defn.source_location)
            case models.Intrinsic(op=AVMOp.global_, immediates=["ZeroAddress"]):
                return models.BytesConstant(
                    value=_decode_address(algo_constants.ZERO_ADDRESS),
                    encoding=AVMBytesEncoding.base32,
                    source_location=byte_arg.source_location,
                )
    elif isinstance(byte_arg, models.Constant):
        if isinstance(byte_arg, models.BytesConstant):
            return byte_arg
        if isinstance(byte_arg, models.BigUIntConstant):
            return models.BytesConstant(
                value=biguint_bytes_eval(byte_arg.value),
                encoding=AVMBytesEncoding.base16,
                source_location=byte_arg.source_location,
            )
        if isinstance(byte_arg, models.AddressConstant):
            return models.BytesConstant(
                value=_decode_address(byte_arg.value),
                encoding=AVMBytesEncoding.base32,
                source_location=byte_arg.source_location,
            )
        if isinstance(byte_arg, models.MethodConstant):
            return models.BytesConstant(
                value=method_selector_hash(byte_arg.value),
                encoding=AVMBytesEncoding.base16,
                source_location=byte_arg.source_location,
            )
    return None

def _eval_itob(arg: int, loc: SourceLocation | None) -> models.BytesConstant:
    return models.BytesConstant(
        value=arg.to_bytes(8, byteorder="big", signed=False),
        encoding=AVMBytesEncoding.base16,
        source_location=loc,
    )

def _eval_bzero(arg: int, loc: SourceLocation | None) -> models.BytesConstant | None:
    if arg <= 64:
        return models.BytesConstant(
            value=b"\x00" * arg,
            encoding=AVMBytesEncoding.base16,
            source_location=loc,
        )
    return None

def _try_simplify_uint64_unary_op(
    intrinsic: models.Intrinsic, arg: models.Value
) -> models.Value | None:
    op_loc = intrinsic.source_location
    x = _get_int_constant(arg)
    if x is not None:
        if intrinsic.op is AVMOp.not_:
            not_x = 0 if x else 1
            return models.UInt64Constant(value=not_x, source_location=op_loc)
        elif intrinsic.op is AVMOp.bitwise_not:
            inverted = x ^ 0xFFFFFFFFFFFFFFFF
            return models.UInt64Constant(value=inverted, source_location=op_loc)
        elif intrinsic.op is AVMOp.sqrt:
            value = math.isqrt(x)
            return models.UInt64Constant(value=value, source_location=op_loc)
        elif intrinsic.op is AVMOp.bitlen:
            return UInt64Constant(value=x.bit_length(), source_location=op_loc)
        elif intrinsic.op is AVMOp.itob:
            return _eval_itob(x, op_loc)
        elif intrinsic.op is AVMOp.bzero:
            return _eval_bzero(x, op_loc)
        else:
            logger.debug(f"Don't know how to simplify {intrinsic.op.code} of {x}")
    return None

def _try_simplify_bytes_unary_op(
    register_assignments: _RegisterAssignments, intrinsic: models.Intrinsic, arg: models.Value
) -> models.Value | None:
    op_loc = intrinsic.source_location
    if intrinsic.op is AVMOp.bsqrt:
        biguint_const, _, _ = _get_biguint_constant(register_assignments, arg)
        if biguint_const is not None:
            value = math.isqrt(biguint_const)
            return models.BigUIntConstant(value=value, source_location=op_loc)
    else:
        byte_const = _get_byte_constant(register_assignments, arg)
        if byte_const is not None:
            if intrinsic.op is AVMOp.bitwise_not_bytes:
                inverted = bytes([x ^ 0xFF for x in byte_const.value])
                return models.BytesConstant(
                    value=inverted, encoding=byte_const.encoding, source_location=op_loc
                )
            elif intrinsic.op is AVMOp.btoi:
                converted = int.from_bytes(byte_const.value, byteorder="big", signed=False)
                return models.UInt64Constant(value=converted, source_location=op_loc)
            elif intrinsic.op is AVMOp.len_:
                length = len(byte_const.value)
                return models.UInt64Constant(value=length, source_location=op_loc)
            elif intrinsic.op is AVMOp.bitlen:
                converted = int.from_bytes(byte_const.value, byteorder="big", signed=False)
                return UInt64Constant(value=converted.bit_length(), source_location=op_loc)
            else:
                logger.debug(f"Don't know how to simplify {intrinsic.op.code} of {byte_const}")
    return None

def _try_simplify_uint64_binary_op(
    intrinsic: models.Intrinsic, a: models.Value, b: models.Value, *, bool_context: bool = False
) -> models.Value | models.Intrinsic | None:
    op = intrinsic.op
    c: models.Value | int | None = None

    if a == b:
        match op:
            case AVMOp.sub:
                c = 0
            case AVMOp.eq | AVMOp.lte | AVMOp.gte:
                c = 1
            case AVMOp.neq | AVMOp.lt | AVMOp.gt:
                c = 0
            case AVMOp.div_floor:
                c = 1
            case AVMOp.bitwise_xor:
                c = 0
            case AVMOp.bitwise_and | AVMOp.bitwise_or:
                c = a
    if c is None:
        a_const = _get_int_constant(a)
        b_const = _get_int_constant(b)

        # a >= 0 <-> 1
        if b_const == 0 and op == AVMOp.gte:  # noqa: SIM114
            c = 1
        # 0 <= b <-> 1
        elif a_const == 0 and op == AVMOp.lte:
            c = 1
        elif a_const == 1 and op == AVMOp.mul:
            c = b
        elif b_const == 1 and op == AVMOp.mul:
            c = a
        elif a_const == 0 and op in (AVMOp.add, AVMOp.or_):
            c = b
        elif b_const == 0 and op in (AVMOp.add, AVMOp.sub, AVMOp.or_):
            c = a
        elif 0 in (a_const, b_const) and op in (AVMOp.mul, AVMOp.and_):
            c = 0
        # 0 != b <-> b
        #   OR
        # 0 < b <-> b
        #   OR
        # 1 <= b <-> b
        elif (bool_context or b.ir_type == IRType.bool) and (
            (a_const == 0 and op in (AVMOp.neq, AVMOp.lt)) or (a_const == 1 and op == AVMOp.lte)
        ):
            c = b
        # a != 0 <-> a
        #   OR
        # a > 0 <-> a
        #   OR
        # a >= 1 <-> a
        elif (bool_context or a.ir_type == IRType.bool) and (
            (b_const == 0 and op in (AVMOp.neq, AVMOp.gt)) or (b_const == 1 and op == AVMOp.gte)
        ):
            c = a
        elif a_const is not None and b_const is not None:
            match op:
                case AVMOp.add:
                    c = a_const + b_const
                case AVMOp.sub:
                    c = a_const - b_const
                case AVMOp.mul:
                    c = a_const * b_const
                case AVMOp.div_floor if b_const != 0:
                    c = a_const // b_const
                case AVMOp.mod if b_const != 0:
                    c = a_const % b_const
                case AVMOp.lt:
                    c = 1 if a_const < b_const else 0
                case AVMOp.lte:
                    c = 1 if a_const <= b_const else 0
                case AVMOp.gt:
                    c = 1 if a_const > b_const else 0
                case AVMOp.gte:
                    c = 1 if a_const >= b_const else 0
                case AVMOp.eq:
                    c = 1 if a_const == b_const else 0
                case AVMOp.neq:
                    c = 1 if a_const != b_const else 0
                case AVMOp.and_:
                    c = int(a_const and b_const)
                case AVMOp.or_:
                    c = int(a_const or b_const)
                case AVMOp.shl:
                    c = (a_const << b_const) % (2**64)
                case AVMOp.shr:
                    c = a_const >> b_const
                case AVMOp.exp:
                    if a_const == 0 and b_const == 0:
                        return None  # would fail at runtime, lets hope this is unreachable 😬
                    c = a_const**b_const
                case AVMOp.bitwise_or:
                    c = a_const | b_const
                case AVMOp.bitwise_and:
                    c = a_const & b_const
                case AVMOp.bitwise_xor:
                    c = a_const ^ b_const
                case _:
                    logger.debug(
                        f"Don't know how to simplify {a_const} {intrinsic.op.code} {b_const}"
                    )
        # 0 == b <-> !b
        elif a_const == 0 and op == AVMOp.eq:
            return attrs.evolve(intrinsic, op=AVMOp.not_, args=[b])
        # a == 0 <-> !a
        elif b_const == 0 and op == AVMOp.eq:
            return attrs.evolve(intrinsic, op=AVMOp.not_, args=[a])
    if not isinstance(c, int):
        return c
    if c < 0:
        # Value cannot be folded as it would result in a negative uint64
        return None
    return models.UInt64Constant(value=c, source_location=intrinsic.source_location)

def _try_simplify_bytes_binary_op(
    register_assignments: _RegisterAssignments,
    intrinsic: models.Intrinsic,
    a: models.Value,
    b: models.Value,
) -> models.Value | None:
    op = intrinsic.op
    op_loc = intrinsic.source_location
    c: models.Value | int | None = None

    if a == b:
        match op:
            case AVMOp.sub_bytes:
                c = 0
            case AVMOp.eq_bytes | AVMOp.eq | AVMOp.lte_bytes | AVMOp.gte_bytes:
                c = 1
            case AVMOp.neq_bytes | AVMOp.neq | AVMOp.lt_bytes | AVMOp.gt_bytes:
                c = 0
            case AVMOp.div_floor_bytes:
                c = 1
            case AVMOp.bitwise_xor_bytes:
                c = 0
            case AVMOp.bitwise_and_bytes | AVMOp.bitwise_or_bytes:
                c = a
    if c is None:
        a_const, a_const_bytes, a_encoding = _get_biguint_constant(register_assignments, a)
        b_const, b_const_bytes, b_encoding = _get_biguint_constant(register_assignments, b)
        if a_const == 1 and op == AVMOp.mul_bytes:
            c = b
        elif b_const == 1 and op == AVMOp.mul_bytes:
            c = a
        elif a_const == 0 and op == AVMOp.add_bytes:
            c = b
        elif b_const == 0 and op in (AVMOp.add_bytes, AVMOp.sub_bytes):
            c = a
        elif 0 in (a_const, b_const) and op == AVMOp.mul_bytes:
            c = 0
        elif a_const is not None and b_const is not None:
            if typing.TYPE_CHECKING:
                assert a_const_bytes is not None
                assert b_const_bytes is not None
                assert a_encoding is not None
                assert b_encoding is not None
            match op:
                case AVMOp.add_bytes:
                    c = a_const + b_const
                case AVMOp.sub_bytes:
                    c = a_const - b_const
                case AVMOp.mul_bytes:
                    c = a_const * b_const
                case AVMOp.div_floor_bytes:
                    c = a_const // b_const
                case AVMOp.mod_bytes if b_const != 0:
                    c = a_const % b_const
                case AVMOp.lt_bytes:
                    c = 1 if a_const < b_const else 0
                case AVMOp.lte_bytes:
                    c = 1 if a_const <= b_const else 0
                case AVMOp.gt_bytes:
                    c = 1 if a_const > b_const else 0
                case AVMOp.gte_bytes:
                    c = 1 if a_const >= b_const else 0
                case AVMOp.eq_bytes:
                    c = 1 if a_const == b_const else 0
                case AVMOp.neq_bytes:
                    c = 1 if a_const != b_const else 0
                case AVMOp.eq:
                    c = 1 if a_const_bytes == b_const_bytes else 0
                case AVMOp.neq:
                    c = 1 if a_const_bytes != b_const_bytes else 0
                case AVMOp.bitwise_or_bytes:
                    return models.BytesConstant(
                        value=_byte_wise(operator.or_, a_const_bytes, b_const_bytes),
                        encoding=_choose_encoding(a_encoding, b_encoding),
                        source_location=op_loc,
                    )
                case AVMOp.bitwise_and_bytes:
                    return models.BytesConstant(
                        value=_byte_wise(operator.and_, a_const_bytes, b_const_bytes),
                        encoding=_choose_encoding(a_encoding, b_encoding),
                        source_location=op_loc,
                    )
                case AVMOp.bitwise_xor_bytes:
                    return models.BytesConstant(
                        value=_byte_wise(operator.xor, a_const_bytes, b_const_bytes),
                        encoding=_choose_encoding(a_encoding, b_encoding),
                        source_location=op_loc,
                    )
                case _:
                    logger.debug(
                        f"Don't know how to simplify {a_const} {intrinsic.op.code} {b_const}"
                    )
    if not isinstance(c, int):
        return c
    if c < 0:
        # don't fold to a negative
        return None
    # could look at StackType of op_signature.returns, but some are StackType.any
    if op in (
        AVMOp.eq_bytes,
        AVMOp.eq,
        AVMOp.neq_bytes,
        AVMOp.neq,
        AVMOp.lt_bytes,
        AVMOp.lte_bytes,
        AVMOp.gt_bytes,
        AVMOp.gte_bytes,
    ):
        return models.UInt64Constant(value=c, source_location=intrinsic.source_location)
    else:
        return models.BigUIntConstant(value=c, source_location=intrinsic.source_location)


import typing
from collections import defaultdict

import attrs

from puya.context import CompileContext
from puya.ir import models
from puya.ir.avm_ops import AVMOp

_ARRAY_FIELDS: typing.Final = frozenset(
    (
        "ApplicationArgs",
        "Accounts",
        "Assets",
        "Applications",
        "ApprovalProgramPages",
        "ClearStateProgramPages",
    )
)
_ZERO_DEFAULTS: typing.Final = frozenset(
    (
        "ExtraProgramPages",
        "GlobalNumUint",
        "GlobalNumByteSlice",
        "LocalNumUint",
        "LocalNumByteSlice",
    )
)

def elide_itxn_field_calls(_context: CompileContext, subroutine: models.Subroutine) -> bool:
    modified = False
    for block in subroutine.body:
        if _elide_within_block(block):
            modified = True
    return modified

def _elide_within_block(block: models.BasicBlock) -> bool:
    groups = list[_ITxnGroup]()
    current_group: _ITxnGroup | None = None
    for op_idx, op in enumerate(block.ops):
        if isinstance(op, models.Intrinsic):
            match op.op:
                case AVMOp.itxn_begin | AVMOp.itxn_next:
                    groups.append(current_group := _ITxnGroup(has_start=True))
                case AVMOp.itxn_submit:
                    current_group = None
                case AVMOp.itxn_field:
                    (field_im,) = op.immediates
                    assert isinstance(field_im, str)
                    if field_im not in _ARRAY_FIELDS:
                        if current_group is None:
                            groups.append(current_group := _ITxnGroup(has_start=False))
                        current_group.sets[field_im].append(op_idx)
    if not groups:
        return False
    remove_indexes = set[int]()
    for group in groups:
        for field_im, indexes in group.sets.items():
            final_idx = indexes.pop()
            remove_indexes.update(indexes)
            if group.has_start and field_im in _ZERO_DEFAULTS:
                match block.ops[final_idx]:
                    case models.Intrinsic(args=[models.UInt64Constant(value=0)]):
                        remove_indexes.add(final_idx)
    if not remove_indexes:
        return False
    block.ops = [op for idx, op in enumerate(block.ops) if idx not in remove_indexes]
    return True

@attrs.define
class _ITxnGroup:
    has_start: bool
    sets: defaultdict[str, list[int]] = attrs.field(factory=lambda: defaultdict(list))


from collections.abc import Callable, Collection, Iterable

import attrs

from puya import log
from puya.context import ArtifactCompileContext
from puya.ir import models
from puya.ir.optimize.assignments import copy_propagation
from puya.ir.optimize.collapse_blocks import remove_empty_blocks, remove_linear_jump
from puya.ir.optimize.compiled_reference import replace_compiled_references
from puya.ir.optimize.constant_propagation import constant_replacer
from puya.ir.optimize.context import IROptimizationContext
from puya.ir.optimize.control_op_simplification import simplify_control_ops
from puya.ir.optimize.dead_code_elimination import (
    remove_unreachable_blocks,
    remove_unused_subroutines,
    remove_unused_variables,
)
from puya.ir.optimize.inlining import analyse_subroutines_for_inlining, perform_subroutine_inlining
from puya.ir.optimize.inner_txn import inner_txn_field_replacer
from puya.ir.optimize.intrinsic_simplification import intrinsic_simplifier
from puya.ir.optimize.itxn_field_elision import elide_itxn_field_calls
from puya.ir.optimize.repeated_code_elimination import repeated_expression_elimination
from puya.ir.to_text_visitor import render_program
from puya.utils import attrs_extend

MAX_PASSES = 100
SubroutineOptimizerCallable = Callable[[IROptimizationContext, models.Subroutine], bool]
logger = log.get_logger(__name__)

@attrs.define(kw_only=True)
class SubroutineOptimization:
    id: str
    desc: str
    _optimize: SubroutineOptimizerCallable
    loop: bool

    @classmethod
    def from_function(
        cls, func: SubroutineOptimizerCallable, *, loop: bool = False
    ) -> "SubroutineOptimization":
        func_name = func.__name__
        func_desc = func_name.replace("_", " ").title().strip()
        return SubroutineOptimization(id=func_name, desc=func_desc, optimize=func, loop=loop)

    def optimize(self, context: IROptimizationContext, ir: models.Subroutine) -> bool:
        did_modify = self._optimize(context, ir)
        if did_modify:
            while self.loop and self._optimize(context, ir):
                pass
        return did_modify

def get_subroutine_optimizations(optimization_level: int) -> Iterable[SubroutineOptimization]:
    if optimization_level:
        return [
            SubroutineOptimization.from_function(perform_subroutine_inlining),
            SubroutineOptimization.from_function(_split_parallel_copies),
            SubroutineOptimization.from_function(constant_replacer),
            SubroutineOptimization.from_function(copy_propagation),
            SubroutineOptimization.from_function(intrinsic_simplifier),
            SubroutineOptimization.from_function(elide_itxn_field_calls),
            SubroutineOptimization.from_function(remove_unused_variables),
            SubroutineOptimization.from_function(inner_txn_field_replacer),
            SubroutineOptimization.from_function(replace_compiled_references),
            SubroutineOptimization.from_function(simplify_control_ops, loop=True),
            SubroutineOptimization.from_function(remove_linear_jump),
            SubroutineOptimization.from_function(remove_empty_blocks),
            SubroutineOptimization.from_function(remove_unreachable_blocks),
            SubroutineOptimization.from_function(repeated_expression_elimination),
        ]
    else:
        return [
            SubroutineOptimization.from_function(perform_subroutine_inlining),
            SubroutineOptimization.from_function(_split_parallel_copies),
            SubroutineOptimization.from_function(constant_replacer),
            SubroutineOptimization.from_function(remove_unused_variables),
            SubroutineOptimization.from_function(inner_txn_field_replacer),
            SubroutineOptimization.from_function(replace_compiled_references),
        ]

def _split_parallel_copies(_ctx: ArtifactCompileContext, sub: models.Subroutine) -> bool:
    # not an optimisation, but simplifies other optimisation code
    any_modified = False
    for block in sub.body:
        ops = list[models.Op]()
        modified = False
        for op in block.ops.copy():
            if isinstance(op, models.Assignment) and isinstance(op.source, models.ValueTuple):
                for dst, src in zip(op.targets, op.source.values, strict=True):
                    modified = True
                    ops.append(
                        models.Assignment(
                            targets=[dst],
                            source=src,
                            source_location=op.source_location,
                        )
                    )
            else:
                ops.append(op)
        if modified:
            any_modified = True
            block.ops = ops
    return any_modified

def optimize_program_ir(
    context: ArtifactCompileContext,
    program: models.Program,
    *,
    routable_method_ids: Collection[str] | None,
) -> None:
    pipeline = get_subroutine_optimizations(context.options.optimization_level)
    opt_context = attrs_extend(IROptimizationContext, context, expand_all_bytes=False)
    for pass_num in range(1, MAX_PASSES + 1):
        program_modified = False
        logger.debug(f"Begin optimization pass {pass_num}/{MAX_PASSES}")
        analyse_subroutines_for_inlining(opt_context, program, routable_method_ids)
        for subroutine in program.all_subroutines:
            logger.debug(f"Optimizing subroutine {subroutine.id}")
            for optimizer in pipeline:
                logger.debug(f"Optimizer: {optimizer.desc}")
                if optimizer.optimize(opt_context, subroutine):
                    program_modified = True
            subroutine.validate_with_ssa()
        if remove_unused_subroutines(program):
            program_modified = True
        if not program_modified:
            logger.debug(f"No optimizations performed in pass {pass_num}, ending loop")
            break
        if context.options.output_optimization_ir:
            render_program(context, program, qualifier="ssa.opt")


import functools
import operator
import typing
from collections.abc import Sequence

import attrs

from puya import log
from puya.context import CompileContext
from puya.ir import models
from puya.ir.avm_ops import AVMOp
from puya.ir.optimize.assignments import copy_propagation
from puya.ir.optimize.dead_code_elimination import PURE_AVM_OPS
from puya.ir.visitor import NoOpIRVisitor

logger = log.get_logger(__name__)

@attrs.frozen
class IntrinsicData:
    op: AVMOp
    immediates: tuple[str | int, ...]
    args: tuple[models.Value, ...]

    @classmethod
    def from_op(cls, op: models.Intrinsic) -> typing.Self:
        return cls(
            op=op.op,
            immediates=tuple(op.immediates),
            args=tuple(op.args),
        )

def repeated_expression_elimination(
    context: CompileContext, subroutine: models.Subroutine
) -> bool:
    any_modified = False
    dom = compute_dominators(subroutine)
    modified = True
    while modified:
        modified = False
        block_asserted = dict[models.BasicBlock, set[models.Value]]()
        block_const_intrinsics = dict[
            models.BasicBlock, dict[IntrinsicData, Sequence[models.Register]]
        ]()
        for block in subroutine.body:
            visitor = RCEVisitor(block)
            for op in block.ops.copy():
                modified = bool(op.accept(visitor)) or modified
            block_asserted[block] = visitor.asserted
            block_const_intrinsics[block] = visitor.const_intrinsics

        for block in subroutine.body:
            dominators = dom[block]
            if dominators:
                visitor = RCEVisitor(
                    block,
                    # if there is a single dominator then reduce will return the original
                    # collection of that dominator, so copy to ensure original collections are
                    # not modified
                    const_intrinsics=functools.reduce(
                        operator.or_, (block_const_intrinsics[b] for b in dominators)
                    ).copy(),
                    asserted=functools.reduce(
                        operator.or_, (block_asserted[b] for b in dominators)
                    ).copy(),
                )
                for op in block.ops.copy():
                    modified = bool(op.accept(visitor)) or modified
        if modified:
            any_modified = True
            copy_propagation(context, subroutine)
    return any_modified

def compute_dominators(
    subroutine: models.Subroutine,
) -> dict[models.BasicBlock, list[models.BasicBlock]]:
    all_blocks = set(subroutine.body)
    dom = {b: all_blocks if b.predecessors else {b} for b in subroutine.body}
    changes = True
    while changes:
        changes = False
        for block in reversed(subroutine.body):
            if block.predecessors:
                pred_dom = functools.reduce(set.intersection, (dom[p] for p in block.predecessors))
                new = pred_dom | {block}
                if new != dom[block]:
                    dom[block] = new
                    changes = True
    return {
        b: sorted(dom_set - {b}, key=lambda a: typing.cast(int, a.id))
        for b, dom_set in dom.items()
    }

@attrs.define
class RCEVisitor(NoOpIRVisitor[bool]):
    block: models.BasicBlock
    const_intrinsics: dict[IntrinsicData, Sequence[models.Register]] = attrs.field(factory=dict)
    asserted: set[models.Value] = attrs.field(factory=set)

    _assignment: models.Assignment | None = None

    def visit_assignment(self, ass: models.Assignment) -> bool | None:
        self._assignment = ass
        remove = ass.source.accept(self)
        self._assignment = None
        return remove

    def visit_intrinsic_op(self, intrinsic: models.Intrinsic) -> bool:
        modified = False
        if (ass := self._assignment) is not None:
            # only consider ops with stack args because they're much more likely to
            # produce extra stack manipulations
            if intrinsic.args and intrinsic.op.code in PURE_AVM_OPS:
                key = IntrinsicData.from_op(intrinsic)
                try:
                    existing = self.const_intrinsics[key]
                except KeyError:
                    self.const_intrinsics[key] = ass.targets
                else:
                    logger.debug(
                        f"Replacing redundant declaration {ass}"
                        f" with copy of existing registers {existing}"
                    )
                    modified = True
                    if len(existing) == 1:
                        ass.source = existing[0]
                    else:
                        current_idx = self.block.ops.index(ass)
                        self.block.ops[current_idx : current_idx + 1] = [
                            models.Assignment(
                                targets=[dst], source=src, source_location=ass.source_location
                            )
                            for dst, src in zip(ass.targets, existing, strict=True)
                        ]
        elif intrinsic.op.code == "assert":
            (assert_arg,) = intrinsic.args
            if assert_arg in self.asserted:
                logger.debug(f"Removing redundant assert of {assert_arg}")
                modified = True
                self.block.ops.remove(intrinsic)
            else:
                self.asserted.add(assert_arg)
        return modified


from collections.abc import Set

import attrs

from puya.ir import models
from puya.ir.visitor import IRTraverser

@attrs.frozen
class RegisterReadCollector(IRTraverser):
    _used_registers: list[models.Register] = attrs.field(factory=list, init=False)

    @property
    def used_registers(self) -> Set[models.Register]:
        return dict.fromkeys(self._used_registers).keys()

    def visit_register(self, reg: models.Register) -> None:
        self._used_registers.append(reg)

    def visit_assignment(self, ass: models.Assignment) -> None:
        ass.source.accept(self)

    def visit_phi(self, phi: models.Phi) -> None:
        for arg in phi.args:
            self.visit_phi_argument(arg)


from collections.abc import Sequence

import attrs

from puya import log
from puya.errors import InternalError
from puya.ir import models as ir
from puya.ir.types_ import IRType
from puya.ir.visitor_mutator import IRMutator
from puya.parse import SourceLocation

logger = log.get_logger(__name__)

class BraunSSA:
    """
    Constructs CFG+SSA IR directly from an AST like-structure, without having to go
    through an intermediate non-SSA IR and then compute dominance frontiers etc.

    The primary reference is https://pp.ipd.kit.edu/uploads/publikationen/braun13cc.pdf,
    but also of interest is https://pp.ipd.kit.edu/uploads/publikationen/buchwald16cc.pdf,
    in which the algorithm is converted into a functional programming form,
    and formally verified.
    """

    def __init__(
        self,
        all_blocks: Sequence[ir.BasicBlock],
        live_variables: Sequence[ir.Register],
        active_block: ir.BasicBlock,
    ) -> None:
        self._all_blocks = all_blocks  # note: this is a shared reference with BlocksBuilder
        self._sealed_blocks = set[ir.BasicBlock]()
        self._current_def = dict[str, dict[ir.BasicBlock, ir.Register]]()
        self._incomplete_phis = dict[ir.BasicBlock, list[ir.Phi]]()
        self._variable_versions = dict[str, int]()
        # initialize any live variables at the start of the subroutines, i.e. parameters
        for parameter in live_variables:
            self.write_variable(parameter.name, active_block, parameter)
            self._variable_versions[parameter.name] = 1

    def has_version(self, variable: str) -> bool:
        return variable in self._variable_versions

    def write_variable(self, variable: str, block: ir.BasicBlock, value: ir.Register) -> None:
        self._current_def.setdefault(variable, {})[block] = value

    def has_write(self, variable: str, block: ir.BasicBlock) -> bool:
        try:
            self._current_def[variable][block]
        except KeyError:
            return False
        else:
            return True

    def read_variable(self, variable: str, ir_type: IRType, block: ir.BasicBlock) -> ir.Register:
        try:
            result = self._current_def[variable][block]
        except KeyError:
            result = self._read_variable_recursive(variable, ir_type, block)
        return result

    def _read_variable_recursive(
        self, variable: str, ir_type: IRType, block: ir.BasicBlock
    ) -> ir.Register:
        value: ir.Register
        if not self.is_sealed(block):
            logger.debug(
                f"Looking for '{variable}' in an unsealed block creating an incomplete "
                f"Phi: {block}"
            )
            # incomplete CFG
            phi = self._new_phi(block, variable, ir_type)
            self._incomplete_phis.setdefault(block, []).append(phi)
            value = phi.register
        elif len(block.predecessors) == 1:
            value = self.read_variable(variable, ir_type, block.predecessors[0])
        else:
            # break any potential cycles with empty Phi
            phi = self._new_phi(block, variable, ir_type)
            self.write_variable(variable, block, phi.register)
            value = self._add_phi_operands(phi, block)
        self.write_variable(variable, block, value)
        return value

    def new_register(
        self, name: str, ir_type: IRType, location: SourceLocation | None
    ) -> ir.Register:
        version = self._variable_versions.get(name, 0)
        self._variable_versions[name] = version + 1
        return ir.Register(source_location=location, name=name, version=version, ir_type=ir_type)

    def is_sealed(self, block: ir.BasicBlock) -> bool:
        return block in self._sealed_blocks

    def verify_complete(self) -> None:
        unsealed = [b for b in self._all_blocks if not self.is_sealed(b)]
        if unsealed:
            raise InternalError(
                f"{len(unsealed)} block/s were not sealed: " + ", ".join(map(str, unsealed))
            )
        incomplete_phis = [p for block_phis in self._incomplete_phis.values() for p in block_phis]
        if incomplete_phis:
            # if we get here, there is a bug in this algorithm itself
            raise InternalError(
                f"{len(incomplete_phis)} phi node/s are incomplete: "
                + ", ".join(map(str, incomplete_phis))
            )

    def _new_phi(self, block: ir.BasicBlock, variable: str, ir_type: IRType) -> ir.Phi:
        reg = self.new_register(variable, ir_type, location=None)
        phi = ir.Phi(register=reg)
        block.phis.append(phi)

        logger.debug(
            f"Created Phi assignment: {phi} while trying to resolve '{variable}' in {block}"
        )
        return phi

    def _add_phi_operands(self, phi: ir.Phi, block: ir.BasicBlock) -> ir.Register:
        variable = phi.register.name
        for block_pred in block.predecessors:
            pred_variable = self.read_variable(variable, phi.ir_type, block_pred)
            phi.args.append(
                ir.PhiArgument(
                    value=pred_variable,
                    through=block_pred,
                )
            )
            logger.debug(f"Added {pred_variable} to Phi node: {phi} in {block_pred}")
            attrs.validate(phi)
        trivial_replacements = TrivialPhiRemover.try_remove(phi, self._all_blocks)
        if not trivial_replacements:
            return phi.register
        _, result = trivial_replacements[0]

        variable_def = self._current_def[variable]
        for removed_phi, replacement_memory in trivial_replacements:
            for unsealed_block, incomplete_phis in self._incomplete_phis.items():
                if removed_phi in incomplete_phis:
                    raise InternalError(
                        f"removed an incomplete phi {removed_phi} from block {unsealed_block}!!"
                    )
            if replacement_memory.name != variable:
                raise InternalError(
                    "Tangled phi web created during SSA construction",
                    replacement_memory.source_location,
                )
            if removed_phi.register == result:
                result = replacement_memory
            # ensure replacement is also recorded in variable cache
            replacements = 0
            for block_def, arg in variable_def.items():
                if arg == removed_phi.register:
                    variable_def[block_def] = replacement_memory
                    replacements += 1
            logger.debug(
                f"Replaced trivial Phi node: {removed_phi} ({removed_phi.register})"
                f" with {replacement_memory} in current definition for {replacements} blocks"
            )
        return result

    def seal_block(self, block: ir.BasicBlock) -> None:
        logger.debug(f"Sealing {block}")
        if self.is_sealed(block):
            raise InternalError(f"Cannot seal a block that was already sealed: {block}")

        for predecessor in block.predecessors:
            if not predecessor.terminated:
                raise InternalError(f"All predecessors must be terminated before sealing: {block}")

        phi_list = self._incomplete_phis.pop(block, [])
        for phi in phi_list:
            self._add_phi_operands(phi, block)

        self._sealed_blocks.add(block)

@attrs.define(kw_only=True)
class TrivialPhiRemover(IRMutator):
    _find: ir.Phi
    _replacement: ir.Register
    collected: list[ir.Phi] = attrs.field(factory=list)
    replaced: int = 0

    @classmethod
    def try_remove(
        cls, phi: ir.Phi, blocks: Sequence[ir.BasicBlock]
    ) -> list[tuple[ir.Phi, ir.Register]]:
        try:
            replacement, *other = {a.value for a in phi.non_self_args}
        except ValueError:
            # # per https://pp.ipd.kit.edu/uploads/publikationen/buchwald16cc.pdf, section 4.1
            logger.warning(f"Variable {phi.register.name} potentially used before assignment")
            return []

        result = []
        # if other, then phi merges >1 value: not trivial
        if not other:
            result.append((phi, replacement))
            # replace users of phi with replacement
            logger.debug(f"Replacing trivial Phi node: {phi} ({phi.register}) with {replacement}")
            # this approach is not very efficient, but is simpler than tracking all usages
            # and good enough for now
            replacer = cls(find=phi, replacement=replacement)
            for block in blocks:
                replacer.visit_block(block)
            # recursively check/replace all phi users, which may now be trivial
            for phi_user in replacer.collected:
                result.extend(cls.try_remove(phi_user, blocks))
        return result

    def visit_register(self, reg: ir.Register) -> ir.Register:
        if reg != self._find.register:
            return reg
        self.replaced += 1
        return self._replacement

    def visit_phi(self, phi: ir.Phi) -> ir.Phi | None:
        if phi is self._find:
            logger.debug(f"Deleting Phi assignment: {phi}")
            return None
        prior_replace_count = self.replaced
        result = super().visit_phi(phi)
        assert result is phi, "phi instance changed"
        if self.replaced > prior_replace_count:
            self.collected.append(result)
        return result


import contextlib
import typing
from collections.abc import Iterator, Sequence

from puya import log
from puya.context import ArtifactCompileContext
from puya.ir import models
from puya.ir.types_ import IRType
from puya.ir.utils import format_bytes, format_error_comment
from puya.ir.visitor import IRVisitor
from puya.utils import make_path_relative_to_cwd

logger = log.get_logger(__name__)

class ToTextVisitor(IRVisitor[str]):
    @typing.override
    def visit_assignment(self, op: models.Assignment) -> str:
        targets = ", ".join(f"{r.accept(self)}: {r.ir_type.name}" for r in op.targets)
        if len(op.targets) > 1:
            targets = f"({targets})"
        source = op.source.accept(self)
        return f"let {targets} = {source}"

    @typing.override
    def visit_register(self, op: models.Register) -> str:
        return op.local_id

    @typing.override
    def visit_undefined(self, op: models.Undefined) -> str:
        return "undefined"

    @typing.override
    def visit_uint64_constant(self, op: models.UInt64Constant) -> str:
        return f"{op.value}u" if not op.teal_alias else op.teal_alias

    @typing.override
    def visit_biguint_constant(self, op: models.BigUIntConstant) -> str:
        return f"{op.value}b"

    @typing.override
    def visit_bytes_constant(self, op: models.BytesConstant) -> str:
        return format_bytes(op.value, op.encoding)

    @typing.override
    def visit_address_constant(self, op: models.AddressConstant) -> str:
        return f"addr {op.value}"

    @typing.override
    def visit_method_constant(self, op: models.MethodConstant) -> str:
        return f'method "{op.value}"'

    @typing.override
    def visit_itxn_constant(self, op: models.ITxnConstant) -> str:
        return f"{op.ir_type.name}({op.value})"

    @typing.override
    def visit_compiled_contract_reference(self, const: models.CompiledContractReference) -> str:
        return (
            ", ".join(
                (
                    f"compiled_contract({const.artifact!r}",
                    f"field={const.field.name}",
                    f"program_page={const.program_page}",
                    *(
                        f"{var}={val.accept(self)}"
                        for var, val in const.template_variables.items()
                    ),
                )
            )
            + ")"
        )

    @typing.override
    def visit_compiled_logicsig_reference(self, const: models.CompiledLogicSigReference) -> str:
        return (
            ", ".join(
                (
                    f"compiled_logicsig({const.artifact!r}",
                    *(
                        f"{var}={val.accept(self)}"
                        for var, val in const.template_variables.items()
                    ),
                )
            )
            + ")"
        )

    @typing.override
    def visit_intrinsic_op(self, intrinsic: models.Intrinsic) -> str:
        callee = intrinsic.op.code
        immediates = list(map(str, intrinsic.immediates))
        if immediates:
            callee = f"({' '.join((callee, *immediates))})"
        args = " ".join(arg.accept(self) for arg in intrinsic.args)
        if args:
            callee = f"({callee} {args})"
        if intrinsic.error_message:
            callee += f" // {format_error_comment(intrinsic.op.code, intrinsic.error_message)}"
        return callee

    @typing.override
    def visit_inner_transaction_field(self, field: models.InnerTransactionField) -> str:
        group = field.group_index.accept(self)
        array_access = f"[{field.array_index.accept(self)}]" if field.array_index else ""
        return f"itxn[{group}].{field.field}{array_access}"

    @typing.override
    def visit_invoke_subroutine(self, op: models.InvokeSubroutine) -> str:
        args = ", ".join(a.accept(self) for a in op.args)
        return f"{op.target.id}({args})"

    @typing.override
    def visit_conditional_branch(self, op: models.ConditionalBranch) -> str:
        return f"goto {op.condition.accept(self)} ? {op.non_zero} : {op.zero}"

    @typing.override
    def visit_goto(self, op: models.Goto) -> str:
        return f"goto {op.target}"

    @typing.override
    def visit_goto_nth(self, op: models.GotoNth) -> str:
        blocks = ", ".join(map(str, op.blocks))
        return f"goto_nth [{blocks}][{op.value.accept(self)}] else goto {op.default}"

    @typing.override
    def visit_switch(self, op: models.Switch) -> str:
        cases = {k.accept(self): str(b) for k, b in op.cases.items()}
        cases["*"] = str(op.default)
        map_ = ", ".join(f"{k} => {v}" for k, v in cases.items())
        return f"switch {op.value.accept(self)} {{{map_}}}"

    @typing.override
    def visit_subroutine_return(self, op: models.SubroutineReturn) -> str:
        results = " ".join(r.accept(self) for r in op.result)
        return f"return {results}"

    @typing.override
    def visit_template_var(self, deploy_var: models.TemplateVar) -> str:
        return f"TemplateVar[{deploy_var.ir_type.name}]({deploy_var.name})"

    @typing.override
    def visit_program_exit(self, op: models.ProgramExit) -> str:
        return f"exit {op.result.accept(self)}"

    @typing.override
    def visit_fail(self, op: models.Fail) -> str:
        if op.error_message:
            return f"fail // {op.error_message}"
        return "fail"

    @typing.override
    def visit_phi(self, op: models.Phi) -> str:
        r = op.register
        target = f"{r.accept(self)}: {r.ir_type.name}"
        if op.args:
            args = ", ".join(a.accept(self) for a in op.args)
            source = f"φ({args})"
        else:
            source = "undefined"
        return f"let {target} = {source}"

    @typing.override
    def visit_phi_argument(self, op: models.PhiArgument) -> str:
        return f"{op.value.accept(self)} <- {op.through}"

    @typing.override
    def visit_value_tuple(self, tup: models.ValueTuple) -> str:
        return "(" + ", ".join(val.accept(self) for val in tup.values) + ")"

class TextEmitter:
    def __init__(self) -> None:
        self.lines = list[str]()
        self._indent = 0

    def append(self, line: str) -> None:
        self.lines.append((self._indent * " ") + line)

    @contextlib.contextmanager
    def indent(self, spaces: int = 4) -> Iterator[None]:
        self._indent += spaces
        try:
            yield
        finally:
            self._indent -= spaces

def _render_block_name(block: models.BasicBlock) -> str:
    result = f"{block}: // "
    if block.comment:
        result += f"{block.comment}_"
    result += f"L{block.source_location.line}"
    return result

def _render_body(emitter: TextEmitter, blocks: Sequence[models.BasicBlock]) -> None:
    renderer = ToTextVisitor()
    for block in blocks:
        assert block.terminated
        emitter.append(_render_block_name(block))
        with emitter.indent():
            for op in block.all_ops:
                emitter.append(op.accept(renderer))

def render_program(
    context: ArtifactCompileContext, program: models.Program, *, qualifier: str
) -> None:
    path = context.build_output_path(program.kind, qualifier, "ir")
    if path is None:
        return
    emitter = TextEmitter()
    emitter.append(f"main {program.main.id}:")
    with emitter.indent():
        _render_body(emitter, program.main.body)
    for sub in program.subroutines:
        emitter.append("")
        args = ", ".join(f"{r.name}: {r.ir_type.name}" for r in sub.parameters)
        match sub.returns:
            case []:
                returns = "void"
            case [IRType(name=returns)]:
                pass
            case _ as ir_types:
                returns = f"<{', '.join(t.name for t in ir_types)}>"
        emitter.append(f"subroutine {sub.id}({args}) -> {returns}:")
        with emitter.indent():
            _render_body(emitter, sub.body)
    path.write_text("\n".join(emitter.lines), encoding="utf-8")
    logger.debug(f"Output IR to {make_path_relative_to_cwd(path)}")


import enum
import typing
from collections.abc import Sequence

from puya.avm import AVMType
from puya.awst import (
    nodes as awst_nodes,
    wtypes,
)
from puya.awst.nodes import BytesEncoding
from puya.errors import CodeError, InternalError
from puya.ir.avm_ops_models import StackType
from puya.parse import SourceLocation

@enum.unique
class AVMBytesEncoding(enum.StrEnum):
    unknown = enum.auto()
    base16 = enum.auto()
    base32 = enum.auto()
    base64 = enum.auto()
    utf8 = enum.auto()

@enum.unique
class IRType(enum.StrEnum):
    bytes = enum.auto()
    uint64 = enum.auto()
    bool = enum.auto()
    biguint = enum.auto()
    itxn_group_idx = enum.auto()  # the group index of the result
    itxn_field_set = enum.auto()  # a collection of fields for a pending itxn submit

    @property
    def avm_type(self) -> typing.Literal[AVMType.uint64, AVMType.bytes]:
        maybe_result = self.maybe_avm_type
        if not isinstance(maybe_result, AVMType):
            raise InternalError(f"{maybe_result} cannot be mapped to AVM stack type")
        return maybe_result

    @property
    def maybe_avm_type(self) -> typing.Literal[AVMType.uint64, AVMType.bytes] | str:
        match self:
            case IRType.uint64 | IRType.bool:
                return AVMType.uint64
            case IRType.bytes | IRType.biguint:
                return AVMType.bytes
            case IRType.itxn_group_idx | IRType.itxn_field_set:
                return self.name
            case _:
                typing.assert_never(self)

def bytes_enc_to_avm_bytes_enc(bytes_encoding: BytesEncoding) -> AVMBytesEncoding:
    try:
        return AVMBytesEncoding(bytes_encoding.value)
    except ValueError as ex:
        raise InternalError(f"Unhandled BytesEncoding: {bytes_encoding}") from ex

def wtype_to_ir_type(
    expr_or_wtype: wtypes.WType | awst_nodes.Expression,
    source_location: SourceLocation | None = None,
) -> IRType:
    if isinstance(expr_or_wtype, awst_nodes.Expression):
        return wtype_to_ir_type(
            expr_or_wtype.wtype, source_location=source_location or expr_or_wtype.source_location
        )
    else:
        wtype = expr_or_wtype
    match wtype:
        case wtypes.bool_wtype:
            return IRType.bool
        case wtypes.biguint_wtype:
            return IRType.biguint
        case wtypes.WInnerTransaction():
            return IRType.itxn_group_idx
        case wtypes.WInnerTransactionFields():
            return IRType.itxn_field_set
        case wtypes.void_wtype:
            raise InternalError("can't translate void wtype to irtype", source_location)
        # case wtypes.state_key:
        #     return IRType.state_key  # TODO
        # case wtypes.box_key:
        #     return IRType.box_key  # TODO
    match wtype.scalar_type:
        case AVMType.uint64:
            return IRType.uint64
        case AVMType.bytes:
            return IRType.bytes
        case None:
            raise CodeError(
                f"unsupported nested/compound wtype encountered: {wtype}", source_location
            )
        case _:
            typing.assert_never(wtype.scalar_type)

def get_wtype_arity(wtype: wtypes.WType) -> int:
    """Returns the number of values this wtype represents on the stack"""
    if isinstance(wtype, wtypes.WTuple):
        return sum_wtypes_arity(wtype.types)
    else:
        return 1

def sum_wtypes_arity(types: Sequence[wtypes.WType]) -> int:
    return sum(map(get_wtype_arity, types))

def wtype_to_ir_types(
    expr_or_wtype: wtypes.WType | awst_nodes.Expression,
    source_location: SourceLocation | None = None,
) -> list[IRType]:
    if isinstance(expr_or_wtype, awst_nodes.Expression):
        wtype = expr_or_wtype.wtype
    else:
        wtype = expr_or_wtype
    if wtype == wtypes.void_wtype:
        return []
    elif isinstance(wtype, wtypes.WTuple):
        return [
            ir_type
            for wtype in wtype.types
            for ir_type in wtype_to_ir_types(wtype, source_location)
        ]
    else:
        return [wtype_to_ir_type(wtype, source_location)]

def stack_type_to_avm_type(stack_type: StackType) -> AVMType:
    match stack_type:
        case StackType.uint64 | StackType.bool | StackType.asset | StackType.application:
            return AVMType.uint64
        case (
            StackType.bytes
            | StackType.bigint
            | StackType.box_name
            | StackType.address
            | StackType.state_key
        ):
            return AVMType.bytes
        case StackType.any | StackType.address_or_index:
            return AVMType.any

def stack_type_to_ir_type(stack_type: StackType) -> IRType | None:
    match stack_type:
        case StackType.bool:
            return IRType.bool
        case StackType.bigint:
            return IRType.biguint
        case StackType.uint64 | StackType.asset | StackType.application:
            return IRType.uint64
        case StackType.bytes | StackType.box_name | StackType.address | StackType.state_key:
            return IRType.bytes
        case StackType.any | StackType.address_or_index:
            return None
        case _:
            typing.assert_never(stack_type)


import base64
from collections.abc import Sequence

from puya.awst import (
    nodes as awst_nodes,
    wtypes,
)
from puya.ir.types_ import AVMBytesEncoding

def escape_utf8_string(s: str) -> str:
    """Escape a UTF-8 string for use in TEAL assembly.

    Args:
        s: A UTF-8 string to escape.

    Returns:
        An escaped version of the input string. This version will be surrounded in double quotes,
        all special characters (such as \\n) will be escaped with additional backslashes, and all
        Unicode characters beyond the latin-1 encoding will be encoded in hex escapes (e.g. \\xf0).
    """
    # The point of this conversion is to escape all special characters and turn all Unicode
    # characters into hex-escaped characters in the input string.
    #
    # The first step breaks up large Unicode characters into multiple UTF-8 hex characters:
    #     s_1 = s.encode("utf-8").decode("latin-1")
    # e.g. "\n 😀" => "\n ð\x9f\x98\x80"
    #
    # The next step escapes all special characters:
    #     s_1.encode("unicode-escape").decode("latin-1")
    # e.g. "\n ð\x9f\x98\x80" => "\\n \\xf0\\x9f\\x98\\x80"
    #
    # If we skipped the first step we would end up with Unicode codepoints instead of hex escaped
    # characters, which TEAL assembly cannot process:
    #     s.encode("unicode-escape").decode("latin-1")
    # e.g. "\n 😀" => "\\n \\U0001f600'"
    s = s.encode("utf-8").decode("latin-1").encode("unicode-escape").decode("latin-1")

    # Escape double quote characters (not covered by unicode-escape) but leave in single quotes
    s = s.replace('"', '\\"')

    # Surround string in double quotes
    return '"' + s + '"'

def format_bytes(b: bytes, encoding: AVMBytesEncoding) -> str:
    match encoding:
        case AVMBytesEncoding.utf8:
            return escape_utf8_string(b.decode())
        case AVMBytesEncoding.base32:
            return base64.b32encode(b).decode("ascii").rstrip("=")
        case AVMBytesEncoding.base64:
            return base64.b64encode(b).decode("ascii")
        case AVMBytesEncoding.base16 | AVMBytesEncoding.unknown:
            return f"0x{b.hex()}"

def format_tuple_index(var_type: wtypes.WTuple, base_name: str, index_or_name: int | str) -> str:
    # If a named tuple is indexed numerical, convert this to the item name
    if isinstance(index_or_name, int) and var_type.names is not None:
        index_or_name = var_type.names[index_or_name]
    return f"{base_name}.{index_or_name}"

def lvalue_items(tup: awst_nodes.TupleExpression) -> Sequence[awst_nodes.Lvalue]:
    items = list[awst_nodes.Lvalue]()
    for item in tup.items:
        assert isinstance(item, awst_nodes.Lvalue)
        items.append(item)
    return items

def format_error_comment(op: str, error_message: str) -> str:
    if op in ("err", "assert"):
        return error_message
    else:
        return f"on error: {error_message}"


import abc
import typing

from puya.context import CompileContext
from puya.ir.avm_ops_models import RunMode
from puya.ir.models import Contract, LogicSignature, ModuleArtifact, Program
from puya.ir.visitor import IRTraverser

class DestructuredIRValidator(IRTraverser, abc.ABC):
    def __init__(
        self,
        context: CompileContext,
        program: Program,
        run_mode: typing.Literal[RunMode.app, RunMode.lsig],
    ):
        self.context = context
        self.current_run_mode = run_mode
        self.active_program = program

    @classmethod
    def validate(cls, context: CompileContext, artifact: ModuleArtifact) -> None:
        match artifact:
            case LogicSignature() as l_sig:
                cls.validate_logic_sig(context, l_sig)
            case Contract() as contract:
                cls.validate_contract(context, contract)
            case _:
                typing.assert_never(artifact)

    @classmethod
    def validate_logic_sig(cls, context: CompileContext, logic_sig: LogicSignature) -> None:
        validator = cls(context, logic_sig.program, RunMode.lsig)
        for sub in logic_sig.program.all_subroutines:
            validator.visit_all_blocks(sub.body)

    @classmethod
    def validate_contract(cls, context: CompileContext, contract: Contract) -> None:
        for program in contract.all_programs():
            validator = cls(context, program, RunMode.app)
            for sub in program.all_subroutines:
                validator.visit_all_blocks(sub.body)


import typing
from collections.abc import Iterable

from puya import log
from puya.ir.models import (
    CompiledContractReference,
    CompiledLogicSigReference,
    Constant,
    TemplateVar,
    Value,
)
from puya.ir.validation._base import DestructuredIRValidator

logger = log.get_logger(__name__)

class CompileReferenceValidator(DestructuredIRValidator):
    @typing.override
    def visit_compiled_contract_reference(self, const: CompiledContractReference) -> None:
        _log_non_constant_values(const.template_variables.values())

    @typing.override
    def visit_compiled_logicsig_reference(self, const: CompiledLogicSigReference) -> None:
        _log_non_constant_values(const.template_variables.values())

def _log_non_constant_values(values: Iterable[Value]) -> None:
    for value in values:
        if isinstance(value, Constant):
            continue
        if isinstance(value, TemplateVar):
            logger.error(
                "nested template variables are not supported",
                location=value.source_location,
            )
        else:
            logger.error(
                "non-constant template value",
                location=value.source_location,
            )


import typing

from puya import log
from puya.ir.models import InnerTransactionField, ITxnConstant
from puya.ir.validation._base import DestructuredIRValidator

logger = log.get_logger(__name__)

class ITxnResultFieldValidator(DestructuredIRValidator):
    @typing.override
    def visit_itxn_constant(self, const: ITxnConstant) -> None:
        logger.info(
            "Potential cause of field access with non-constant group index",
            location=const.source_location,
        )

    @typing.override
    def visit_inner_transaction_field(self, field: InnerTransactionField) -> None:
        logger.error(
            "Inner transaction field access with non constant group index,"
            " to resolve move field access to same code path where the inner transaction is"
            " submitted",
            location=field.source_location,
        )


import attrs

from puya.context import CompileContext
from puya.ir.models import ModuleArtifact
from puya.ir.validation.compile_reference_validator import CompileReferenceValidator
from puya.ir.validation.itxn_result_field_validator import ITxnResultFieldValidator
from puya.ir.validation.min_avm_version_validator import MinAvmVersionValidator
from puya.ir.validation.op_run_mode_validator import OpRunModeValidator

def validate_module_artifact(context: CompileContext, artifact: ModuleArtifact) -> None:
    attrs.validate(artifact)
    for validator_cls in (
        OpRunModeValidator,
        MinAvmVersionValidator,
        ITxnResultFieldValidator,
        CompileReferenceValidator,
    ):
        validator_cls.validate(context, artifact)


import typing

from puya import log
from puya.ir.models import Intrinsic
from puya.ir.validation._base import DestructuredIRValidator

logger = log.get_logger(__name__)

class MinAvmVersionValidator(DestructuredIRValidator):
    @typing.override
    def visit_intrinsic_op(self, intrinsic: Intrinsic) -> None:
        program_avm_version = self.active_program.avm_version
        op_avm_version = intrinsic.op_variant.min_avm_version
        if op_avm_version > program_avm_version:
            op_desc = intrinsic.op.value
            # if variant min version differs from op min version, then include variant enum
            if op_avm_version != intrinsic.op.min_avm_version:
                op_desc += f" {intrinsic.op_variant.enum}"
            logger.error(
                f"Opcode {op_desc!r} requires a min AVM version of "
                f"{op_avm_version} but the target AVM version is"
                f" {program_avm_version}",
                location=intrinsic.source_location,
            )


import typing

from puya import log
from puya.ir.avm_ops_models import RunMode
from puya.ir.models import Intrinsic
from puya.ir.validation._base import DestructuredIRValidator

logger = log.get_logger(__name__)

class OpRunModeValidator(DestructuredIRValidator):
    @typing.override
    def visit_intrinsic_op(self, intrinsic: Intrinsic) -> None:
        match intrinsic.op_variant.supported_modes:
            case RunMode.any:
                pass
            case RunMode.app:
                if self.current_run_mode != RunMode.app:
                    logger.warning(
                        f"The operation {intrinsic} is only allowed in smart contracts",
                        location=intrinsic.source_location,
                    )
            case RunMode.lsig:
                if self.current_run_mode != RunMode.lsig:
                    logger.warning(
                        f"The operation {intrinsic} is only allowed in logic signatures",
                        location=intrinsic.source_location,
                    )
            case _:
                typing.assert_never(intrinsic.op_variant.supported_modes)


# ruff: noqa: ARG002

from __future__ import annotations  # needed to break import cycle

import typing as t
from abc import ABC, abstractmethod

if t.TYPE_CHECKING:
    from collections.abc import Iterable

    import puya.ir.models

class IRVisitor[T](ABC):
    @abstractmethod
    def visit_assignment(self, ass: puya.ir.models.Assignment) -> T: ...

    @abstractmethod
    def visit_register(self, reg: puya.ir.models.Register) -> T: ...

    @abstractmethod
    def visit_undefined(self, val: puya.ir.models.Undefined) -> T: ...

    @abstractmethod
    def visit_uint64_constant(self, const: puya.ir.models.UInt64Constant) -> T: ...

    @abstractmethod
    def visit_biguint_constant(self, const: puya.ir.models.BigUIntConstant) -> T: ...

    @abstractmethod
    def visit_bytes_constant(self, const: puya.ir.models.BytesConstant) -> T: ...

    @abstractmethod
    def visit_compiled_contract_reference(
        self, const: puya.ir.models.CompiledContractReference
    ) -> T: ...

    @abstractmethod
    def visit_compiled_logicsig_reference(
        self, const: puya.ir.models.CompiledLogicSigReference
    ) -> T: ...

    @abstractmethod
    def visit_address_constant(self, const: puya.ir.models.AddressConstant) -> T: ...

    @abstractmethod
    def visit_method_constant(self, const: puya.ir.models.MethodConstant) -> T: ...

    @abstractmethod
    def visit_itxn_constant(self, const: puya.ir.models.ITxnConstant) -> T: ...

    @abstractmethod
    def visit_phi(self, phi: puya.ir.models.Phi) -> T: ...

    @abstractmethod
    def visit_phi_argument(self, arg: puya.ir.models.PhiArgument) -> T: ...

    @abstractmethod
    def visit_intrinsic_op(self, intrinsic: puya.ir.models.Intrinsic) -> T: ...

    @abstractmethod
    def visit_inner_transaction_field(
        self, intrinsic: puya.ir.models.InnerTransactionField
    ) -> T: ...

    @abstractmethod
    def visit_invoke_subroutine(self, callsub: puya.ir.models.InvokeSubroutine) -> T: ...

    @abstractmethod
    def visit_value_tuple(self, tup: puya.ir.models.ValueTuple) -> T: ...

    @abstractmethod
    def visit_conditional_branch(self, branch: puya.ir.models.ConditionalBranch) -> T: ...

    @abstractmethod
    def visit_goto(self, goto: puya.ir.models.Goto) -> T: ...

    @abstractmethod
    def visit_goto_nth(self, goto_nth: puya.ir.models.GotoNth) -> T: ...

    @abstractmethod
    def visit_switch(self, switch: puya.ir.models.Switch) -> T: ...

    @abstractmethod
    def visit_subroutine_return(self, retsub: puya.ir.models.SubroutineReturn) -> T: ...

    @abstractmethod
    def visit_program_exit(self, exit_: puya.ir.models.ProgramExit) -> T: ...

    @abstractmethod
    def visit_fail(self, fail: puya.ir.models.Fail) -> T: ...

    @abstractmethod
    def visit_template_var(self, deploy_var: puya.ir.models.TemplateVar) -> T: ...

class IRTraverser(IRVisitor[None]):
    def visit_all_blocks(self, blocks: Iterable[puya.ir.models.BasicBlock]) -> None:
        for block in blocks:
            self.visit_block(block)

    def visit_block(self, block: puya.ir.models.BasicBlock) -> None:
        for op in list(block.all_ops):  # make a copy in case visitors need to modify ops
            op.accept(self)

    def visit_assignment(self, ass: puya.ir.models.Assignment) -> None:
        for target in ass.targets:
            target.accept(self)
        ass.source.accept(self)

    def visit_register(self, reg: puya.ir.models.Register) -> None:
        pass

    def visit_undefined(self, val: puya.ir.models.Undefined) -> None:
        pass

    def visit_uint64_constant(self, const: puya.ir.models.UInt64Constant) -> None:
        pass

    def visit_biguint_constant(self, const: puya.ir.models.BigUIntConstant) -> None:
        pass

    def visit_bytes_constant(self, const: puya.ir.models.BytesConstant) -> None:
        pass

    def visit_address_constant(self, const: puya.ir.models.AddressConstant) -> None:
        pass

    def visit_template_var(self, deploy_var: puya.ir.models.TemplateVar) -> None:
        pass

    def visit_method_constant(self, const: puya.ir.models.MethodConstant) -> None:
        pass

    def visit_compiled_contract_reference(
        self, const: puya.ir.models.CompiledContractReference
    ) -> None:
        for var in const.template_variables.values():
            var.accept(self)

    def visit_compiled_logicsig_reference(
        self, const: puya.ir.models.CompiledLogicSigReference
    ) -> None:
        for var in const.template_variables.values():
            var.accept(self)

    def visit_itxn_constant(self, const: puya.ir.models.ITxnConstant) -> None:
        pass

    def visit_phi(self, phi: puya.ir.models.Phi) -> None:
        phi.register.accept(self)
        for arg in phi.args:
            arg.accept(self)

    def visit_phi_argument(self, arg: puya.ir.models.PhiArgument) -> None:
        arg.value.accept(self)

    def visit_intrinsic_op(self, intrinsic: puya.ir.models.Intrinsic) -> None:
        for arg in intrinsic.args:
            arg.accept(self)

    def visit_inner_transaction_field(self, field: puya.ir.models.InnerTransactionField) -> None:
        field.group_index.accept(self)
        field.is_last_in_group.accept(self)
        if field.array_index:
            field.array_index.accept(self)

    def visit_invoke_subroutine(self, callsub: puya.ir.models.InvokeSubroutine) -> None:
        for arg in callsub.args:
            arg.accept(self)

    def visit_conditional_branch(self, branch: puya.ir.models.ConditionalBranch) -> None:
        branch.condition.accept(self)

    def visit_goto(self, goto: puya.ir.models.Goto) -> None:
        pass

    def visit_goto_nth(self, goto_nth: puya.ir.models.GotoNth) -> None:
        goto_nth.value.accept(self)

    def visit_switch(self, switch: puya.ir.models.Switch) -> None:
        switch.value.accept(self)
        for case in switch.cases:
            case.accept(self)

    def visit_subroutine_return(self, retsub: puya.ir.models.SubroutineReturn) -> None:
        for r in retsub.result:
            r.accept(self)

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

    def visit_fail(self, fail: puya.ir.models.Fail) -> None:
        pass

    def visit_value_tuple(self, tup: puya.ir.models.ValueTuple) -> None:
        for v in tup.values:
            v.accept(self)

class NoOpIRVisitor[T](IRVisitor[T | None]):
    def visit_assignment(self, ass: puya.ir.models.Assignment) -> T | None:
        return None

    def visit_register(self, reg: puya.ir.models.Register) -> T | None:
        return None

    def visit_undefined(self, val: puya.ir.models.Undefined) -> T | None:
        return None

    def visit_uint64_constant(self, const: puya.ir.models.UInt64Constant) -> T | None:
        return None

    def visit_biguint_constant(self, const: puya.ir.models.BigUIntConstant) -> T | None:
        return None

    def visit_bytes_constant(self, const: puya.ir.models.BytesConstant) -> T | None:
        return None

    def visit_address_constant(self, const: puya.ir.models.AddressConstant) -> T | None:
        return None

    def visit_method_constant(self, const: puya.ir.models.MethodConstant) -> T | None:
        return None

    def visit_itxn_constant(self, const: puya.ir.models.ITxnConstant) -> T | None:
        return None

    def visit_compiled_contract_reference(
        self, const: puya.ir.models.CompiledContractReference
    ) -> T | None:
        return None

    def visit_compiled_logicsig_reference(
        self, const: puya.ir.models.CompiledLogicSigReference
    ) -> T | None:
        return None

    def visit_phi(self, phi: puya.ir.models.Phi) -> T | None:
        return None

    def visit_template_var(self, deploy_var: puya.ir.models.TemplateVar) -> T | None:
        return None

    def visit_phi_argument(self, arg: puya.ir.models.PhiArgument) -> T | None:
        return None

    def visit_intrinsic_op(self, intrinsic: puya.ir.models.Intrinsic) -> T | None:
        return None

    def visit_inner_transaction_field(
        self, field: puya.ir.models.InnerTransactionField
    ) -> T | None:
        return None

    def visit_invoke_subroutine(self, callsub: puya.ir.models.InvokeSubroutine) -> T | None:
        return None

    def visit_value_tuple(self, tup: puya.ir.models.ValueTuple) -> T | None:
        return None

    def visit_conditional_branch(self, branch: puya.ir.models.ConditionalBranch) -> T | None:
        return None

    def visit_goto(self, goto: puya.ir.models.Goto) -> T | None:
        return None

    def visit_goto_nth(self, goto_nth: puya.ir.models.GotoNth) -> T | None:
        return None

    def visit_switch(self, switch: puya.ir.models.Switch) -> T | None:
        return None

    def visit_subroutine_return(self, retsub: puya.ir.models.SubroutineReturn) -> T | None:
        return None

    def visit_program_exit(self, exit_: puya.ir.models.ProgramExit) -> T | None:
        return None

    def visit_fail(self, fail: puya.ir.models.Fail) -> T | None:
        return None


from collections.abc import Iterable, Mapping

import attrs

from puya.ir import models
from puya.ir.visitor_mutator import IRMutator

@attrs.define
class MemoryReplacer(IRMutator):
    _replacements: Mapping[models.Register, models.Register]
    replaced: int = 0

    @classmethod
    def apply(
        cls,
        blocks: Iterable[models.BasicBlock],
        *,
        replacements: Mapping[models.Register, models.Register],
    ) -> int:
        if not replacements:
            return 0
        replacer = cls(replacements=replacements)
        for block in blocks:
            replacer.visit_block(block)
        return replacer.replaced

    def visit_register(self, reg: models.Register) -> models.Register:
        try:
            replacement = self._replacements[reg]
        except KeyError:
            return reg
        self.replaced += 1
        return replacement