decode-guard.js

// APE A2 WebGPU AR ship engine — guarded-decode constraint stack, pure JS.
//
// 1:1 port of the PT-gate-frozen Python modules (behavior is prereg-bound;
// no "improvements"):
//   scripts/ape_a2/a2_cont_decode_constrained.py  (DigitState, walk_piece,
//     digit mask modes, BAND_HI length band)          sha a1c687c0…
//   scripts/ape_a2/a2_name_copy_guard.py          (NameTables, NameState,
//     name_violations R1/R2a/R2b, partial viability)  sha f2f53741…
//   scripts/ape_a2/a2_ab_decode_nameguard.py      (tier loop 16/256/24000,
//     v3 empty-piece rule, stats accounting)          sha 3e3084c2…
//   scripts/ape_a2/a2_eval_decode.py              (toks, repeat_patterns)
//
// Python-semantics helpers are deliberate: str.casefold ≈ toLowerCase for
// this corpus (proven empirically by the P1 replay + selftest equivalence,
// not assumed), str.isspace ≠ JS \s (adds \x1c-\x1f \x85, excludes ),
// str.isalpha = \p{L}, str.isupper = \p{Lu}\p{Lt}, \d = \p{Nd}, len() =
// code points. No DOM/Node dependencies — runs in browser and Node.

export const EOS_ID = 2;
export const NEG = -1.0e30;
export const TIERS = [16, 256, 24000];
export const BAND_HI = 1.75;
export const MAX_NEW_TOKENS = 224;

// a2_eval_decode.PUNCT (repeat/toks stripping)
export const PUNCT_REPEAT = '“”"\'‘’.,!?;:…()[]«»-–—*';
// a2_name_copy_guard.PUNCT (name-token stripping; superset with CJK brackets)
export const PUNCT_NAME = '“”"\'‘’.,!?;:…()[]«»-–—*《》〈〉「」『』【】‹›';
export const SENT_BREAK = new Set([...'.!?…:;—–-"“”\'‘’«»()[]*《》〈〉「」『』【】‹›']);

// --- Python-semantics primitives -------------------------------------------

const PY_WS = /[\t\n\x0b\x0c\r\x1c\x1d\x1e\x1f \x85\u00a0\u1680\u2000-\u200a\u2028\u2029\u202f\u205f\u3000]/u;
const RE_LETTER = /\p{L}/u;
const RE_UPPER = /[\p{Lu}\p{Lt}]/u;
const RE_ND = /\p{Nd}/u;

export const isPySpace = (ch) => PY_WS.test(ch);
export const casefold = (s) => s.toLowerCase();
export const cpLen = (s) => [...s].length;

export function pyStrip(text) {
  const cps = [...text];
  let start = 0;
  let end = cps.length;
  while (start < end && isPySpace(cps[start])) start += 1;
  while (end > start && isPySpace(cps[end - 1])) end -= 1;
  return cps.slice(start, end).join('');
}

export function pySplit(text) {
  const out = [];
  let buf = '';
  for (const ch of text) {
    if (isPySpace(ch)) {
      if (buf) { out.push(buf); buf = ''; }
    } else buf += ch;
  }
  if (buf) out.push(buf);
  return out;
}

function stripChars(raw, punctSet) {
  const cps = [...raw];
  let start = 0;
  let end = cps.length;
  while (start < end && punctSet.has(cps[start])) start += 1;
  while (end > start && punctSet.has(cps[end - 1])) end -= 1;
  return [cps.slice(start, end).join(''), cps.slice(0, start).join(''), cps.slice(end).join('')];
}

const PUNCT_REPEAT_SET = new Set([...PUNCT_REPEAT]);
const PUNCT_NAME_SET = new Set([...PUNCT_NAME]);

// --- toks / repeat_patterns (a2_eval_decode) --------------------------------

export function toks(text) {
  const out = [];
  for (let token of pySplit(casefold(text.normalize('NFC')))) {
    token = stripChars(token, PUNCT_REPEAT_SET)[0];
    if (token) out.push(token);
  }
  return out;
}

const SEP = '\u0000';

export function repeatPatterns(text) {
  const tokens = toks(text);
  const patterns = new Set();
  for (let i = 0; i < tokens.length - 2; i += 1) {
    if (tokens[i] === tokens[i + 1] && tokens[i] === tokens[i + 2]) patterns.add(tokens[i]);
  }
  for (let size = 2; size <= 8; size += 1) {
    outer: for (let i = 0; i + 2 * size <= tokens.length; i += 1) {
      for (let j = 0; j < size; j += 1) {
        if (tokens[i + j] !== tokens[i + size + j]) continue outer;
      }
      patterns.add(tokens.slice(i, i + size).join(SEP));
    }
  }
  return patterns;
}

