bikini-patchwork/patchwork/abyss.py

from __future__ import annotations

import ast
import base64
import json
import random
from dataclasses import dataclass
from typing import Any

from .util import gen_bytes

DISPATCH_NAME = "__pw_ab_dispatch__"
ASSETS_NAME = "__pw_ab_assets__"

_OPS = (
    "CONST",
    "LOAD",
    "STORE",
    "POP",
    "DUP",
    "BIN",
    "UNARY",
    "COMPARE_CHAIN",
    "JUMP",
    "JUMP_IF_FALSE",
    "JUMP_IF_TRUE_KEEP",
    "JUMP_IF_FALSE_KEEP",
    "CALL",
    "GET_ATTR",
    "SUBSCR",
    "BUILD_SLICE",
    "BUILD_LIST",
    "BUILD_TUPLE",
    "BUILD_SET",
    "BUILD_DICT",
    "RETURN",
    "GET_ITER",
    "FOR_ITER",
    "UNPACK",
    "BUILD_STRING",
    "FORMAT_VALUE",
)

_BIN_OPS: tuple[tuple[type[ast.operator], str], ...] = (
    (ast.Add, "add"),
    (ast.Sub, "sub"),
    (ast.Mult, "mul"),
    (ast.MatMult, "matmul"),
    (ast.Div, "truediv"),
    (ast.FloorDiv, "floordiv"),
    (ast.Mod, "mod"),
    (ast.Pow, "pow"),
    (ast.LShift, "lshift"),
    (ast.RShift, "rshift"),
    (ast.BitOr, "or"),
    (ast.BitXor, "xor"),
    (ast.BitAnd, "and"),
)

_UNARY_OPS: tuple[tuple[type[ast.unaryop], str], ...] = (
    (ast.Invert, "invert"),
    (ast.Not, "not"),
    (ast.UAdd, "pos"),
    (ast.USub, "neg"),
)

_COMPARE_OPS: tuple[tuple[type[ast.cmpop], str], ...] = (
    (ast.Eq, "eq"),
    (ast.NotEq, "ne"),
    (ast.Lt, "lt"),
    (ast.LtE, "le"),
    (ast.Gt, "gt"),
    (ast.GtE, "ge"),
    (ast.Is, "is"),
    (ast.IsNot, "is_not"),
    (ast.In, "in"),
    (ast.NotIn, "not_in"),
)


class UnsupportedAbyssNode(ValueError):
    pass


@dataclass(frozen=True)
class EncodedAbyssAssets:
    payload: str
    key: str
    opcodes: dict[str, int]


_JUMP_OPS = {
    "JUMP",
    "JUMP_IF_FALSE",
    "JUMP_IF_TRUE_KEEP",
    "JUMP_IF_FALSE_KEEP",
    "FOR_ITER",
}

_PACKET_LAYOUTS = (
    (0, 1, 2),
    (0, 2, 1),
    (1, 0, 2),
    (1, 2, 0),
    (2, 0, 1),
    (2, 1, 0),
)


class _Label:
    def __init__(self) -> None:
        self.index: int | None = None


class _Emitter:
    def __init__(self) -> None:
        self.instructions: list[list[Any]] = []

    def label(self) -> _Label:
        return _Label()

    def mark(self, label: _Label) -> None:
        label.index = len(self.instructions)

    def emit(self, op: str, *args: Any) -> None:
        self.instructions.append([op, *args])

    def resolve(self) -> list[list[Any]]:
        resolved: list[list[Any]] = []
        for inst in self.instructions:
            op, *args = inst
            new_args: list[Any] = []
            for arg in args:
                if isinstance(arg, _Label):
                    if arg.index is None:
                        raise UnsupportedAbyssNode("internal unresolved VM label")
                    new_args.append(arg.index)
                else:
                    new_args.append(arg)
            resolved.append([op, *new_args])
        return resolved


def _op_name(op: ast.AST, table: tuple[tuple[type[Any], str], ...]) -> str:
    for cls, name in table:
        if isinstance(op, cls):
            return name
    raise UnsupportedAbyssNode(f"unsupported operator {type(op).__name__}")


def _strip_docstring(body: list[ast.stmt]) -> list[ast.stmt]:
    if (
        body
        and isinstance(body[0], ast.Expr)
        and isinstance(body[0].value, ast.Constant)
        and isinstance(body[0].value.value, str)
    ):
        return body[1:]
    return body


def _argument_names(args: ast.arguments) -> set[str]:
    names = {arg.arg for arg in args.posonlyargs + args.args + args.kwonlyargs}
    if args.vararg is not None:
        names.add(args.vararg.arg)
    if args.kwarg is not None:
        names.add(args.kwarg.arg)
    return names


def _target_names(target: ast.AST) -> set[str]:
    if isinstance(target, ast.Name):
        return {target.id}
    if isinstance(target, (ast.Tuple, ast.List)):
        names: set[str] = set()
        for elt in target.elts:
            names.update(_target_names(elt))
        return names
    return set()


