Cockatrice/webclient/buf.gen.plugin.mjs

// @ts-check
/**
 * Custom protoc-gen-es sibling plugin. Emits `src/generated/index.ts`, a
 * single rollup that re-exports every generated `_pb` module and adds
 * `MessageInitShape<T>` param aliases for every `Command_*` message.
 *
 * Wired into `buf.gen.yaml` as a second local plugin. Runs with the same
 * descriptor set protoc-gen-es consumes, so output always tracks the protos.
 */
import { createEcmaScriptPlugin, runNodeJs } from '@bufbuild/protoplugin';

const HEADER = [
  '// @generated by protoc-gen-data. DO NOT EDIT.',
  '// Rollup of all proto modules + MessageInitShape param aliases for every Command_*,',
  '// plus type maps for Response/Event extensions grouped by scope.',
  '/* eslint-disable */',
  '',
  '',
].join('\n');

const inner = createEcmaScriptPlugin({
  name: 'protoc-gen-data',
  version: 'v0.1.0',
  generateTs(schema) {
    const f = schema.generateFile('index.ts');

    const MessageInitShape = f.import('MessageInitShape', '@bufbuild/protobuf', true);
    const MessageType = f.import('Message', '@bufbuild/protobuf', true);
    const GenExtensionType = f.import('GenExtension', '@bufbuild/protobuf/codegenv2', true);

    const sortedFiles = [...schema.files].sort((a, b) => a.name.localeCompare(b.name));

    for (const file of sortedFiles) {
      f.print('export * from ', f.string(`./proto/${file.name}_pb`), ';');
    }
    f.print();

    const commandMessages = [];
    for (const file of sortedFiles) {
      for (const msg of file.messages) {
        if (msg.name.startsWith('Command_')) {
          commandMessages.push(msg);
        }
      }
    }
    commandMessages.sort((a, b) => a.name.localeCompare(b.name));

    // importSchema() resolves paths relative to this plugin's `out` dir, which
    // yields `./<name>_pb` — but the _pb files live under ./proto/ (protoc-gen-es's
    // out). Build the import path explicitly so it points inside the proto subdir.
    for (const msg of commandMessages) {
      const alias = msg.name.slice('Command_'.length) + 'Params';
      const schemaName = `${msg.name}Schema`;
      const schemaSym = f.import(schemaName, `./proto/${msg.file.name}_pb`, true);
      f.print('export type ', alias, ' = ', MessageInitShape, '<typeof ', schemaSym, '>;');
    }
    f.print();

    // ── Type maps for Response/Event extensions, grouped by extendee ────────
    //
    // Scans all messages for nested `extend` declarations and groups them by
    // which message they extend (Response, SessionEvent, RoomEvent, GameEvent).
    // Emits one `interface *Map { TypeName: TypeName; ... }` per scope.

    /** @type {Map<string, import('@bufbuild/protobuf').DescMessage[]>} */
    const extendeeGroups = new Map();

    for (const file of sortedFiles) {
      for (const msg of file.messages) {
        for (const ext of msg.nestedExtensions) {
          const target = ext.extendee.name;
          const group = extendeeGroups.get(target);
          if (group) {
            group.push(msg);
          } else {
            extendeeGroups.set(target, [msg]);
          }
        }
      }
    }

    /** @type {[string, string, import('@bufbuild/protobuf').DescMessage | null][]} */
    const maps = [
      ['ResponseMap', 'Response', null],
      ['SessionEventMap', 'SessionEvent', null],
      ['RoomEventMap', 'RoomEvent', null],
      ['GameEventMap', 'GameEvent', null],
    ];

    // Resolve the base extendee message for maps that need the base type included
    for (const file of sortedFiles) {
      for (const msg of file.messages) {
        for (const entry of maps) {
          if (msg.name === entry[1]) {
            entry[2] = msg;
          }
        }
      }
    }

    for (const [mapName, extendeeName, baseMsg] of maps) {
      const msgs = (extendeeGroups.get(extendeeName) || []).slice();
      msgs.sort((a, b) => a.name.localeCompare(b.name));

      if (msgs.length === 0 && !baseMsg) continue;

      f.print('export interface ', mapName, ' {');

      // Include the base extendee type itself (e.g. Response in ResponseMap)
      if (baseMsg) {
        const sym = f.import(baseMsg.name, `./proto/${baseMsg.file.name}_pb`, true);
        f.print('  ', baseMsg.name, ': ', sym, ';');
      }

      for (const msg of msgs) {
        const sym = f.import(msg.name, `./proto/${msg.file.name}_pb`, true);
        f.print('  ', msg.name, ': ', sym, ';');
      }

      f.print('}');
      f.print();
    }

    // Generic extension registry infrastructure. Consolidates the three
    // near-duplicate registry types and helpers that used to live in
    // src/websocket/services/protobuf-types.ts into one generic pair.
    // Specialised aliases (Session/Room/Game) still live in protobuf-types.ts
    // because GameExtensionRegistry needs GameEventMeta — a hand-written
    // domain type whose import would create a generated/ ↔ types/ cycle.
    f.print('export type RegistryEntry<V, T extends ', MessageType, ', M = unknown> = [');
    f.print('  ', GenExtensionType, '<T, V>,');
    f.print('  (value: V, meta: M) => void,');
    f.print('];');
    f.print();
    // Return type widens V to `unknown` so the heterogeneous entries that
    // callers build can be stored in a homogeneous `RegistryEntry<unknown, T, M>[]`
    // array. This is the actual value-add over a bare tuple literal.
    f.print('export function makeEntry<T extends ', MessageType, ', V, M = unknown>(');
    f.print('  ext: ', GenExtensionType, '<T, V>,');
    f.print('  handler: (value: V, meta: M) => void,');
    f.print('): RegistryEntry<unknown, T, M> {');
    f.print('  return [ext, handler] as unknown as RegistryEntry<unknown, T, M>;');
    f.print('}');
  },
});

// Skip f.preamble() above and inject a custom rollup-aware header here instead —
// preamble() would write "@generated from file X.proto" which is misleading for
// a rollup file built from every input proto.
/** @type {import('@bufbuild/protoplugin').Plugin} */
const plugin = {
  name: inner.name,
  version: inner.version,
  run(request) {
    const response = inner.run(request);
    for (const file of response.file) {
      if (file.name === 'index.ts' && typeof file.content === 'string') {
        file.content = HEADER + file.content;
      }
    }
    return response;
  },
};

runNodeJs(plugin);