composed_encoding.py

"""Composed multiply+reduce CoT — the DEPLOYABLE surface.

divcot validates the reduction in isolation by feeding N=x*y as input; a real
submission only gets (x, y, p) with x,y < p. This encoding chains the two
validated scaffolds in one generated trace:

  1. schoolbook multiply as a running accumulator: after consuming each y-limb
     y_j the model emits acc_j = acc_{j-1} + x*y_j*B^j (2W limbs, LSB-first so
     the product+add carry chain is local). acc_{W-1} = N.
  2. copy-reverse: re-emit N MSB-first (attention reversal — staging the
     dividend so the division block looks exactly like validated divcot).
  3. long division by p: quotient limbs MSB-first with running-remainder
     install targets, then the answer (= final remainder) LSB-first.

Surface:
  "A" x(W,LSB) "B" y(W,LSB) "M" p(W,MSB) "="
      [ "a" acc_j(2W,LSB) ] * W
      "N" N(2W,MSB)  Q(2W)  "R" ans(W,LSB) "\n"

Tokens/example = 2W^2 + 9W + 7 (W=5: 102; base-100 tier 5 W=10: 297).

Install targets (NTP-aligned):
  acarry  carry into each emitted acc limb (the multiply-accumulate state)
  rem{j}  MSB limb j of the running remainder at each quotient position
  qhat    leading-limbs quotient estimate at each quotient position
  ans     limb k of (x*y) mod p at each answer position

Compliance: identical status to divcot/karatsuba — fixed trace SHAPE, every
value generated by trained weights; randomising weights collapses accuracy.
"""
from __future__ import annotations

from limbs import limb_char, limb_str, parse_limbs, to_limbs
from divcot2_encoding import long_division, qhat_estimate

AMARK, BMARK, DIVMARK, EQ, ACC, NMARK, REVMARK, NL = "A", "B", "M", "=", "a", "N", "R", "\n"
PLSB = "m"


def prompt_str(x: int, y: int, p: int, W: int, base: int, subpad: bool = False,
               mulonly: bool = False) -> str:
    """subpad adds p LSB-first to the prompt: the MSB field serves quotient
    estimation, the LSB field aligns with the LSB-emitted qd*p / remainder
    chains (per-step diagnostic: subtraction digits need REVERSED access into
    the MSB-only field — the aligned acc stage learned, the division didn't)."""
    if mulonly:                                  # tier 0: pure multiplication
        return AMARK + limb_str(x, W, base) + BMARK + limb_str(y, W, base) + EQ
    s = (AMARK + limb_str(x, W, base) + BMARK + limb_str(y, W, base)
         + DIVMARK + limb_str(p, W, base, msb_first=True))
    if subpad:
        s += PLSB + limb_str(p, W, base)
    return s + EQ


