scrabble-game/ui/src/lib/dict/dawg.ts

// In-memory reader for the dafsa DAWG binary format, ported byte-for-byte from
// github.com/iliadenisov/dafsa (bits.go / disk.go / dawg.go). It answers the one
// query the local move validator needs: given a word as alphabet-index bytes, is
// it stored in the dictionary and at what insertion index.
//
// The Go reader treats the file as a big-endian, MSB-first bit stream. Every
// field this reader touches on the lookup path is at most ~31 bits wide, so the
// bit reader accumulates into a plain JavaScript number (no BigInt) and stays
// exact. Faithfulness to the Go reader is enforced by dawg.parity.test.ts, which
// checks indexOf against the authoritative Go output across the whole dictionary.

// bitsLen returns the number of bits needed to represent x (Go math/bits.Len).
function bitsLen(x: number): number {
  return x === 0 ? 0 : 32 - Math.clz32(x);
}

/**
 * Dawg is a read-only view over a serialized dafsa dictionary held in memory.
 * Construct it from the raw bytes of a `.dawg` file, then call {@link indexOf}.
 */
export class Dawg {
  private readonly bytes: Uint8Array;
  private p = 0; // current position, in bits

  private readonly cbits: number;
  private readonly abits: number;
  private readonly numEdges: number;
  private readonly wbits: number;
  private readonly firstNodeOffset: number;
  private readonly hasEmptyWord: boolean;

  /** Number of words stored (dafsa NumAdded). */
  readonly numAdded: number;
  /** Number of graph nodes (dafsa NumNodes). */
  readonly numNodes: number;

  // Scratch for the last edge resolved by getEdge, mirroring dafsa's edgeEnd +
  // final flag. Reused to keep the lookup path allocation-free.
  private eNode = 0;
  private eCount = 0;
  private eFinal = false;

  constructor(bytes: Uint8Array) {
    this.bytes = bytes;

    // Reject anything that is not a serialized dawg — e.g. an HTML error page or a
    // truncated download routed here by mistake. The 32-bit big-endian header size is
    // the total byte length; if it does not match, the blob is not a dawg. Without
    // this guard a non-dawg blob would parse into a bogus reader that silently reports
    // every word as missing (so the caller must throw here to fall back to the network).
    const declaredSize = ((bytes[0] << 24) | (bytes[1] << 16) | (bytes[2] << 8) | bytes[3]) >>> 0;
    if (declaredSize !== bytes.length) {
      throw new Error(`dawg: not a dawg blob (size header ${declaredSize} != ${bytes.length} bytes)`);
    }

    // Header: 32-bit size (skipped — the whole file is already in memory), then
    // cbits, abits, the language code string, and the word/node/edge counts.
    this.p = 32;
    this.cbits = this.readBits(8);
    this.abits = this.readBits(8);
    this.skipString(); // language code — not needed to walk index bytes

    const numAdded = this.readUnsigned();
    const numNodes = this.readUnsigned();
    this.numEdges = this.readUnsigned();
    this.firstNodeOffset = this.p;
    this.hasEmptyWord = this.readBits(1) === 1;

    this.numAdded = numAdded;
    this.numNodes = numNodes;
    this.wbits = bitsLen(numAdded);
  }

  /**
   * indexOf returns the insertion index of the given word (as alphabet-index
   * bytes), or -1 if the word was never added. Mirrors dafsa IndexOfB.
   */
  indexOf(word: ArrayLike<number>): number {
    let skipped = 0;
    let node = 0; // rootNode
    let final = this.hasEmptyWord;
    for (let i = 0; i < word.length; i++) {
      if (!this.getEdge(node, word[i])) {
        return -1;
      }
      node = this.eNode;
      final = this.eFinal;
      skipped += this.eCount;
    }
    return final ? skipped : -1;
  }

  /** has reports whether the word (as alphabet-index bytes) is in the dictionary. */
  has(word: ArrayLike<number>): boolean {
    return this.indexOf(word) >= 0;
  }

  // --- Step-by-step traversal (the move generator's primitive) ---------------
  //
  // A `Node` is a bit offset into the graph; 0 denotes the root (which resolves
  // to firstNodeOffset). These mirror dafsa's traverse.go Cursor (Root/Final/
  // Next/Arcs) over the same bitstream this reader already decodes, so the ported
  // generator can drive the automaton one transition at a time. Single-threaded
  // JS shares this reader's position across calls; every method re-seeks to its
  // node on entry, and arcs brackets the callback with a save/restore, so nested
  // use during a walk is safe. Mirrors dafsa (*Cursor).

  /** root returns the start state of the automaton. */
  root(): number {
    return 0;
  }

  /** final reports whether node is an accepting state (a stored word ends there). */
  final(node: number): boolean {
    if (this.numEdges <= 0) {
      return this.hasEmptyWord && node === 0;
    }
    this.p = node === 0 ? this.firstNodeOffset : node;
    return this.readBits(1) === 1;
  }

