snippets/serializer/plugin/register.ts

import { compileObject, compileUnion } from './codegen.ts';
import type { AnySchema, ObjectSchema, UnionSchema } from './descriptors.ts';
import type { Reader, Writer } from './io.ts';

export interface Codec<T = unknown> {
  readonly id: number;
  readonly name: string;
  readonly encode: (w: Writer, v: T) => void;
  readonly decode: (r: Reader) => T;
}

type AnyCodec = Codec<any>;

const byName = new Map<string, AnyCodec>();
const byId = new Map<number, AnyCodec>();

function fnv1a16(s: string): number {
  let h = 0x811c9dc5;
  for (let i = 0; i < s.length; i++) {
    h ^= s.charCodeAt(i);
    h = Math.imul(h, 0x01000193);
  }
  return ((h >>> 16) ^ h) & 0xffff;
}

function collectRefDeps(schema: AnySchema, acc: Map<string, ObjectSchema>): void {
  switch (schema.kind) {
    case 'ref': {
      const target = schema.thunk();
      if (!acc.has(target.name)) {
        acc.set(target.name, target);
        for (const f of Object.values(target.fields)) collectRefDeps(f, acc);
      }
      return;
    }
    case 'object':
      for (const f of Object.values(schema.fields)) collectRefDeps(f, acc);
      return;
    case 'array':
    case 'optional':
      collectRefDeps(schema.elem, acc);
      return;
    case 'tuple':
      for (const e of schema.elems) collectRefDeps(e, acc);
      return;
    case 'union':
      for (const v of Object.values(schema.variants)) collectRefDeps(v, acc);
      return;
    default:
      return;
  }
}

function sanIdent(name: string): string {
  return name.replace(/[^A-Za-z0-9_]/g, '_');
}

/**
 * Compile a schema into a `Codec` and register it in the lookup tables. Used
 * internally by `type(...)` / `oneOf(...)`. Idempotent by schema name.
 */
export function register<T = unknown>(schema: ObjectSchema | UnionSchema): Codec<T> {
  const existing = byName.get(schema.name);
  if (existing) return existing as Codec<T>;

  // Only `ref` boundaries require a separately registered codec; nested
  // objects/unions/arrays/etc. are inlined into the generated function.
  const refTargets = new Map<string, ObjectSchema>();
  if (schema.kind === 'object') {
    for (const f of Object.values(schema.fields)) collectRefDeps(f, refTargets);
  } else {
    for (const v of Object.values(schema.variants)) collectRefDeps(v, refTargets);
  }
  for (const dep of refTargets.values()) {
    if (!byName.has(dep.name)) register(dep);
  }

  const cg = schema.kind === 'object' ? compileObject(schema) : compileUnion(schema);

  const depsObj: Record<string, Function> = {};
  for (const [local, info] of cg.deps) {
    const target = byName.get(info.targetName);
    if (!target) throw new Error(`Dep not registered: ${info.targetName}`);
    depsObj[local] = info.mode === 'enc' ? target.encode : target.decode;
  }

  const closureObj: Record<string, unknown> = {};
  for (const [k, v] of cg.closure) closureObj[k] = v;

  const encDeps: string[] = [];
  const decDeps: string[] = [];
  for (const [local, info] of cg.deps) {
    const line = `const ${local} = deps[${JSON.stringify(local)}];`;
    if (info.mode === 'enc') encDeps.push(line);
    else decDeps.push(line);
  }

  const closureDecls = [...cg.closure.keys()]
    .map((k) => `const ${k} = closure[${JSON.stringify(k)}];`)
    .join('\n');

  const fname = sanIdent(schema.name);

  const encSrc = `${encDeps.join('\n')}
${closureDecls}
return function encode_${fname}(w, o) {
  let pos = w.pos;
  let buf = w.buf;
  let view = w.view;
  ${cg.encodeBody}
  w.pos = pos;
};`;

  const decSrc = `${decDeps.join('\n')}
${closureDecls}
return function decode_${fname}(r) {
  let pos = r.pos;
  const buf = r.buf;
  const view = r.view;
  ${cg.decodeBody}
};`;

  let encFn: (w: Writer, v: T) => void;
  let decFn: (r: Reader) => T;
  try {
    encFn = new Function('deps', 'closure', encSrc)(depsObj, closureObj);
    decFn = new Function('deps', 'closure', decSrc)(depsObj, closureObj);
  } catch (e) {
    throw new Error(
      `Codegen failed for "${schema.name}": ${(e as Error).message}\n--- encode source ---\n${encSrc}\n--- decode source ---\n${decSrc}`,
    );
  }

  const id = fnv1a16(schema.name);
  const idExisting = byId.get(id);
  if (idExisting) {
    throw new Error(
      `Schema ID collision: "${schema.name}" and "${idExisting.name}" both hash to 0x${id.toString(16)}`,
    );
  }

  const codec: Codec<T> = Object.freeze({
    id,
    name: schema.name,
    encode: encFn,
    decode: decFn,
  });

  byName.set(schema.name, codec);
  byId.set(id, codec);
  return codec;
}

/** Reset the global codec registry. Test helper. */
export function clearRegistry(): void {
  byName.clear();
  byId.clear();
}

/**
 * Internal: AOT-generated codecs use this to inject themselves into the
 * runtime registry on module load. Not meant for user code.
 */
export function __registerPrecompiled<T = unknown>(codec: Codec<T>): Codec<T> {
  const existing = byName.get(codec.name);
  if (existing) return existing as Codec<T>;
  const idExisting = byId.get(codec.id);
  if (idExisting && idExisting.name !== codec.name) {
    throw new Error(
      `Schema ID collision: "${codec.name}" and "${idExisting.name}" both hash to 0x${codec.id.toString(16)}`,
    );
  }
  byName.set(codec.name, codec as AnyCodec);
  byId.set(codec.id, codec as AnyCodec);
  return codec;
}