def gen_len(W: int, scratch: bool = False, cursor: bool = False,
            subpad: bool = False, stepidx: bool = False,
            skiptriv: bool = False, subnum: bool = False,
            bemit: bool = False, srt: bool = False,
            mulonly: bool = False) -> int:
    """Tokens generated after '=': W acc blocks, N restage, quotient, answer.
    scratch=True: each quotient limb is followed by the running remainder
    (W limbs, LSB-first) — externalizes the division state, which the model
    cannot hold internally past W~3 (composed tier-3 diagnostic: first
    quotient digit right, then collapse).
    cursor=True (requires scratch): each step block starts with the consumed
    dividend limb, copied from the restaged N — an explicit progress anchor.
    Diagnostic at scratch step-4k: mid-sequence remainder drift + the model
    losing count of steps (9 blocks emitted instead of 2W); copying gives the
    same position mechanism that put restage at 1.00.
    subpad=True (requires cursor): each step also emits qd*p (W+1 limbs, LSB)
    before the remainder — externalizes the multiply-subtract that the
    per-step diagnostic showed failing (steps 0-4 = copies, perfect; real
    division steps 5+ at 0.04-0.15 even teacher-forced). Tier-2's compact
    success could memorize qd*p over 48 primes; tier 3's 400 primes need the
    generic circuit, so generate it like the (perfectly learned) acc rows.
    stepidx=True (requires cursor): each step opens with its 2-digit ASCII
    index, and the restage pairs every dividend limb with its index — a
    content-addressable name per step. Tier-4 diagnostic (W=10, 20 steps):
    steps 0-9 at ~1.00, cliff to ~0.45 at 10+; rem fails worst (0.36 late)
    despite all-local operands, so the wall is step MISBINDING under
    relative-position (ALiBi) attention — the number of identical-looking
    distractor blocks grows with step index. Indexed steps turn 'previous
    step's remainder' and 'my dividend limb' into content lookups.
    skiptriv=True (requires stepidx): x,y < p guarantees Q = N//p < p, so the
    first W division steps ALWAYS emit qd=0 — provably trivial filler that
    doubles the distractor count, dilutes gradient (250 of 828 tokens at
    W=10), and pushes real steps ~250 tokens further from the prompt's p
    field. Skip them: division = W real steps keeping their TRUE indices
    (so restage lookups still match), starting from r_prev = N's top half.
    Also re-emits p LSB-first ('m' + W limbs) right after the restage, so
    qd*p rows read a local p copy instead of a ~700-token fetch (tier-4
    stepidx diagnostic: qdp 0.78-0.84 was the residual fetch weakness).
    subnum=True (requires subpad): each step copies its numerator
    num = r_prev*base + cursor limb (W+1 limbs, LSB) between cursor and qd,
    so qd / qd*p / rem are all WITHIN-step local arithmetic. skiptriv
    diagnostic: qdp hit 0.9949 once p became a local read while rem
    plateaued at 0.43 fetching the previous step's rem block -- every pure
    copy in this stack reaches ~1.00, every fetch+arithmetic fusion stalls;
    decouple them (the subpad lesson, applied to the subtraction).
    bemit=True (requires subnum): the remainder block becomes (borrow, limb)
    pairs -- the borrow bit is EMITTED as an ASCII 0/1 token before each rem
    limb, so b_k sits in context when rem[k] is predicted and b_{k+1} is a
    local function of (num[k], qdp[k], b_k), all adjacent. v4 diagnostic:
    rem errors were 58% unstructured noise, 0.54 even at b=0 -- the
    install-only borrow supervision never built a circuit; emit the chain
    (the scratch lesson, applied to the borrow state)."""
    if mulonly:                                  # acc rows + answer copy
        return W * (1 + 2 * W) + 1 + 2 * W
    restage = 1 + (3 if stepidx else 1) * 2 * W
    pcopy = (1 + W) if skiptriv else 0
    if srt:
        # stepi + cursor + num complement (W+2) + qd + |qd|*p (W+1) + pairs
        per_step = 2 + 1 + (W + 2) + 1 + (W + 1) + 2 * (W + 2)
        final = 1 + 2 * W                        # sign + correction pairs
        return (W * (1 + 2 * W) + restage + pcopy + W * per_step
                + final + 1 + W)
    per_step = ((2 if stepidx else 0) + (1 if cursor else 0)
                + ((W + 1) if subnum else 0) + 1
                + ((W + 1) if subpad else 0)
                + ((2 * W) if bemit else (W if scratch else 0)))
    n_steps = W if skiptriv else 2 * W
    return W * (1 + 2 * W) + restage + pcopy + n_steps * per_step + 1 + W


def _srt_qd(num: int, p: int, base: int) -> int:
    """Signed SRT digit from |num|'s top-4 / p's top-3 limbs (exact integer
    rounding). Guarantees |num - qd*p| < p given |num| < base*p (verified:
    0 violations / 500k tier-4 selections, worst |s|/p = 0.54)."""
    if num == 0:
        return 0
    sgn = 1 if num > 0 else -1
    an = abs(num)
    nw = 1
    while base ** nw <= an:
        nw += 1
    pw = 1
    while base ** pw <= p:
        pw += 1
    n4 = an // base ** max(0, nw - 4)
    p3 = p // base ** max(0, pw - 3)
    sh = max(0, nw - 4) - max(0, pw - 3)
    if sh >= 0:
        a, b = n4 * base ** sh, p3
    else:
        a, b = n4, p3 * base ** -sh
    q = (2 * a + b) // (2 * b)                   # round(a/b), exact
    return sgn * min(base - 1, q)