class _ScopeCollector(ast.NodeVisitor):
    def __init__(self, args: ast.arguments) -> None:
        self.local_names = set(_argument_names(args))
        self.global_names: set[str] = set()

    def visit_Global(self, node: ast.Global) -> None:
        self.global_names.update(node.names)

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
        return

    visit_AsyncFunctionDef = visit_FunctionDef

    def visit_ClassDef(self, node: ast.ClassDef) -> None:
        return

    def visit_Lambda(self, node: ast.Lambda) -> None:
        return

    def visit_Assign(self, node: ast.Assign) -> None:
        for target in node.targets:
            self.local_names.update(_target_names(target))
        self.visit(node.value)

    def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
        self.local_names.update(_target_names(node.target))
        if node.value is not None:
            self.visit(node.value)

    def visit_AugAssign(self, node: ast.AugAssign) -> None:
        self.local_names.update(_target_names(node.target))
        self.visit(node.value)

    def visit_For(self, node: ast.For) -> None:
        self.local_names.update(_target_names(node.target))
        self.visit(node.iter)
        for stmt in node.body + node.orelse:
            self.visit(stmt)

    def visit_NamedExpr(self, node: ast.NamedExpr) -> None:
        self.local_names.update(_target_names(node.target))
        self.visit(node.value)

    def visit_GeneratorExp(self, node: ast.GeneratorExp) -> None:
        for generator in node.generators:
            self.local_names.update(_target_names(generator.target))
            self.visit(generator.iter)
            for condition in generator.ifs:
                self.visit(condition)
        self.visit(node.elt)

    visit_ListComp = visit_GeneratorExp

    def finalize(self) -> tuple[set[str], set[str]]:
        return self.local_names - self.global_names, self.global_names