  /**
   * next follows the edge labelled ch (an alphabet index) from node, returning the
   * destination node, or -1 when no such edge exists.
   */
  next(node: number, ch: number): number {
    return this.getEdge(node, ch) ? this.eNode : -1;
  }

  /**
   * arcs calls fn for each out-edge of node in ascending label order, passing the
   * edge's label, its destination node and whether that destination is accepting.
   * It stops early if fn returns false. Mirrors dafsa (*Cursor).Arcs.
   */
  arcs(node: number, fn: (label: number, dest: number, final: boolean) => boolean): void {
    if (this.numEdges <= 0) {
      return;
    }
    this.p = node === 0 ? this.firstNodeOffset : node;
    this.readBits(1); // node final flag — not needed here
    const fallthrough = this.readBits(1);

    if (fallthrough === 1) {
      const label = this.readBits(this.cbits);
      // The reader now sits at the destination node, whose first bit is its final flag.
      const dest = this.p;
      const final = this.readBits(1) === 1;
      fn(label, dest, final);
      return;
    }

    const nskiplen = bitsLen(this.wbits);
    let nskip = 0;
    let numEdges = 1;
    if (this.readBits(1) !== 1) {
      // not a single edge
      numEdges = this.readUnsigned();
      nskip = this.readBits(nskiplen);
    }

    for (let i = 0; i < numEdges; i++) {
      const label = this.readBits(this.cbits);
      if (i > 0) {
        this.readBits(nskip); // per-edge skip count, unused for traversal
      }
      const dest = this.readBits(this.abits);
      const resume = this.p;
      this.p = dest;
      const final = this.readBits(1) === 1;
      if (!fn(label, dest, final)) {
        return;
      }
      this.p = resume;
    }
  }

  // getEdge resolves the outgoing edge for ch from the node at the given bit
  // offset. On success it fills eNode/eCount/eFinal and returns true. Mirrors
  // dafsa (*dawg).getEdge.
  private getEdge(node: number, ch: number): boolean {
    if (this.numEdges <= 0) {
      return false;
    }
    const pos = node === 0 ? this.firstNodeOffset : node;
    this.p = pos;
    const nodeFinal = this.readBits(1);
    const fallthrough = this.readBits(1);

    if (fallthrough === 1) {
      const edgeCh = this.readBits(this.cbits);
      if (edgeCh === ch) {
        this.eCount = nodeFinal;
        this.eNode = this.p; // the fallthrough target is the physically next node
        this.eFinal = this.readBits(1) === 1;
        return true;
      }
      return false;
    }

    const singleEdge = this.readBits(1);
    let numEdges = 1;
    const nskiplen = bitsLen(this.wbits);
    let nskip = 0;
    if (singleEdge !== 1) {
      numEdges = this.readUnsigned();
      nskip = this.readBits(nskiplen);
    }

    const base = this.p; // bit offset of the first edge record
    const recordBits = this.cbits + nskip + this.abits;

    // Binary search over the edges, which are sorted ascending by character.
    let high = numEdges;
    let low = -1;
    while (high - low > 1) {
      const probe = (high + low) >> 1;
      // The first edge omits its (zero) skip field, so every later record is
      // shifted back by nskip bits.
      let seekTo = base + probe * recordBits;
      if (probe > 0) {
        seekTo -= nskip;
      }
      this.p = seekTo;
      const edgeCh = this.readBits(this.cbits);
      const cmp = edgeCh - ch;
      if (cmp === 0) {
        this.eCount = probe > 0 ? this.readBits(nskip) : nodeFinal;
        this.eNode = this.readBits(this.abits);
        this.p = this.eNode;
        this.eFinal = this.readBits(1) === 1;
        return true;
      } else if (cmp < 0) {
        low = probe;
      } else {
        high = probe;
      }
    }
    return false;
  }

  // readBits reads n (<= 32) bits MSB-first from the current position and
  // advances it. Mirrors dafsa (*bitSeeker).ReadBits for the widths on this path.
  private readBits(n: number): number {
    let result = 0;
    let p = this.p;
    const bytes = this.bytes;
    while (n > 0) {
      const byteIndex = p >>> 3;
      const bitInByte = p & 7;
      const avail = 8 - bitInByte;
      const take = avail < n ? avail : n;
      const shift = avail - take;
      const chunk = (bytes[byteIndex] >>> shift) & ((1 << take) - 1);
      result = result * (1 << take) + chunk;
      p += take;
      n -= take;
    }
    this.p = p;
    return result;
  }

  // readUnsigned reads a dafsa "7code" varint (7 payload bits per byte, high bit
  // continues). Mirrors dafsa readUnsigned.
  private readUnsigned(): number {
    let result = 0;
    for (;;) {
      const d = this.readBits(8);
      result = result * 128 + (d & 0x7f);
      if ((d & 0x80) === 0) {
        break;
      }
    }
    return result;
  }

  // skipString advances past a 7code-length-prefixed byte string.
  private skipString(): void {
    const n = this.readUnsigned();
    this.p += n * 8;
  }
}