def build_example(x: int, y: int, p: int, W: int, base: int, scratch: bool = False,
                  cursor: bool = False, subpad: bool = False, stepidx: bool = False,
                  skiptriv: bool = False, subnum: bool = False,
                  bemit: bool = False, srt: bool = False, mulonly: bool = False,
                  srt_dither: float = 0.0):
    """Return (text, ann). x,y < p < base**W (mulonly: x,y < base**W)."""
    assert not subnum or subpad, "subnum requires subpad"
    assert not bemit or subnum, "bemit requires subnum"
    assert not srt or (skiptriv and bemit), "srt requires skiptriv+bemit"
    assert not srt or base == 10, "srt: signed-digit tokens are base-10 only"
    assert not cursor or scratch, "cursor requires scratch"
    assert not subpad or cursor, "subpad requires cursor"
    assert not stepidx or cursor, "stepidx requires cursor"
    assert not skiptriv or stepidx, "skiptriv requires stepidx"
    assert not stepidx or 2 * W <= 100, "stepidx: 2-digit indices cap at 100 steps"
    assert not skiptriv or (x < p and y < p), "skiptriv: Q < p needs x,y < p"
    N = x * y
    y_limbs = to_limbs(y, W, base)
    x_limbs = to_limbs(x, W, base)
    parts, ann_g = [], []                       # ann_g: (gen_index, var, val)
    pos = 0

    acc = 0
    for j, yj in enumerate(y_limbs):            # multiply-accumulate rows
        prev = acc
        acc = prev + x * yj * base ** j
        parts.append(ACC)
        pos += 1
        prev_l = to_limbs(prev, 2 * W, base)
        c = 0
        for i in range(2 * W):                  # carry into limb i of acc_j
            ann_g.append((pos + i, "acarry", c))
            xi = x_limbs[i - j] if 0 <= i - j < W else 0
            c = (prev_l[i] + xi * yj + c) // base
        parts.append(limb_str(acc, 2 * W, base))
        pos += 2 * W
    assert acc == N

    if mulonly:                                  # tier 0: answer = N itself
        parts.append(NMARK)
        pos += 1
        n_l = to_limbs(N, 2 * W, base)
        for k in range(2 * W):
            ann_g.append((pos + k, "ans", n_l[k]))
        parts.append(limb_str(N, 2 * W, base) + NL)
        pos += 2 * W + 1
        prompt = prompt_str(x, y, p, W, base, mulonly=True)
        text = prompt + "".join(parts)
        assert len(text) == len(prompt) + gen_len(W, mulonly=True) + 1
        ann = [dict() for _ in range(len(text))]
        for gi, var, val in ann_g:
            ann[len(prompt) + gi - 1][var] = int(val)
        return text, ann

    n_msb = to_limbs(N, 2 * W, base, msb_first=True)
    if stepidx:                                  # restage as (index, limb) pairs
        parts.append(NMARK + "".join(f"{i:02d}" + limb_char(n_msb[i])
                                     for i in range(2 * W)))
        pos += 1 + 3 * 2 * W
    else:
        parts.append(NMARK + limb_str(N, 2 * W, base, msb_first=True))  # restage MSB
        pos += 1 + 2 * W

    if skiptriv:                                 # local p copy for the qd*p rows
        parts.append(PLSB + limb_str(p, W, base))
        pos += 1 + W

    if srt:
        # SRT-style signed-digit division: qd in {-9..9} chosen by a top-4/
        # top-3 leading-limbs lookup; remainder kept in (-p, p) as 10's-
        # complement limbs. Redundancy absorbs lookup error: wherever the
        # ratio sits near a digit boundary BOTH neighbours are valid
        # (|num - qd*p| < p either way), so the knife-edge cases that broke
        # exact-compare qd (0.96 ceiling, all errors off-by-one at margin
        # ~1e-3) become don't-cares. Verified: 0 invariant violations /
        # 500k selections, worst |s|/p = 0.54. num = 10*s + d in complement
        # is a pure shift-append COPY (no arithmetic).
        K = W + 2
        r = N // base ** W                       # after the W trivial steps
        for i in range(W, 2 * W):
            d = n_msb[i]
            num = r * base + d
            parts.append(f"{i:02d}" + limb_char(d))      # stepi + cursor
            pos += 3
            num_c = to_limbs((num + base ** K) % base ** K, K, base)
            parts.append("".join(limb_char(v) for v in num_c))
            pos += K
            qd = _srt_qd(num, p, base)
            if srt_dither:
                # Deterministic pseudo-random alternate-digit training: SRT's
                # redundancy makes qd+-1 often equally valid, but the model
                # only ever saw label-digit continuations -- free-running it
                # sometimes picks the OTHER valid digit and must continue
                # consistently from its own choice. Dithered traces teach
                # exactly those continuations. (hash-derived: reproducible)
                h = hash((x, y, p, i)) & 0xFFFF
                if h < int(srt_dither * 0x10000):
                    alt = qd + (1 if h & 1 else -1)
                    if abs(alt) <= base - 1 and abs(num - alt * p) < p:
                        qd = alt
            assert abs(num - qd * p) < p, (x, y, p, i, num, qd)
            ann_g.append((pos, "qhat", qd + base - 1))   # 19-class signed label
            parts.append(chr(0x90 + base - 1 + qd))
            pos += 1
            aq = abs(qd) * p
            p_l = to_limbs(p, W, base) + [0]
            c = 0
            for k in range(W + 1):               # carry into limb k of |qd|*p
                ann_g.append((pos + k, "spcarry", c))
                c = (abs(qd) * p_l[k] + c) // base
            parts.append(limb_str(aq, W + 1, base))
            pos += W + 1
            s = num - qd * p
            s_c = to_limbs((s + base ** K) % base ** K, K, base)
            aq_l = to_limbs(aq, K, base)
            ch = 0
            for k in range(K):                   # (chain-bit, limb) pairs
                parts.append(str(ch) + limb_char(s_c[k]))
                if qd >= 0:                      # subtract: borrow chain
                    ch = 1 if num_c[k] - aq_l[k] - ch < 0 else 0
                else:                            # negative digit: add carries
                    ch = (num_c[k] + aq_l[k] + ch) // base
            pos += 2 * K
            r = s
        neg = 1 if r < 0 else 0                  # final: ans = r + p if r<0
        parts.append(str(neg))
        pos += 1
        r_c = to_limbs((r + base ** (W + 1)) % base ** (W + 1), W + 1, base)
        pn_l = to_limbs(p * neg, W + 1, base)
        answer = r + p * neg
        a_l = to_limbs(answer, W, base)
        ch = 0
        for k in range(W):                       # (carry, limb) correction pairs
            parts.append(str(ch) + limb_char(a_l[k]))
            ch = (r_c[k] + pn_l[k] + ch) // base
        pos += 2 * W
        assert answer == N % p, (x, y, p)
        q_limbs = []                             # skip the classic loop below

    if not srt:
        q_limbs, rems, answer = long_division(N, p, W, base)
    r_prev = 0
    for i, (qd, r) in enumerate(zip(q_limbs, rems) if not srt else ()):
        if skiptriv and i < W:                   # provably qd=0 (Q < p): skip
            assert qd == 0
            r_prev = r
            continue
        if stepidx:                              # the step announces its own name
            parts.append(f"{i:02d}")
            pos += 2
        if cursor:                               # copy the consumed dividend limb
            parts.append(limb_char(n_msb[i]))
            pos += 1
        if subnum:                               # numerator copy: r_prev*B + limb
            num = r_prev * base + n_msb[i]
            parts.append(limb_str(num, W + 1, base))
            pos += W + 1
        # Install the TRUE quotient digit, not the leading-limbs estimate: at
        # tier 4 (10-digit p) qhat_estimate disagrees with qd in 46.6% of
        # steps (top-2/top-1 degrades with divisor width), so the probe was
        # supervising noise — the model reached qd 0.975 by correcting it,
        # but the install should shape where the true digit is computed.
        ann_g.append((pos, "qhat", qd))
        if not scratch:                          # compact: remainder is internal state
            rstr = to_limbs(r, W, base, msb_first=True)
            for j in range(W):
                ann_g.append((pos, f"rem{j}", rstr[j]))
        parts.append(limb_char(qd))
        pos += 1
        if subpad:                               # emit qd*p LSB-first (generic mult)
            qdp = qd * p
            p_l = to_limbs(p, W, base) + [0]
            c = 0
            for k in range(W + 1):               # carry into limb k of qd*p
                ann_g.append((pos + k, "spcarry", c))
                c = (qd * p_l[k] + c) // base
            parts.append(limb_str(qdp, W + 1, base))
            pos += W + 1
        if scratch:                              # emit remainder LSB-first (local borrow)
            num_l = [n_msb[i]] + to_limbs(r_prev, W, base)
            qdp_l = to_limbs(qd * p, W + 1, base)
            if bemit:                            # (borrow, limb) pairs: chain in-context
                r_l = to_limbs(r, W, base)
                b = 0
                for k in range(W):
                    parts.append(str(b) + limb_char(r_l[k]))
                    b = 1 if num_l[k] - qdp_l[k] - b < 0 else 0
                pos += 2 * W
            else:
                if subpad:                       # borrow chain of num - qd*p
                    b = 0
                    for k in range(W):
                        ann_g.append((pos + k, "sborrow", b))
                        b = 1 if num_l[k] - qdp_l[k] - b < 0 else 0
                parts.append(limb_str(r, W, base))
                pos += W
        r_prev = r
    parts.append(REVMARK)
    pos += 1
    ans_limbs = to_limbs(answer, W, base)
    for k in range(W):
        ann_g.append((pos + k, "ans", ans_limbs[k]))
    parts.append(limb_str(answer, W, base) + NL)
    pos += W + 1

    prompt = prompt_str(x, y, p, W, base, subpad)
    text = prompt + "".join(parts)
    assert len(text) == len(prompt) + gen_len(W, scratch, cursor, subpad, stepidx,
                                              skiptriv, subnum, bemit, srt) + 1, \
        (len(text), len(prompt), gen_len(W, scratch, cursor, subpad, stepidx,
                                         skiptriv, subnum, bemit, srt))
    ann = [dict() for _ in range(len(text))]
    base_i = len(prompt)
    for gi, var, val in ann_g:
        ann[base_i + gi - 1][var] = int(val)    # NTP alignment
    return text, ann