class AbyssCompiler:
    def __init__(self) -> None:
        self.emitter = _Emitter()
        self.constants: list[Any] = []
        self.local_names: set[str] = set()
        self.global_names: set[str] = set()
        self.external_names: set[str] = set()
        self.loop_stack: list[tuple[_Label, _Label, int]] = []
        self._temp_counter = 0

    def compile(self, node: ast.FunctionDef) -> dict[str, Any]:
        self._reject_function_shape(node)
        collector = _ScopeCollector(node.args)
        for stmt in node.body:
            collector.visit(stmt)
        self.local_names, self.global_names = collector.finalize()

        body = _strip_docstring(node.body)
        if not body:
            self.emitter.emit("CONST", self._const(None))
            self.emitter.emit("RETURN")
        else:
            for stmt in body:
                self._stmt(stmt)
            self.emitter.emit("CONST", self._const(None))
            self.emitter.emit("RETURN")

        return {
            "name": node.name,
            "code": self.emitter.resolve(),
            "consts": list(self.constants),
            "globals": sorted(self.global_names),
            "locals": sorted(self.local_names),
            "externals": sorted(self.external_names),
        }

    def _reject_function_shape(self, node: ast.FunctionDef) -> None:
        for child in ast.walk(node):
            if isinstance(child, (ast.Yield, ast.YieldFrom, ast.Await)):
                raise UnsupportedAbyssNode("generators, coroutines, and await are not supported")
            if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
                if child is not node:
                    raise UnsupportedAbyssNode("nested functions/classes are not supported")

    def _const(self, value: Any) -> int:
        idx = len(self.constants)
        self.constants.append(value)
        return idx

    def _temp(self) -> str:
        self._temp_counter += 1
        name = f"_pw_ab_tmp_{self._temp_counter}"
        self.local_names.add(name)
        return name

    def _stmt(self, node: ast.stmt) -> None:
        if isinstance(node, ast.Return):
            if node.value is None:
                self.emitter.emit("CONST", self._const(None))
            else:
                self._expr(node.value)
            self.emitter.emit("RETURN")
            return
        if isinstance(node, ast.Assign):
            self._expr(node.value)
            for index, target in enumerate(node.targets):
                if index < len(node.targets) - 1:
                    self.emitter.emit("DUP")
                self._store_target(target)
            return
        if isinstance(node, ast.AnnAssign):
            if node.value is not None:
                self._expr(node.value)
                self._store_target(node.target)
            return
        if isinstance(node, ast.AugAssign):
            if not isinstance(node.target, ast.Name):
                raise UnsupportedAbyssNode("augmented assignment only supports local/global names")
            self._expr(ast.Name(id=node.target.id, ctx=ast.Load()))
            self._expr(node.value)
            self.emitter.emit("BIN", _op_name(node.op, _BIN_OPS))
            self.emitter.emit("STORE", node.target.id)
            return
        if isinstance(node, ast.Expr):
            self._expr(node.value)
            self.emitter.emit("POP")
            return
        if isinstance(node, ast.Pass):
            return
        if isinstance(node, ast.Global):
            return
        if isinstance(node, ast.If):
            else_label = self.emitter.label()
            end_label = self.emitter.label()
            self._expr(node.test)
            self.emitter.emit("JUMP_IF_FALSE", else_label)
            for stmt in node.body:
                self._stmt(stmt)
            self.emitter.emit("JUMP", end_label)
            self.emitter.mark(else_label)
            for stmt in node.orelse:
                self._stmt(stmt)
            self.emitter.mark(end_label)
            return
        if isinstance(node, ast.While):
            start_label = self.emitter.label()
            end_label = self.emitter.label()
            self.emitter.mark(start_label)
            self._expr(node.test)
            self.emitter.emit("JUMP_IF_FALSE", end_label)
            self.loop_stack.append((start_label, end_label, 0))
            for stmt in node.body:
                self._stmt(stmt)
            self.loop_stack.pop()
            self.emitter.emit("JUMP", start_label)
            self.emitter.mark(end_label)
            if node.orelse:
                raise UnsupportedAbyssNode("while/else is not supported")
            return
        if isinstance(node, ast.For):
            if node.orelse:
                raise UnsupportedAbyssNode("for/else is not supported")
            start_label = self.emitter.label()
            end_label = self.emitter.label()
            self._expr(node.iter)
            self.emitter.emit("GET_ITER")
            self.emitter.mark(start_label)
            self.emitter.emit("FOR_ITER", end_label)
            self._store_target(node.target)
            self.loop_stack.append((start_label, end_label, 1))
            for stmt in node.body:
                self._stmt(stmt)
            self.loop_stack.pop()
            self.emitter.emit("JUMP", start_label)
            self.emitter.mark(end_label)
            return
        if isinstance(node, ast.Break):
            if not self.loop_stack:
                raise UnsupportedAbyssNode("break outside loop")
            _, end_label, cleanup = self.loop_stack[-1]
            for _ in range(cleanup):
                self.emitter.emit("POP")
            self.emitter.emit("JUMP", end_label)
            return
        if isinstance(node, ast.Continue):
            if not self.loop_stack:
                raise UnsupportedAbyssNode("continue outside loop")
            start_label, _, _ = self.loop_stack[-1]
            self.emitter.emit("JUMP", start_label)
            return
        raise UnsupportedAbyssNode(f"unsupported statement {type(node).__name__}")

    def _store_target(self, target: ast.AST) -> None:
        if isinstance(target, ast.Name):
            self.emitter.emit("STORE", target.id)
            return
        if isinstance(target, (ast.Tuple, ast.List)):
            self.emitter.emit("UNPACK", len(target.elts))
            for elt in target.elts:
                self._store_target(elt)
            return
        raise UnsupportedAbyssNode(f"unsupported assignment target {type(target).__name__}")

    def _expr(self, node: ast.expr) -> None:
        if isinstance(node, ast.Constant):
            self.emitter.emit("CONST", self._const(node.value))
            return
        if isinstance(node, ast.Name):
            if isinstance(node.ctx, ast.Load):
                if node.id not in self.local_names or node.id in self.global_names:
                    self.external_names.add(node.id)
                self.emitter.emit("LOAD", node.id)
                return
            raise UnsupportedAbyssNode("name expression is not load context")
        if isinstance(node, ast.BinOp):
            self._expr(node.left)
            self._expr(node.right)
            self.emitter.emit("BIN", _op_name(node.op, _BIN_OPS))
            return
        if isinstance(node, ast.UnaryOp):
            self._expr(node.operand)
            self.emitter.emit("UNARY", _op_name(node.op, _UNARY_OPS))
            return
        if isinstance(node, ast.BoolOp):
            self._bool_op(node)
            return
        if isinstance(node, ast.Compare):
            self._expr(node.left)
            for comparator in node.comparators:
                self._expr(comparator)
            self.emitter.emit("COMPARE_CHAIN", [_op_name(op, _COMPARE_OPS) for op in node.ops])
            return
        if isinstance(node, ast.IfExp):
            else_label = self.emitter.label()
            end_label = self.emitter.label()
            self._expr(node.test)
            self.emitter.emit("JUMP_IF_FALSE", else_label)
            self._expr(node.body)
            self.emitter.emit("JUMP", end_label)
            self.emitter.mark(else_label)
            self._expr(node.orelse)
            self.emitter.mark(end_label)
            return
        if isinstance(node, ast.Call):
            if any(isinstance(arg, ast.Starred) for arg in node.args):
                raise UnsupportedAbyssNode("*args calls are not supported")
            if any(keyword.arg is None for keyword in node.keywords):
                raise UnsupportedAbyssNode("**kwargs calls are not supported")
            self._expr(node.func)
            for arg in node.args:
                self._expr(arg)
            keyword_names: list[str] = []
            for keyword in node.keywords:
                if keyword.arg is None:
                    raise UnsupportedAbyssNode("**kwargs calls are not supported")
                keyword_names.append(keyword.arg)
                self._expr(keyword.value)
            self.emitter.emit("CALL", len(node.args), keyword_names)
            return
        if isinstance(node, ast.Attribute):
            self._expr(node.value)
            self.emitter.emit("GET_ATTR", node.attr)
            return
        if isinstance(node, ast.Subscript):
            self._expr(node.value)
            self._expr(node.slice)
            self.emitter.emit("SUBSCR")
            return
        if isinstance(node, ast.Slice):
            for part in (node.lower, node.upper, node.step):
                if part is None:
                    self.emitter.emit("CONST", self._const(None))
                else:
                    self._expr(part)
            self.emitter.emit("BUILD_SLICE")
            return
        if isinstance(node, ast.List):
            self._sequence(node.elts, "BUILD_LIST")
            return
        if isinstance(node, (ast.GeneratorExp, ast.ListComp)):
            self._comprehension(node.elt, node.generators)
            return
        if isinstance(node, ast.Tuple):
            self._sequence(node.elts, "BUILD_TUPLE")
            return
        if isinstance(node, ast.Set):
            self._sequence(node.elts, "BUILD_SET")
            return
        if isinstance(node, ast.Dict):
            for key, value in zip(node.keys, node.values):
                if key is None:
                    raise UnsupportedAbyssNode("dictionary unpacking is not supported")
                self._expr(key)
                self._expr(value)
            self.emitter.emit("BUILD_DICT", len(node.keys))
            return
        if isinstance(node, ast.JoinedStr):
            for value in node.values:
                if isinstance(value, ast.Constant) and isinstance(value.value, str):
                    self.emitter.emit("CONST", self._const(value.value))
                elif isinstance(value, ast.FormattedValue):
                    self._formatted_value(value)
                else:
                    raise UnsupportedAbyssNode("unsupported f-string part")
            self.emitter.emit("BUILD_STRING", len(node.values))
            return
        if isinstance(node, ast.FormattedValue):
            self._formatted_value(node)
            return
        if isinstance(node, ast.NamedExpr):
            if not isinstance(node.target, ast.Name):
                raise UnsupportedAbyssNode("walrus target must be a name")
            self._expr(node.value)
            self.emitter.emit("DUP")
            self._store_target(node.target)
            return
        raise UnsupportedAbyssNode(f"unsupported expression {type(node).__name__}")

    def _comprehension(self, elt: ast.expr, generators: list[ast.comprehension]) -> None:
        if len(generators) != 1:
            raise UnsupportedAbyssNode("only single-generator comprehensions are supported")
        generator = generators[0]
        if generator.is_async:
            raise UnsupportedAbyssNode("async comprehensions are not supported")
        result_name = self._temp()
        start_label = self.emitter.label()
        end_label = self.emitter.label()
        self.emitter.emit("BUILD_LIST", 0)
        self.emitter.emit("STORE", result_name)
        self._expr(generator.iter)
        self.emitter.emit("GET_ITER")
        self.emitter.mark(start_label)
        self.emitter.emit("FOR_ITER", end_label)
        self._store_target(generator.target)
        for condition in generator.ifs:
            self._expr(condition)
            self.emitter.emit("JUMP_IF_FALSE", start_label)
        self.emitter.emit("LOAD", result_name)
        self.emitter.emit("GET_ATTR", "append")
        self._expr(elt)
        self.emitter.emit("CALL", 1, [])
        self.emitter.emit("POP")
        self.emitter.emit("JUMP", start_label)
        self.emitter.mark(end_label)
        self.emitter.emit("LOAD", result_name)

    def _bool_op(self, node: ast.BoolOp) -> None:
        if not node.values:
            raise UnsupportedAbyssNode("empty boolean operation")
        end_label = self.emitter.label()
        jump_op = "JUMP_IF_TRUE_KEEP" if isinstance(node.op, ast.Or) else "JUMP_IF_FALSE_KEEP"
        for index, value in enumerate(node.values):
            self._expr(value)
            if index < len(node.values) - 1:
                self.emitter.emit(jump_op, end_label)
        self.emitter.mark(end_label)

    def _sequence(self, elts: list[ast.expr], op: str) -> None:
        for elt in elts:
            if isinstance(elt, ast.Starred):
                raise UnsupportedAbyssNode("starred literals are not supported")
            self._expr(elt)
        self.emitter.emit(op, len(elts))

    def _formatted_value(self, node: ast.FormattedValue) -> None:
        self._expr(node.value)
        if node.format_spec is not None:
            self._expr(node.format_spec)
            self.emitter.emit("FORMAT_VALUE", node.conversion, True)
        else:
            self.emitter.emit("FORMAT_VALUE", node.conversion, False)