export function hasNewRepeat(text, draftPatterns) {
  for (const p of repeatPatterns(text)) if (!draftPatterns.has(p)) return true;
  return false;
}

// --- digit-copy guard (a2_cont_decode_constrained) --------------------------

export function digitRuns(text) {
  const runs = new Map();
  let buf = '';
  for (const ch of text) {
    if (RE_ND.test(ch)) buf += ch;
    else if (buf) { runs.set(buf, (runs.get(buf) ?? 0) + 1); buf = ''; }
  }
  if (buf) runs.set(buf, (runs.get(buf) ?? 0) + 1);
  return runs;
}

// Simulate appending `text`; returns [legal, consumedRuns[], newTrailing].
export function walkPiece(text, remaining, trailing) {
  const consumed = [];
  const scratch = new Map(remaining);
  for (const ch of text) {
    if (RE_ND.test(ch)) {
      trailing += ch;
      // Python iterates Counter KEYS here — entries decremented to zero
      // within this walk still count as viable prefixes (the close branch
      // rejects them via the <= 0 check). Do NOT filter on count > 0.
      let viable = false;
      for (const [run] of scratch) {
        if (run.startsWith(trailing)) { viable = true; break; }
      }
      if (!viable) return [false, consumed, trailing];
    } else if (trailing) {
      if ((scratch.get(trailing) ?? 0) <= 0) return [false, consumed, trailing];
      scratch.set(trailing, scratch.get(trailing) - 1);
      consumed.push(trailing);
      trailing = '';
    }
  }
  return [true, consumed, trailing];
}

export class DigitState {
  constructor(draft) {
    this.remaining = digitRuns(draft);
    this.trailing = '';
    this.consumed = new Map();
    this.broken = false;
  }

  key() {
    const parts = [];
    for (const [run, count] of [...this.remaining].sort((a, b) => (a[0] < b[0] ? -1 : a[0] > b[0] ? 1 : 0))) {
      for (let i = 0; i < count; i += 1) parts.push(run);
    }
    return `${parts.join(',')}|${this.trailing}`;
  }

  closable() {
    return !this.trailing || (this.remaining.get(this.trailing) ?? 0) > 0;
  }

  apply(text) {
    const [legal, consumed, trailing] = walkPiece(text, this.remaining, this.trailing);
    for (const run of consumed) {
      this.remaining.set(run, (this.remaining.get(run) ?? 0) - 1);
      this.consumed.set(run, (this.consumed.get(run) ?? 0) + 1);
    }
    for (const [run, count] of [...this.remaining]) if (count <= 0) this.remaining.delete(run);
    this.trailing = trailing;
    if (!legal) this.broken = true;
  }

  emitted() {
    const total = new Map(this.consumed);
    if (this.trailing) total.set(this.trailing, (total.get(this.trailing) ?? 0) + 1);
    return total;
  }
}

// Digit mask decision — mirrors masks_for (ban / allow / fallback), memoized
// by digit-state key. pieceTable: {pieceText, digitIds, bannedBytes}.
export function masksFor(state, pieceTable, memo) {
  const key = state.key();
  const hit = memo.get(key);
  if (hit) return hit;
  const allowed = [];
  for (const id of pieceTable.digitIds) {
    if (walkPiece(pieceTable.pieceText[id], state.remaining, state.trailing)[0]) allowed.push(id);
  }
  let entry;
  if (state.closable()) {
    const allowedSet = new Set(allowed);
    entry = { mode: 'ban', ids: pieceTable.digitIds.filter((t) => !allowedSet.has(t)).concat(pieceTable.bannedBytes) };
  } else if (allowed.length) {
    entry = { mode: 'allow', ids: allowed };
  } else {
    entry = { mode: 'fallback', ids: [] };
  }
  memo.set(key, entry);
  return entry;
}

// --- name copy-guard (a2_name_copy_guard) -----------------------------------

function splitNameToken(raw) {
  return stripChars(raw, PUNCT_NAME_SET);
}

function isCapital(core) {
  for (const ch of core) {
    if (RE_LETTER.test(ch)) return RE_UPPER.test(ch);
  }
  return false;
}

export function tokMeta(text) {
  const out = [];
  for (const raw of pySplit(text)) {
    const [core, lead, trail] = splitNameToken(raw);
    out.push({ raw, core, lead, trail, cap: isCapital(core) });
  }
  return out;
}