def var_specs(W: int, base: int, scratch: bool = False, subpad: bool = False,
              srt: bool = False, mulonly: bool = False):
    if mulonly:                                  # tier 0: multiply only
        return [("ans", base), ("acarry", base)]
    specs = [("ans", base), ("qhat", 2 * base - 1 if srt else base),
             ("acarry", base)]
    if subpad:                                  # multiply-carry + subtract-borrow chains
        specs += [("spcarry", base), ("sborrow", 2)]
    if not scratch:                             # scratch emits remainders: no probe
        specs += [(f"rem{j}", base) for j in range(W)]
    return specs


def decode_answer(gen_chars: str, W: int, base: int, scratch: bool = False,
                  cursor: bool = False, subpad: bool = False,
                  stepidx: bool = False, skiptriv: bool = False,
                  subnum: bool = False, bemit: bool = False,
                  srt: bool = False, mulonly: bool = False) -> int:
    if mulonly:                                  # answer = 2W limbs after 'N'
        off = gen_len(W, mulonly=True) - 2 * W
        return parse_limbs(gen_chars[off:off + 2 * W], base)
    # the answer block after the R-mark is W plain limbs in every mode
    off = gen_len(W, scratch, cursor, subpad, stepidx, skiptriv, subnum, bemit,
                  srt) - W
    return parse_limbs(gen_chars[off:off + W], base)


