harness/pybytecode_core/rep.py

#!/usr/bin/env python3
"""INPUT REPRESENTATION v2 — stop destroying information.

Phase 2 measured the tuned model reproducing docstrings 1/80 times, and I nearly filed it as
hallucination. It was my bug: the v1 disassembly emits ONLY the instruction stream, and a function's
docstring lives in `co_consts[0]` where NO opcode ever references it (CPython sets `__doc__`
implicitly at function-creation). The model was asked to reconstruct information it could not see.

WHAT IS ACTUALLY DESTROYED — measured (probe_consts.py), not assumed:
  * co_consts[0] docstrings: DESTROYED. Every one of them.
  * every other unreferenced const: only the implicit `None`. Carries nothing.
  * co_names: 0 unreferenced. Every name ALREADY appears in the instruction stream via argrepr
    (LOAD_GLOBAL/LOAD_ATTR/STORE_NAME...). Dumping a NAMES table would add zero information and
    cost input length -- so v2 does NOT dump it. (The brief asked for it; the measurement says no.)
  * signature shape: posonly/kwonly/vararg/kwarg boundaries were flattened into a flat arg list.
    Recoverable from co_posonlyargcount/co_kwonlyargcount/co_flags -- v1 just threw it away. This
    matters now that the corpus admits richer signatures than the pilot's.
  * default values: NOT destroyed (LOAD_CONST'd on the enclosing frame before MAKE_FUNCTION).

DESIGN — the rep must be a pure function of the code object, because at inference we have only a
.pyc. Everything below is read straight off the CodeType. Nothing is sourced from the .py file.

  CODE qualname(a, b, /, c, *args, d, **kw)     <- full signature shape
    DOC 'the docstring'                          <- co_consts[0], the destroyed information
    <instructions>
    EXC ...
  END

DOC is emitted as a first-class line rather than a raw CONSTS dump: a raw table would re-print
every already-referenced const (they appear inline at their LOAD_CONST) to recover exactly one
missing entry -- pure token inflation for no information. We emit precisely what was lost.
"""
from __future__ import annotations
import dis, types

JUMP_OPS = set(dis.hasjrel) | set(dis.hasjabs)

CO_VARARGS, CO_VARKEYWORDS = 0x04, 0x08


def _labels_for(co: types.CodeType) -> dict[int, str]:
    pts: set[int] = set()
    for i in dis.get_instructions(co):
        if i.opcode in JUMP_OPS and isinstance(i.argval, int):
            pts.add(i.argval)
    for e in dis._parse_exception_table(co):  # noqa: SLF001
        pts.update((e.start, e.end, e.target))   # e.end included — see the EXC note below
    return {off: f"L{n}" for n, off in enumerate(sorted(pts), 1)}


def _child_names(co: types.CodeType) -> dict[int, str]:
    return {id(c): c.co_qualname for c in co.co_consts if isinstance(c, types.CodeType)}


def _fmt_arg(i, labels: dict[int, str], kids: dict[int, str]) -> str:
    if i.arg is None:
        return ""
    if i.opcode in JUMP_OPS and isinstance(i.argval, int):
        return labels.get(i.argval, f"@{i.argval}")
    v = i.argval
    if isinstance(v, types.CodeType):
        return f"<code {kids.get(id(v), v.co_qualname)}>"
    if i.opname in ("LOAD_CONST", "RETURN_CONST", "KW_NAMES"):
        return repr(v)
    if i.argrepr:
        return i.argrepr
    return str(i.arg)


def signature(co: types.CodeType) -> str:
    """The full signature SHAPE, straight off the code object: posonly '/', kwonly '*',
    *args and **kwargs. v1 flattened all of this into a bare comma list."""
    n_pos, n_all = co.co_posonlyargcount, co.co_argcount
    n_kw = co.co_kwonlyargcount
    names = list(co.co_varnames)
    parts: list[str] = []
    idx = 0
    for k in range(n_all):
        parts.append(names[k])
        if k + 1 == n_pos and n_pos:
            parts.append("/")
    idx = n_all
    star_done = False
    if co.co_flags & CO_VARARGS:
        # *args sits right after the positional block
        va = names[n_all + n_kw]
        parts.append(f"*{va}")
        star_done = True
    if n_kw and not star_done:
        parts.append("*")
    for k in range(n_kw):
        parts.append(names[n_all + k])
    if co.co_flags & CO_VARKEYWORDS:
        off = n_all + n_kw + (1 if co.co_flags & CO_VARARGS else 0)
        parts.append(f"**{names[off]}")
    return ", ".join(parts)


CO_OPTIMIZED = 0x01
_CONST_OPS = ("LOAD_CONST", "RETURN_CONST", "KW_NAMES")


def docstring_of_published(co: types.CodeType) -> str | None:
    """The rule used to build the PUBLISHED benchmarks. Kept verbatim; do not 'fix' in place.

    It is wrong -- see docstring_of() -- but the published `bench.jsonl` inputs, the v3 training
    corpus and every scored generation were produced with it. Reproducing a published number
    requires reproducing the input that produced it, so this stays.
    """
    if co.co_consts and isinstance(co.co_consts[0], str):
        return co.co_consts[0]
    return None