class AbyssTransformer(ast.NodeTransformer):
    def __init__(
        self,
        rng: random.Random,
        *,
        targets: set[str] | None = None,
        auto: bool = False,
        dispatch_name: str = DISPATCH_NAME,
    ) -> None:
        self.rng = rng
        self.targets = set(targets or ())
        self.auto = auto
        self.dispatch_name = dispatch_name
        self.assets: list[dict[str, Any]] = []
        self.keep_names: set[str] = {dispatch_name}
        self.protected: list[str] = []
        self.skipped: list[tuple[str, str]] = []
        self._matched_targets: set[str] = set()
        self._class_stack: list[str] = []
        self._function_depth = 0

    @property
    def explicit(self) -> bool:
        return bool(self.targets)

    def protect(self, tree: ast.Module) -> ast.Module:
        new_tree = self.visit(tree)
        if not isinstance(new_tree, ast.Module):
            raise TypeError("expected module")
        missing = sorted(self.targets - self._matched_targets)
        if missing:
            raise ValueError(f"abyss target(s) not found: {', '.join(missing)}")
        return new_tree

    def visit_ClassDef(self, node: ast.ClassDef) -> ast.AST:
        self._class_stack.append(node.name)
        node.body = [self.visit(stmt) for stmt in node.body]
        self._class_stack.pop()
        return node

    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.AST:
        qualname = ".".join([*self._class_stack, node.name]) if self._class_stack else node.name
        target_names = {node.name, qualname}
        targeted = bool(self.targets & target_names)
        desired = self._function_depth == 0 and (targeted or (self.auto and not self.explicit))
        if targeted:
            self._matched_targets.update(self.targets & target_names)
        if not desired:
            self._function_depth += 1
            node = self.generic_visit(node)
            self._function_depth -= 1
            return node

        try:
            asset = AbyssCompiler().compile(node)
        except UnsupportedAbyssNode as exc:
            if targeted:
                raise ValueError(f"abyss target {qualname} is unsupported: {exc}") from exc
            self.skipped.append((qualname, str(exc)))
            self._function_depth += 1
            node = self.generic_visit(node)
            self._function_depth -= 1
            return node

        fid = len(self.assets)
        self.assets.append(asset)
        self.keep_names.update(asset["externals"])
        self.protected.append(qualname)
        node.body = self._wrapper_body(fid, _function_docstring(node))
        return node

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AST:
        qualname = ".".join([*self._class_stack, node.name]) if self._class_stack else node.name
        target_names = {node.name, qualname}
        targeted = bool(self.targets & target_names)
        if targeted:
            self._matched_targets.update(self.targets & target_names)
            raise ValueError(f"abyss target {qualname} is unsupported: async functions are not supported")
        return node

    def _wrapper_body(self, fid: int, docstring: str | None) -> list[ast.stmt]:
        call = ast.Call(
            func=ast.Name(id=self.dispatch_name, ctx=ast.Load()),
            args=[
                ast.Constant(value=fid),
                ast.Call(func=ast.Name(id="locals", ctx=ast.Load()), args=[], keywords=[]),
            ],
            keywords=[],
        )
        body: list[ast.stmt] = []
        if docstring is not None:
            body.append(ast.Expr(value=ast.Constant(value=docstring)))
        body.append(ast.Return(value=call))
        return body