export function exemptAt(tokens, index) {
  if (tokens[index].lead) return true;
  let scan = index - 1;
  while (scan >= 0) {
    const token = tokens[scan];
    if (token.trail && [...token.trail].some((ch) => SENT_BREAK.has(ch))) return true;
    if (!token.core) {
      if ([...token.raw].some((ch) => SENT_BREAK.has(ch))) return true;
      scan -= 1;
      continue;
    }
    return false;
  }
  return true;
}

const PAIR = (a, b) => `${a}${SEP}${b}`;

export class NameTables {
  constructor(draft) {
    const tokens = tokMeta(draft);
    this.cores = new Set();
    for (const t of tokens) if (t.core) this.cores.add(t.core);
    this.coresList = [...this.cores].sort();
    this.capPairs = new Set();
    this.foldPairs = new Set();
    this.capNext = new Map();
    this.nextFolds = new Map();
    for (let i = 1; i < tokens.length; i += 1) {
      const prev = tokens[i - 1];
      const cur = tokens[i];
      if (prev.core && cur.core && !prev.trail && !cur.lead) {
        if (prev.cap && cur.cap) {
          this.capPairs.add(PAIR(prev.core, cur.core));
          if (!this.capNext.has(prev.core)) this.capNext.set(prev.core, new Set());
          this.capNext.get(prev.core).add(cur.core);
        }
        this.foldPairs.add(PAIR(casefold(prev.core), casefold(cur.core)));
        const pf = casefold(prev.core);
        if (!this.nextFolds.has(pf)) this.nextFolds.set(pf, new Set());
        this.nextFolds.get(pf).add(casefold(cur.core));
      }
    }
  }
}

// Post-hoc checker — the frozen specification (design-doc rules).
export function nameViolations(draft, candidate) {
  const tables = new NameTables(draft);
  const tokens = tokMeta(candidate);
  const hits = [];
  const meaningful = [];
  const idxMap = [];
  for (let i = 0; i < tokens.length; i += 1) {
    if (tokens[i].core) { meaningful.push(tokens[i]); idxMap.push(i); }
  }
  for (let m = 0; m < meaningful.length; m += 1) {
    const token = meaningful[m];
    const i = idxMap[m];
    const ex = exemptAt(tokens, i);
    const prev = m > 0 ? meaningful[m - 1] : null;
    const prevI = m > 0 ? idxMap[m - 1] : null;
    const adjacent = prev !== null && prevI === i - 1 && !prev.trail && !token.lead;
    const foldPair = prev ? PAIR(casefold(prev.core), casefold(token.core)) : null;
    if (token.cap && !ex) {
      const strict = tables.cores.has(token.core);
      const chain = adjacent && prev.cap && tables.foldPairs.has(foldPair);
      if (!strict && !chain) hits.push(['R1', token.core]);
    }
    if (adjacent && prev.cap && token.cap && !ex && !exemptAt(tokens, prevI)) {
      if (!tables.capPairs.has(PAIR(prev.core, token.core)) && !tables.foldPairs.has(foldPair)) {
        hits.push(['R2a', `${prev.core} ${token.core}`]);
      }
    }
    if (adjacent && prev.cap && !exemptAt(tokens, prevI)
        && casefold(prev.core) === casefold(token.core) && !tables.foldPairs.has(foldPair)) {
      hits.push(['R2b', `${prev.core} ${token.core}`]);
    }
  }
  return hits;
}

class Cursor {
  constructor() {
    this.buf = '';
    this.exemptNext = true;
    this.prevCore = null;
    this.prevCap = false;
    this.prevTrailEmpty = true;
    this.prevExempt = true;
    this.prevAdjacent = false;
  }

  clone() {
    const other = new Cursor();
    other.buf = this.buf;
    other.exemptNext = this.exemptNext;
    other.prevCore = this.prevCore;
    other.prevCap = this.prevCap;
    other.prevTrailEmpty = this.prevTrailEmpty;
    other.prevExempt = this.prevExempt;
    other.prevAdjacent = this.prevAdjacent;
    return other;
  }
}

export class NameState {
  constructor(draft) {
    this.tables = new NameTables(draft);
    this.cur = new Cursor();
    this.broken = false;
  }