def docstring_of(co: types.CodeType) -> str | None:
    """co_consts[0] iff it is the IMPLICIT docstring slot -- the information v1 destroyed.

    `docstring_of_published` asked only "is co_consts[0] a string", which is not the question:
    whether a slot-0 string is a docstring depends on WHAT KIND OF SCOPE the code object is.
    Measured over the two published benchmarks, the published rule emitted a DOC line for 98 code
    objects that have no implicit docstring: 94 on csn-600 (13.4% of its DOC lines) and all 4 on
    mbpp-383. A DOC line is presented to the model as recovered ground truth, so a wrong one is
    not noise -- it is an instruction to reproduce a docstring that was never in the source.

    What CPython 3.12 actually does, measured rather than assumed:

    * A FUNCTION / lambda / async function with no docstring RESERVES slot 0 and fills it with
      `None`; with a docstring, slot 0 is the docstring. So for these scopes "slot 0 is a string"
      is already exactly right, and the published rule was never wrong here. (This is why the
      defect is invisible at function level: 2 of 605 function-scope emissions.)

    * A GENERATOR EXPRESSION is CO_OPTIMIZED like a function but does NOT reserve slot 0 -- its
      first ordinary literal lands there. `sum(x for x in xs if x != '=')` yields
      `co_consts[0] == '='`, and the published rule reports `'='` as a docstring. A genexp has no
      `__doc__` at all. 8 such emissions on csn-600, 3 on mbpp-383.

    * A MODULE or CLASS BODY stores `__doc__` EXPLICITLY (`LOAD_CONST <doc>; STORE_NAME __doc__`),
      so its docstring is already visible in the instruction stream and nothing was destroyed --
      emitting DOC duplicates it. And when there is no docstring, slot 0 is simply the first
      const, which may be a DEFAULT ARGUMENT value reached only inside the defaults tuple
      (`co_consts = ('WMAP5', <code>, None, ('WMAP5',))`) -- so it is not even referenced
      directly, and an "is slot 0 loaded" test would still call it a docstring. 86 emissions on
      csn-600, 1 on mbpp-383.

    Hence: function-like scope, and not one of the angle-bracketed compiler-generated scopes.
    `<genexpr>`, `<listcomp>`, `<setcomp>`, `<dictcomp>` and `<lambda>` can never collide with a
    user identifier, so co_name is a sound discriminator. A lambda cannot carry a docstring
    either (its body is a single expression), so excluding it costs nothing.

    Exact on both published benchmarks against `ast.get_docstring` ground truth: 0 spurious,
    0 missed, over 2,158 code objects. See test_rep_docstring.py.
    """
    if not (co.co_flags & CO_OPTIMIZED):
        return None                          # module / class body: explicit, nothing was lost
    if co.co_name.startswith("<"):
        return None                          # <genexpr>/<listcomp>/<lambda>: no __doc__ exists
    if co.co_consts and isinstance(co.co_consts[0], str):
        return co.co_consts[0]
    return None


def disassemble_v2(co: types.CodeType, out: list[str] | None = None, *,
                   doc_rule: str = "published") -> str:
    """Render the code object.

    `doc_rule` selects which DOC rule to apply, and DEFAULTS TO THE PUBLISHED ONE ON PURPOSE.
    The DOC line is part of the model's INPUT: the shipped v3 weights were trained on inputs
    built with `docstring_of_published`, and every published score was measured against those
    inputs. Silently switching the default would change what the public benchmark asks of the
    model and would invalidate the numbers on its own card. Pass doc_rule="fixed" to use the
    corrected rule -- see docstring_of() -- which is the right default only from a retrain
    onward.
    """
    if doc_rule not in ("published", "fixed"):
        raise ValueError(f"doc_rule must be 'published' or 'fixed', got {doc_rule!r}")
    out = [] if out is None else out
    labels, kids = _labels_for(co), _child_names(co)
    out.append(f"CODE {co.co_qualname}({signature(co)})")
    doc = (docstring_of if doc_rule == "fixed" else docstring_of_published)(co)
    if doc is not None:
        out.append(f"  DOC {doc!r}")
    for i in dis.get_instructions(co):
        if i.offset in labels:
            out.append(f"{labels[i.offset]}:")
        a = _fmt_arg(i, labels, kids)
        out.append(f"  {i.opname} {a}".rstrip())
    for e in dis._parse_exception_table(co):  # noqa: SLF001
        # THE `end` IS LOAD-BEARING AND v2.0 OMITTED IT. Found at n=279: the model produced
        #     try: A            (except: pass) else: B; C
        # where the original was
        #     try: A; B; C      (except: pass)
        # CPython compiles an `else` block OUTSIDE the try range, so BOTH forms emit the SAME
        # instruction stream and differ ONLY in how far the exception table's range extends. With
        # `end` omitted, the two programs were BYTE-IDENTICAL in the model's input -- it could not
        # possibly tell them apart, and it guessed wrong. Exactly the docstring bug again: the model
        # was blamed for hallucinating information I had deleted.
        s = labels.get(e.start, f"@{e.start}")
        en = labels.get(e.end, f"@{e.end}")
        t = labels.get(e.target, f"@{e.target}")
        out.append(f"  EXC try={s}..{en} -> handler={t} depth={e.depth} lasti={e.lasti}")
    out.append("END")
    for c in co.co_consts:
        if isinstance(c, types.CodeType):
            disassemble_v2(c, out, doc_rule=doc_rule)
    return "\n".join(out)