def _function_docstring(node: ast.FunctionDef) -> str | None:
    if (
        node.body
        and isinstance(node.body[0], ast.Expr)
        and isinstance(node.body[0].value, ast.Constant)
        and isinstance(node.body[0].value.value, str)
    ):
        return node.body[0].value.value
    return None


def _encode_const(value: Any) -> dict[str, Any]:
    if value is None:
        return {"t": "none"}
    if value is Ellipsis:
        return {"t": "ellipsis"}
    if isinstance(value, bool):
        return {"t": "bool", "v": value}
    if isinstance(value, int):
        return {"t": "int", "v": str(value)}
    if isinstance(value, float):
        return {"t": "float", "v": repr(value)}
    if isinstance(value, str):
        return {"t": "str", "v": value}
    if isinstance(value, bytes):
        return {"t": "bytes", "v": base64.b85encode(value).decode("ascii")}
    raise UnsupportedAbyssNode(f"unsupported constant type {type(value).__name__}")


def _json_bytes(value: Any) -> bytes:
    return json.dumps(value, separators=(",", ":"), sort_keys=True).encode("utf-8")


def _seal_value(value: Any, rng: random.Random, salt: int) -> list[Any]:
    raw = _json_bytes(value)
    key = gen_bytes(rng, rng.randint(9, 24))
    add = rng.randrange(256)
    step = rng.randrange(1, 256, 2)
    encoded = bytes(((byte ^ key[index % len(key)]) + add + ((index + salt) * step)) & 255 for index, byte in enumerate(raw))
    share = gen_bytes(rng, len(encoded))
    other = bytes(left ^ right for left, right in zip(encoded, share))
    return [
        base64.b85encode(share).decode("ascii"),
        base64.b85encode(other).decode("ascii"),
        base64.b85encode(key).decode("ascii"),
        add,
        step,
    ]


def _tag_for(opcode: int, argc: int, index: int, salt: int) -> int:
    return ((opcode * 2654435761) ^ (argc * 2246822519) ^ ((index + salt) * 3266489917)) & 0xFFFFFFFF