  closeToken(cursor, raw) {
    const tables = this.tables;
    const [core, lead, trail] = splitNameToken(raw);
    if (!core) {
      if ([...raw].some((ch) => SENT_BREAK.has(ch))) cursor.exemptNext = true;
      cursor.prevAdjacent = false;
      return [];
    }
    const cap = isCapital(core);
    const exempt = Boolean(lead) || cursor.exemptNext;
    const viols = [];
    const adjacent = cursor.prevAdjacent && cursor.prevTrailEmpty && !lead;
    const foldPair = cursor.prevCore !== null ? PAIR(casefold(cursor.prevCore), casefold(core)) : null;
    if (cap && !exempt) {
      const strict = tables.cores.has(core);
      const chain = adjacent && cursor.prevCap && tables.foldPairs.has(foldPair);
      if (!strict && !chain) viols.push(['R1', core]);
    }
    if (adjacent && cursor.prevCap && cap && !exempt && !cursor.prevExempt) {
      if (!tables.capPairs.has(PAIR(cursor.prevCore, core)) && !tables.foldPairs.has(foldPair)) {
        viols.push(['R2a', `${cursor.prevCore} ${core}`]);
      }
    }
    if (adjacent && cursor.prevCap && !cursor.prevExempt && cursor.prevCore !== null
        && casefold(cursor.prevCore) === casefold(core) && !tables.foldPairs.has(foldPair)) {
      viols.push(['R2b', `${cursor.prevCore} ${core}`]);
    }
    cursor.prevCore = core;
    cursor.prevCap = cap;
    cursor.prevTrailEmpty = trail === '';
    cursor.prevExempt = exempt;
    cursor.prevAdjacent = true;
    cursor.exemptNext = [...trail].some((ch) => SENT_BREAK.has(ch));
    return viols;
  }

  feed(cursor, text) {
    const viols = [];
    for (const ch of text) {
      if (isPySpace(ch)) {
        if (cursor.buf) {
          viols.push(...this.closeToken(cursor, cursor.buf));
          cursor.buf = '';
        }
      } else cursor.buf += ch;
    }
    return viols;
  }

  partialViable(cursor) {
    const buf = cursor.buf;
    if (!buf) return true;
    const [core, lead, trail] = splitNameToken(buf);
    if (!core || !isCapital(core)) return true;
    if (lead || cursor.exemptNext) return true;
    if (trail) {
      // v2 addendum: trailing punctuation freezes the core — must be
      // close-legal NOW.
      const probe = cursor.clone();
      return this.closeToken(probe, buf).length === 0;
    }
    const tables = this.tables;
    const adjacent = cursor.prevAdjacent && cursor.prevTrailEmpty && !lead;
    const prevCapAdj = adjacent && cursor.prevCap && cursor.prevCore !== null;
    if (prevCapAdj && !cursor.prevExempt) {
      const prevCore = cursor.prevCore;
      const prevFold = casefold(prevCore);
      const foldCore = casefold(core);
      for (const nf of tables.nextFolds.get(prevFold) ?? []) {
        if (nf.startsWith(foldCore)) return true;
      }
      for (const w of tables.capNext.get(prevCore) ?? []) {
        if (w.startsWith(core) && casefold(w) !== prevFold) return true;
      }
      return false;
    }
    if (prevCapAdj) {
      const foldCore = casefold(core);
      const prevFold = casefold(cursor.prevCore);
      for (const nf of tables.nextFolds.get(prevFold) ?? []) {
        if (nf.startsWith(foldCore)) return true;
      }
    }
    for (const candidateCore of tables.coresList) {
      if (candidateCore.startsWith(core)) return true;
    }
    return false;
  }

  trial(text, closing = false) {
    if (this.broken) return true;
    const cursor = this.cur.clone();
    const viols = this.feed(cursor, text);
    if (closing && cursor.buf) {
      viols.push(...this.closeToken(cursor, cursor.buf));
      cursor.buf = '';
    }
    if (viols.length) return false;
    if (!closing && !this.partialViable(cursor)) return false;
    return true;
  }

  eosOk() {
    return this.trial('', true);
  }

  apply(text) {
    this.feed(this.cur, text);
  }
}

// Word-by-word NameState close events must agree with the checker (port of
// selftest_equivalence; throws on any disagreement).
export function selftestEquivalence(pairs) {
  let agreeZero = 0;
  let agreeHit = 0;
  pairs.forEach(([draft, candidate], index) => {
    const checkerHits = JSON.stringify(nameViolations(draft, candidate));
    const state = new NameState(draft);
    const cursor = state.cur;
    const incremental = state.feed(cursor, candidate);
    if (cursor.buf) {
      incremental.push(...state.closeToken(cursor, cursor.buf));
      cursor.buf = '';
    }
    if (JSON.stringify(incremental) !== checkerHits) {
      throw new Error(`selftest disagreement at pair ${index}: checker=${checkerHits} incremental=${JSON.stringify(incremental)}`);
    }
    if (checkerHits === '[]') agreeZero += 1;
    else agreeHit += 1;
  });
  return { pairsZero: agreeZero, pairsWithHits: agreeHit };
}

