Tools/jit/_writer.py

"""Utilities for writing JIT build artifacts out to C header files."""

import typing
import math

import _dwarf
import _stencils


def _dump_footer(
    groups: dict[str, _stencils.StencilGroup], symbols: dict[str, int]
) -> typing.Iterator[str]:
    symbol_mask_size = max(math.ceil(len(symbols) / 32), 1)
    yield f'static_assert(SYMBOL_MASK_WORDS >= {symbol_mask_size}, "SYMBOL_MASK_WORDS too small");'
    yield ""
    yield "typedef struct {"
    yield "    void (*emit)("
    yield "        unsigned char *code, unsigned char *data, _PyExecutorObject *executor,"
    yield "        const _PyUOpInstruction *instruction, jit_state *state);"
    yield "    size_t code_size;"
    yield "    size_t data_size;"
    yield "    symbol_mask trampoline_mask;"
    yield "    symbol_mask got_mask;"
    yield "} StencilGroup;"
    yield ""
    yield "static const StencilGroup stencil_groups[MAX_UOP_REGS_ID + 1] = {"
    for opname, group in sorted(groups.items()):
        yield f"    [{opname}] = {group.as_c(opname)},"
    yield "};"
    yield ""
    yield f"static const void * const symbols_map[{max(len(symbols), 1)}] = {{"
    if symbols:
        for symbol, ordinal in symbols.items():
            yield f"    [{ordinal}] = &{symbol},"
    else:
        yield "    0"
    yield "};"


def _dump_stencil(opname: str, group: _stencils.StencilGroup) -> typing.Iterator[str]:
    yield "void"
    yield f"emit_{opname}("
    yield "    unsigned char *code, unsigned char *data, _PyExecutorObject *executor,"
    yield "    const _PyUOpInstruction *instruction, jit_state *state)"
    yield "{"
    for part, stencil in [("code", group.code), ("data", group.data)]:
        for line in stencil.disassembly:
            yield f"    // {line}"
        stripped = stencil.body.rstrip(b"\x00")
        if stripped:
            yield f"    const unsigned char {part}_body[{len(stencil.body)}] = {{"
            for i in range(0, len(stripped), 8):
                row = " ".join(f"{byte:#04x}," for byte in stripped[i : i + 8])
                yield f"        {row}"
            yield "    };"
    # Data is written first (so relaxations in the code work properly):
    for part, stencil in [("data", group.data), ("code", group.code)]:
        if stencil.body.rstrip(b"\x00"):
            yield f"    memcpy({part}, {part}_body, sizeof({part}_body));"
        stencil.holes.sort(key=lambda hole: hole.offset)
        for hole in stencil.holes:
            yield f"    {hole.as_c(part)}"
    yield "}"
    yield ""


def dump(
    groups: dict[str, _stencils.StencilGroup], symbols: dict[str, int]
) -> typing.Iterator[str]:
    """Yield a JIT compiler line-by-line as a C header file."""
    for opname, group in groups.items():
        yield from _dump_stencil(opname, group)
    yield from _dump_footer(groups, symbols)


def dump_unwind_info(
    unwind_info: _dwarf.UnwindInfo | None,
) -> typing.Iterator[str]:
    """Yield JIT unwind information line-by-line as a C header file."""
    if unwind_info is None:
        yield "#define JIT_UNWIND_INFO_SUPPORTED 0"
        return

    yield "#define JIT_UNWIND_INFO_SUPPORTED 1"
    fields = [
        ("JIT_UNWIND_CODE_ALIGNMENT_FACTOR", unwind_info.code_alignment_factor),
        ("JIT_UNWIND_DATA_ALIGNMENT_FACTOR", unwind_info.data_alignment_factor),
        ("JIT_UNWIND_RA_REG", unwind_info.return_address_register),
        ("JIT_UNWIND_CFA_REG", unwind_info.cfa_register),
        ("JIT_UNWIND_CFA_OFFSET", unwind_info.cfa_offset),
        ("JIT_UNWIND_FP_REG", unwind_info.frame_pointer_register),
        ("JIT_UNWIND_FP_OFFSET", unwind_info.frame_pointer_offset),
        ("JIT_UNWIND_RA_OFFSET", unwind_info.return_address_offset),
    ]
    for name, value in fields:
        yield f"#define {name} {value}"