def _encode_packet(inst: list[Any], index: int, rng: random.Random, salt: int) -> list[Any]:
    opcode = inst[0]
    args = inst[1:]
    tag = _tag_for(opcode, len(args), index, salt)
    fields = [
        _seal_value(opcode, rng, salt + index * 5 + 1),
        _seal_value(args, rng, salt + index * 5 + 2),
        _seal_value(tag, rng, salt + index * 5 + 3),
    ]
    layout_id = rng.randrange(len(_PACKET_LAYOUTS))
    layout = _PACKET_LAYOUTS[layout_id]
    return [layout_id, fields[layout[0]], fields[layout[1]], fields[layout[2]]]


def _rebased_instruction(inst: list[Any], base: int) -> list[Any]:
    op = inst[0]
    if op in _JUMP_OPS:
        return [op, base + inst[1], *inst[2:]]
    return list(inst)


def encode_assets(assets: list[dict[str, Any]], rng: random.Random) -> EncodedAbyssAssets:
    values = rng.sample(range(1, 251), len(_OPS))
    opcodes = dict(zip(_OPS, values))
    encoded_funcs: list[dict[str, Any]] = []
    merged_code: list[list[Any]] = []
    salt = rng.randrange(1 << 30)
    for fn_index, asset in enumerate(assets):
        entry = len(merged_code)
        rebased_code = [_rebased_instruction(inst, entry) for inst in asset["code"]]
        merged_code.extend(rebased_code)
        encoded_funcs.append(
            {
                "c": [
                    _seal_value(_encode_const(value), rng, salt + (fn_index + 1) * 100000 + const_index)
                    for const_index, value in enumerate(asset["consts"])
                ],
                "e": _seal_value(entry, rng, salt + 200000 + fn_index),
                "g": _seal_value(asset["globals"], rng, salt + 300000 + fn_index),
                "l": _seal_value(asset["locals"], rng, salt + 400000 + fn_index),
            }
        )
    encoded_packets = [
        _encode_packet([opcodes[inst[0]], *inst[1:]], index, rng, salt)
        for index, inst in enumerate(merged_code)
    ]
    document = _json_bytes({"v": 2, "m": {"s": salt}, "f": encoded_funcs, "p": encoded_packets})
    key = gen_bytes(rng, rng.randint(24, 48))
    encrypted = bytes(byte ^ key[index % len(key)] for index, byte in enumerate(document))
    return EncodedAbyssAssets(
        payload=base64.b85encode(encrypted).decode("ascii"),
        key=base64.b85encode(key).decode("ascii"),
        opcodes=opcodes,
    )