// --- guarded per-row decode step (a2_ab_decode_nameguard tier loop) ---------
//
// One GuardedRow per batch row. Per step the caller:
//   1. asks maskRequest() and applies the digit veto to the raw logits
//      ('fallback' => apply nothing; the row is marked broken internally),
//   2. computes rawTop (argmax of RAW logits), maskedTop (argmax of MASKED
//      logits — argmax, not topk[0]: tie ordering must match the reference)
//      and the masked-descending candidate tiers,
//   3. calls step(rawTop, maskedTop, getTier) where getTier(tierIndex)
//      returns {ids, validCount} — candidate ids in descending masked-logit
//      order, entries at position >= validCount being NEG-vetoed (torch.topk
//      order: vetoed entries always sort after every legal one).
// step() returns the chosen token id (EOS_ID when the row finishes) and
// mutates all guard state exactly like the frozen Python loop.
export class GuardedRow {
  constructor(draft, pieceTable, maskMemo, budget = undefined) {
    this.draft = draft;
    this.pieceTable = pieceTable;
    this.maskMemo = maskMemo;
    this.digit = new DigitState(draft);
    this.name = new NameState(draft);
    this.draftPatterns = repeatPatterns(draft);
    this.sim = '';
    this.outputs = [];
    this.finished = false;
    this.budget = budget === undefined ? Math.trunc(BAND_HI * cpLen(draft)) - 1 : budget;
    this.stats = {
      digitVetoSteps: 0,
      repeatVetoSteps: 0,
      nameVetoSteps: 0,
      nameTierEscalations: 0,
      digitFallback: false,
      nameBroken: false,
      lengthStops: 0,
    };
  }

  // Digit veto for this step; caller applies it BEFORE building tiers.
  maskRequest() {
    if (this.finished || this.digit.broken) return { mode: 'none', ids: [] };
    const entry = masksFor(this.digit, this.pieceTable, this.maskMemo);
    if (entry.mode === 'fallback') {
      this.stats.digitFallback = true;
      this.digit.broken = true;
      return { mode: 'fallback', ids: [] };
    }
    return entry;
  }

  step(rawTop, maskedTop, getTier) {
    if (this.finished) return EOS_ID;
    if (maskedTop !== rawTop) this.stats.digitVetoSteps += 1;

    let pick = null;
    let firstReject = null;
    for (let tierIndex = 0; tierIndex < TIERS.length; tierIndex += 1) {
      const { ids, validCount } = getTier(tierIndex);
      const limit = Math.min(ids.length, validCount);
      for (let i = 0; i < limit; i += 1) {
        const token = ids[i];
        if (token === EOS_ID) {
          if (this.name.broken || this.name.eosOk()) { pick = token; break; }
          if (firstReject === null) firstReject = 'name';
          continue;
        }
        const text = this.pieceTable.pieceText[token];
        if (!text) continue; // v3 addendum: empty-piece specials never escape veto pressure
        if (!(this.name.broken || this.name.trial(text))) {
          if (firstReject === null) firstReject = 'name';
          continue;
        }
        if (hasNewRepeat(this.sim + text, this.draftPatterns)) {
          if (firstReject === null) firstReject = 'repeat';
          continue;
        }
        pick = token;
        break;
      }
      if (pick !== null) break;
      if (tierIndex + 1 < TIERS.length) this.stats.nameTierEscalations += 1;
    }
    if (pick === null) {
      pick = maskedTop;
      this.name.broken = true;
      this.stats.nameBroken = true;
    }
    if (firstReject === 'name' && pick !== maskedTop) this.stats.nameVetoSteps += 1;
    else if (firstReject === 'repeat' && pick !== maskedTop) this.stats.repeatVetoSteps += 1;

    if (
      this.budget !== null
      && pick !== EOS_ID
      && cpLen(pyStrip(this.sim + this.pieceTable.pieceText[pick])) > this.budget
      && (this.digit.broken || this.digit.closable())
      && (this.name.broken || this.name.eosOk())
    ) {
      pick = EOS_ID;
      this.stats.lengthStops += 1;
    }
    if (pick === EOS_ID) {
      this.finished = true;
    } else {
      this.outputs.push(pick);
      const text = this.pieceTable.pieceText[pick];
      this.sim += text;
      if (!this.digit.broken) this.digit.apply(text);
      if (!this.name.broken) this.name.apply(text);
    }
    return pick;
  }
}