if __name__ == "__main__":
    import random
    rng = random.Random(3)
    for base in (10, 100):
        for scratch, cursor, subpad, stepidx, skiptriv, subnum, bemit, srt in (
                (False, False, False, False, False, False, False, False),
                (True, False, False, False, False, False, False, False),
                (True, True, False, False, False, False, False, False),
                (True, True, True, False, False, False, False, False),
                (True, True, True, True, False, False, False, False),
                (True, True, True, True, True, False, False, False),
                (True, True, True, True, True, True, False, False),
                (True, True, True, True, True, True, True, False),
                (True, True, True, True, True, True, True, True)):
            if srt and base != 10:
                continue
            for _ in range(4000):
                W = rng.randint(1, 12)
                p = rng.randrange(max(2, base ** (W - 1)), base ** W)
                x, y = rng.randrange(p), rng.randrange(p)
                text, ann = build_example(x, y, p, W, base, scratch, cursor,
                                          subpad, stepidx, skiptriv, subnum, bemit,
                                          srt)
                plen = len(prompt_str(x, y, p, W, base, subpad))
                assert decode_answer(text[plen:], W, base, scratch, cursor,
                                     subpad, stepidx, skiptriv, subnum, bemit,
                                     srt) \
                    == (x * y) % p, (x, y, p, W, base, scratch, cursor, subpad,
                                     stepidx, skiptriv, subnum, bemit, srt)
                assert len(ann) == len(text)
            print(f"composed base={base} scratch={scratch} cursor={cursor} "
                  f"subpad={subpad} stepidx={stepidx} skiptriv={skiptriv} "
                  f"subnum={subnum} bemit={bemit} srt={srt}: 4000/4000 decode OK")
    print("\ntokens/example (prompt+gen+NL):")
    for tier, Wd in [(3, 5), (4, 10), (5, 20), (6, 39), (7, 78)]:
        for base, W in ((10, Wd), (100, (Wd + 1) // 2)):
            print(f"  tier {tier} base {base:>3} W={W:>3}  compact {3 * W + 4 + gen_len(W) + 1:>6}"
                  f"  scratch {3 * W + 4 + gen_len(W, True) + 1:>6}")