def build_runtime_stmts(encoded: EncodedAbyssAssets, rng: random.Random) -> list[ast.stmt]:
    blocks = _runtime_dispatch_blocks(encoded.opcodes)
    rng.shuffle(blocks)
    dispatch_blocks = "\n".join(blocks)
    layouts_src = repr(_PACKET_LAYOUTS)
    source = f'''
{ASSETS_NAME} = ({encoded.payload!r}, {encoded.key!r})
__pw_ab_cache__ = None

def __pw_ab_open__(_box, _salt):
    _base64 = __import__('base64')
    _json = __import__('json')
    _share = _base64.b85decode(_box[0].encode('ascii'))
    _other = _base64.b85decode(_box[1].encode('ascii'))
    _key = _base64.b85decode(_box[2].encode('ascii'))
    _add = _box[3]
    _step = _box[4]
    _encoded = bytes(_a ^ _b for _a, _b in zip(_share, _other))
    _raw = bytearray(len(_encoded))
    for _i, _b in enumerate(_encoded):
        _raw[_i] = ((_b - _add - ((_i + _salt) * _step)) & 255) ^ _key[_i % len(_key)]
    return _json.loads(bytes(_raw).decode('utf-8'))

def __pw_ab_tag__(_opcode, _argc, _index, _salt):
    return ((_opcode * 2654435761) ^ (_argc * 2246822519) ^ ((_index + _salt) * 3266489917)) & 4294967295

def __pw_ab_packet__(_packet, _index, _meta):
    _layouts = {layouts_src}
    _layout = _layouts[_packet[0] % len(_layouts)]
    _fields = [None, None, None]
    for _pos, _slot in enumerate(_layout):
        _fields[_slot] = _packet[_pos + 1]
    _salt = _meta['s']
    _opcode = __pw_ab_open__(_fields[0], _salt + _index * 5 + 1)
    _args = __pw_ab_open__(_fields[1], _salt + _index * 5 + 2)
    _tag = __pw_ab_open__(_fields[2], _salt + _index * 5 + 3)
    if _tag != __pw_ab_tag__(_opcode, len(_args), _index, _salt):
        raise RuntimeError('invalid abyss packet')
    return [_opcode, *_args]

def __pw_ab_const__(_x):
    _t = _x['t']
    if _t == 'none':
        return None
    if _t == 'ellipsis':
        return Ellipsis
    if _t == 'bool':
        return bool(_x['v'])
    if _t == 'int':
        return int(_x['v'])
    if _t == 'float':
        return float(_x['v'])
    if _t == 'str':
        return _x['v']
    if _t == 'bytes':
        return __import__('base64').b85decode(_x['v'].encode('ascii'))
    raise RuntimeError('invalid abyss constant')

def __pw_ab_load__():
    global __pw_ab_cache__
    if __pw_ab_cache__ is None:
        _base64 = __import__('base64')
        _json = __import__('json')
        _payload, _key = {ASSETS_NAME}
        _enc = _base64.b85decode(_payload.encode('ascii'))
        _raw_key = _base64.b85decode(_key.encode('ascii'))
        _raw = bytes(_b ^ _raw_key[_i % len(_raw_key)] for _i, _b in enumerate(_enc))
        _doc = _json.loads(_raw.decode('utf-8'))
        _meta = _doc['m']
        for _fi, _fn in enumerate(_doc['f']):
            _fn['c'] = [__pw_ab_const__(__pw_ab_open__(_item, _meta['s'] + (_fi + 1) * 100000 + _ci)) for _ci, _item in enumerate(_fn['c'])]
            _fn['e'] = __pw_ab_open__(_fn['e'], _meta['s'] + 200000 + _fi)
            _fn['g'] = __pw_ab_open__(_fn['g'], _meta['s'] + 300000 + _fi)
            _fn['l'] = __pw_ab_open__(_fn['l'], _meta['s'] + 400000 + _fi)
        __pw_ab_cache__ = _doc
    return __pw_ab_cache__

def __pw_ab_get__(_name, _locals, _globals, _builtins, _declared_globals, _declared_locals):
    if _name not in _declared_globals and _name in _locals:
        return _locals[_name]
    if _name not in _declared_globals and _name in _declared_locals:
        raise UnboundLocalError("cannot access local variable '" + _name + "' where it is not associated with a value")
    if _name in _globals:
        return _globals[_name]
    if _name in _builtins:
        return _builtins[_name]
    raise NameError(_name)

def __pw_ab_store__(_name, _value, _locals, _globals, _declared_globals):
    if _name in _declared_globals:
        _globals[_name] = _value
    else:
        _locals[_name] = _value

def __pw_ab_bin__(_op, _a, _b):
    if _op == 'add':
        return _a + _b
    if _op == 'sub':
        return _a - _b
    if _op == 'mul':
        return _a * _b
    if _op == 'matmul':
        return _a @ _b
    if _op == 'truediv':
        return _a / _b
    if _op == 'floordiv':
        return _a // _b
    if _op == 'mod':
        return _a % _b
    if _op == 'pow':
        return _a ** _b
    if _op == 'lshift':
        return _a << _b
    if _op == 'rshift':
        return _a >> _b
    if _op == 'or':
        return _a | _b
    if _op == 'xor':
        return _a ^ _b
    if _op == 'and':
        return _a & _b
    raise RuntimeError('invalid abyss binary op')

def __pw_ab_unary__(_op, _a):
    if _op == 'invert':
        return ~_a
    if _op == 'not':
        return not _a
    if _op == 'pos':
        return +_a
    if _op == 'neg':
        return -_a
    raise RuntimeError('invalid abyss unary op')

def __pw_ab_compare_one__(_op, _a, _b):
    if _op == 'eq':
        return _a == _b
    if _op == 'ne':
        return _a != _b
    if _op == 'lt':
        return _a < _b
    if _op == 'le':
        return _a <= _b
    if _op == 'gt':
        return _a > _b
    if _op == 'ge':
        return _a >= _b
    if _op == 'is':
        return _a is _b
    if _op == 'is_not':
        return _a is not _b
    if _op == 'in':
        return _a in _b
    if _op == 'not_in':
        return _a not in _b
    raise RuntimeError('invalid abyss compare op')

def __pw_ab_compare__(_ops, _values):
    for _idx, _op in enumerate(_ops):
        if not __pw_ab_compare_one__(_op, _values[_idx], _values[_idx + 1]):
            return False
    return True

def __pw_ab_format__(_value, _conversion, _has_spec, _stack):
    _spec = _stack.pop() if _has_spec else ''
    if _conversion == 115:
        _value = str(_value)
    elif _conversion == 114:
        _value = repr(_value)
    elif _conversion == 97:
        _value = ascii(_value)
    return format(_value, _spec)

def __pw_ab_exec__(_doc, _fn, _initial_locals, _globals):
    _builtins = _globals.get('__builtins__', __builtins__)
    if not isinstance(_builtins, dict):
        _builtins = _builtins.__dict__
    _declared_globals = set(_fn.get('g', ()))
    _declared_locals = set(_fn.get('l', ()))
    _locals = dict(_initial_locals)
    _consts = _fn['c']
    _code = _doc['p']
    _meta = _doc['m']
    _stack = []
    _ip = _fn['e']
    while True:
        _inst = __pw_ab_packet__(_code[_ip], _ip, _meta)
        _op = _inst[0]
{dispatch_blocks}
        raise RuntimeError('invalid abyss opcode')

def {DISPATCH_NAME}(_fid, _env):
    _doc = __pw_ab_load__()
    return __pw_ab_exec__(_doc, _doc['f'][_fid], _env, globals())
'''
    return ast.parse(source).body


def _runtime_dispatch_blocks(opcodes: dict[str, int]) -> list[str]:
    return [
        f"""        if _op == {opcodes['CONST']}:
            _stack.append(_consts[_inst[1]])
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['LOAD']}:
            _stack.append(__pw_ab_get__(_inst[1], _locals, _globals, _builtins, _declared_globals, _declared_locals))
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['STORE']}:
            __pw_ab_store__(_inst[1], _stack.pop(), _locals, _globals, _declared_globals)
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['POP']}:
            _stack.pop()
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['DUP']}:
            _stack.append(_stack[-1])
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['BIN']}:
            _b = _stack.pop()
            _a = _stack.pop()
            _stack.append(__pw_ab_bin__(_inst[1], _a, _b))
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['UNARY']}:
            _stack.append(__pw_ab_unary__(_inst[1], _stack.pop()))
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['COMPARE_CHAIN']}:
            _ops = _inst[1]
            _values = [_stack.pop() for _ in range(len(_ops) + 1)]
            _values.reverse()
            _stack.append(__pw_ab_compare__(_ops, _values))
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['JUMP']}:
            _ip = _inst[1]
            continue""",
        f"""        if _op == {opcodes['JUMP_IF_FALSE']}:
            _ip = _inst[1] if not _stack.pop() else _ip + 1
            continue""",
        f"""        if _op == {opcodes['JUMP_IF_TRUE_KEEP']}:
            if _stack[-1]:
                _ip = _inst[1]
            else:
                _stack.pop()
                _ip += 1
            continue""",
        f"""        if _op == {opcodes['JUMP_IF_FALSE_KEEP']}:
            if not _stack[-1]:
                _ip = _inst[1]
            else:
                _stack.pop()
                _ip += 1
            continue""",
        f"""        if _op == {opcodes['CALL']}:
            _argc = _inst[1]
            _kw_names = _inst[2]
            _kw = {{}}
            for _name in reversed(_kw_names):
                _kw[_name] = _stack.pop()
            _args = [_stack.pop() for _ in range(_argc)]
            _args.reverse()
            _func = _stack.pop()
            _stack.append(_func(*_args, **_kw))
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['GET_ATTR']}:
            _stack.append(getattr(_stack.pop(), _inst[1]))
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['SUBSCR']}:
            _key = _stack.pop()
            _obj = _stack.pop()
            _stack.append(_obj[_key])
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['BUILD_SLICE']}:
            _step = _stack.pop()
            _upper = _stack.pop()
            _lower = _stack.pop()
            _stack.append(slice(_lower, _upper, _step))
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['BUILD_LIST']}:
            _items = [_stack.pop() for _ in range(_inst[1])]
            _items.reverse()
            _stack.append(_items)
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['BUILD_TUPLE']}:
            _items = [_stack.pop() for _ in range(_inst[1])]
            _items.reverse()
            _stack.append(tuple(_items))
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['BUILD_SET']}:
            _items = [_stack.pop() for _ in range(_inst[1])]
            _items.reverse()
            _stack.append(set(_items))
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['BUILD_DICT']}:
            _items = []
            for _ in range(_inst[1]):
                _value = _stack.pop()
                _key = _stack.pop()
                _items.append((_key, _value))
            _items.reverse()
            _stack.append(dict(_items))
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['RETURN']}:
            return _stack.pop() if _stack else None""",
        f"""        if _op == {opcodes['GET_ITER']}:
            _stack.append(iter(_stack.pop()))
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['FOR_ITER']}:
            try:
                _stack.append(next(_stack[-1]))
                _ip += 1
            except StopIteration:
                _stack.pop()
                _ip = _inst[1]
            continue""",
        f"""        if _op == {opcodes['UNPACK']}:
            _items = list(_stack.pop())
            if len(_items) != _inst[1]:
                raise ValueError('not enough values to unpack' if len(_items) < _inst[1] else 'too many values to unpack')
            for _item in reversed(_items):
                _stack.append(_item)
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['BUILD_STRING']}:
            _items = [_stack.pop() for _ in range(_inst[1])]
            _items.reverse()
            _stack.append(''.join(str(_item) for _item in _items))
            _ip += 1
            continue""",
        f"""        if _op == {opcodes['FORMAT_VALUE']}:
            _stack.append(__pw_ab_format__(_stack.pop() if not _inst[2] else _stack.pop(-2), _inst[1], _inst[2], _stack))
            _ip += 1
            continue""",
    ]