123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552
import * as ast from "../ast/index.ts";import { Loader } from "../loader/index.ts";import { parse, ParseResult } from "../parser/proto.ts";import { toPojoSet } from "../runtime/array.ts";import {  filterNodesByType,  filterNodesByTypes,  findNodeByType,} from "./ast-util.ts";import { unwrap } from "./comment.ts";import {  evalConstant,  evalIntLit,  evalSignedIntLit,  evalStrLit,} from "./eval-ast-constant.ts";import {  Description,  Enum,  Extend,  File,  Group,  Import,  Message,  MessageField,  Options,  RpcType,  Schema,  Service,  Type,} from "./model.ts";import { scalarValueTypes } from "../runtime/scalar.ts";import {  stringifyFullIdent,  stringifyOptionName,  stringifyType,} from "./stringify-ast-frag.ts";
export interface BuildConfig {  loader: Loader;  files: string[];}export async function build(config: BuildConfig): Promise<Schema> {  return connect(await extract(gather(config)));}
export interface FileInfo {  filePath: string;  parseResult: ParseResult;  file: File;}export async function extract(files: AsyncIterable<FileInfo>) {  const result: Schema = {    files: {},    types: {},    extends: {},    services: {},  };  for await (const { filePath, parseResult, file } of files) {    result.files[filePath] = file;    const typePath = file.package ? "." + file.package : "";    const statements = parseResult.ast.statements;    const services = iterServices(statements, typePath, filePath);    for (const [typePath, service] of services) {      file.servicePaths.push(typePath);      if (!(typePath in result.services)) result.services[typePath] = service;    }    const types = iterTypes(statements, typePath, filePath);    for (const [typePath, type] of types) {      file.typePaths.push(typePath);      if (!(typePath in result.types)) result.types[typePath] = type;    }    // TODO: extends  }  return result;}
export function connect(schema: Schema): Schema {  const importPathToFilePath: {    [importPath: string]: string /* filePath */;  } = {};  for (const [filePath, { importPath }] of Object.entries(schema.files)) {    if (importPath in importPathToFilePath) continue;    importPathToFilePath[importPath] = filePath;  }  for (const file of Object.values(schema.files)) {    for (const entry of file.imports) {      if (entry.importPath in importPathToFilePath) {        entry.filePath = importPathToFilePath[entry.importPath];      }    }  }  for (const [filePath, file] of Object.entries(schema.files)) {    const resolveTypePath = getResolveTypePathFn(schema, filePath);    for (const typePath of file.typePaths) {      const type = schema.types[typePath];      if (type.kind === "enum") continue;      for (const field of Object.values(type.fields)) {        if (field.kind === "map") {          const fieldKeyTypePath = resolveTypePath(            field.keyType,            typePath as `.${string}`,          );          const fieldValueTypePath = resolveTypePath(            field.valueType,            typePath as `.${string}`,          );          if (fieldKeyTypePath) field.keyTypePath = fieldKeyTypePath;          if (fieldValueTypePath) field.valueTypePath = fieldValueTypePath;        } else {          const fieldTypePath = resolveTypePath(            field.type,            typePath as `.${string}`,          );          if (fieldTypePath) field.typePath = fieldTypePath;        }      }    }    for (const servicePath of file.servicePaths) {      const service = schema.services[servicePath];      for (const rpc of Object.values(service.rpcs)) {        const reqTypePath = resolveTypePath(          rpc.reqType.type,          servicePath as `.${string}`,        );        const resTypePath = resolveTypePath(          rpc.resType.type,          servicePath as `.${string}`,        );        if (reqTypePath) rpc.reqType.typePath = reqTypePath;        if (resTypePath) rpc.resType.typePath = resTypePath;      }    }  }  return schema;}
export interface GatherConfig {  files: string[];  loader: Loader;}export async function* gather(  { files, loader }: GatherConfig,): AsyncGenerator<FileInfo> {  const queue = [...files];  const visited: { [importPath: string]: true } = {};  const loaded: { [filePath: string]: true } = {};  while (queue.length) {    const importPath = queue.pop()!;    if (visited[importPath]) continue;    visited[importPath] = true;    const loadResult = await loader.load(importPath);    if (!loadResult) continue;    if (loaded[loadResult.absolutePath]) continue;    loaded[loadResult.absolutePath] = true;    const parseResult = parse(loadResult.data);    const statements = parseResult.ast.statements;    const file: File = {      parseResult,      importPath,      syntax: getSyntax(statements),      package: getPackage(statements),      imports: getImports(statements),      options: getOptions(statements),      typePaths: [],      servicePaths: [],    };    yield { filePath: loadResult.absolutePath, parseResult, file };    queue.push(...file.imports.map(({ importPath }) => importPath));  }}
export function merge(older: Schema, newer: Schema): Schema {  const result = {    files: { ...older.files, ...newer.files },    types: { ...older.types, ...newer.types },    extends: { ...older.extends, ...newer.extends },    services: { ...older.services, ...newer.services },  };  { // types    const olderTypePaths = Object.keys(older.types);    const newerTypePaths = Object.keys(newer.types);    for (const typePath of intersect(olderTypePaths, newerTypePaths)) {      const olderType = older.types[typePath];      const newerType = newer.types[typePath];      result.types[typePath] = mergeType(olderType, newerType);    }  }  { // extends    const olderTypePaths = Object.keys(older.extends);    const newerTypePaths = Object.keys(newer.extends);    for (const typePath of intersect(olderTypePaths, newerTypePaths)) {      const olderExtends = older.extends[typePath];      const newerExtends = newer.extends[typePath];      result.extends[typePath] = mergeExtends(olderExtends, newerExtends);    }  }  { // services    const olderServicePaths = Object.keys(older.services);    const newerServicePaths = Object.keys(newer.services);    for (const servicePath of intersect(olderServicePaths, newerServicePaths)) {      const olderService = older.services[servicePath];      const newerService = newer.services[servicePath];      result.services[servicePath] = mergeService(olderService, newerService);    }  }  return result;}
function mergeType(older: Type, newer: Type): Type {  if (older.kind !== newer.kind) return newer;  if (older.kind === "message") {    const _newer = newer as Message;    const result = { ..._newer };    return result; // TODO  } else {    const _newer = newer as Enum;    const result = { ..._newer };    result.fields = { ...older.fields, ..._newer.fields };    return result;  }}
function mergeExtends(older: Extend[], newer: Extend[]): Extend[] {  return newer; // TODO}
function mergeService(older: Service, newer: Service): Service {  return newer; // TODO}
function intersect<T>(a: T[], b: T[]): T[] {  return a.filter((x) => b.includes(x));}
function getSyntax(statements: ast.Statement[]): File["syntax"] {  const syntaxStatement = findNodeByType(statements, "syntax" as const);  const syntax = syntaxStatement?.syntax.text;  return syntax === "proto3" ? "proto3" : "proto2";}
function getPackage(statements: ast.Statement[]): string {  const packageStatement = findNodeByType(statements, "package" as const);  if (!packageStatement) return "";  return stringifyFullIdent(packageStatement.fullIdent);}
function getImports(statements: ast.Statement[]): Import[] {  const importStatements = filterNodesByType(statements, "import" as const);  return importStatements.map((statement) => {    const kind = (statement.weakOrPublic?.text || "") as Import["kind"];    const importPath = evalStrLit(statement.strLit);    return {      kind,      importPath,    };  });}
function getOptions(nodes?: ast.Node[]): Options {  if (!nodes) return {};  const optionStatements = filterNodesByTypes(nodes, [    "option" as const,    "field-option" as const,  ]);  const result: Options = {};  for (const statement of optionStatements) {    const optionName = stringifyOptionName(statement.optionName);    const optionValue = evalConstant(statement.constant);    result[optionName] = optionValue;  }  return result;}
function* iterServices(  statements: ast.Statement[],  typePath: string,  filePath: string,): Generator<[string, Service]> {  const serviceStatements = filterNodesByType(statements, "service" as const);  for (const statement of serviceStatements) {    const serviceTypePath = typePath + "." + statement.serviceName.text;    const service: Service = {      filePath,      options: getOptions(statement.serviceBody.statements),      description: getDescription(statement),      rpcs: getRpcs(statement.serviceBody.statements),    };    yield [serviceTypePath, service];  }}
function getRpcs(statements: ast.Statement[]): Service["rpcs"] {  const rpcStatements = filterNodesByType(statements, "rpc" as const);  const rpcs: Service["rpcs"] = {};  for (const statement of rpcStatements) {    const options = statement.semiOrRpcBody.type === "rpc-body"      ? getOptions(statement.semiOrRpcBody.statements)      : {};    rpcs[statement.rpcName.text] = {      options,      description: getDescription(statement),      reqType: getRpcType(statement.reqType),      resType: getRpcType(statement.resType),    };  }  return rpcs;}
function getRpcType(rpcType: ast.RpcType): RpcType {  return {    stream: !!rpcType.stream,    type: stringifyType(rpcType.messageType),  };}
function* iterTypes(  statements: ast.Statement[],  typePath: string,  filePath: string,): Generator<[string, Type]> {  for (const statement of statements) {    if (statement.type === "enum") {      yield getEnum(statement, typePath, filePath);    } else if (statement.type === "message") {      const message = getMessage(statement, typePath, filePath);      yield message;      const messageBodyStatement = statement.messageBody.statements;      yield* iterTypes(messageBodyStatement, message[0], filePath);    }  }}
function getMessage(  statement: ast.Message,  typePath: string,  filePath: string,): [string, Message] {  const messageTypePath = typePath + "." + statement.messageName.text;  const statements = statement.messageBody.statements;  const message: Message = {    kind: "message",    filePath,    description: getDescription(statement),    ...getMessageBody(statements),  };  return [messageTypePath, message];}
type MessageBody = Omit<Message, "kind" | "filePath" | "description">;function getMessageBody(  statements: ast.Statement[],): MessageBody {  const fields: Message["fields"] = {};  for (const [fieldNumber, field] of iterMessageFields(statements)) {    fields[fieldNumber] = field;  }  const groups: Message["groups"] = {};  for (    const groupStatement of filterNodesByType(statements, "group" as const)  ) {    const groupName = groupStatement.groupName.text;    groups[groupName] = getGroup(groupStatement);  }  return {    options: getOptions(statements),    fields,    groups,    reservedFieldNumberRanges: [], // TODO    reservedFieldNames: [], // TODO    extensions: [], // TODO  };}
function* iterMessageFields(  statements: ast.Statement[],): Generator<[number, MessageField]> {  for (const statement of statements) {    if (statement.type === "field") {      const fieldNumber = evalIntLit(statement.fieldNumber);      const fieldBase = {        description: getDescription(statement),        name: statement.fieldName.text,        options: getOptions(statement.fieldOptions?.fieldOptionOrCommas),        type: stringifyType(statement.fieldType),      };      if (!statement.fieldLabel) {        yield [fieldNumber, { kind: "normal", ...fieldBase }];      } else {        const kind = statement.fieldLabel.text;        if (          kind === "required" ||          kind === "optional" ||          kind === "repeated"        ) {          yield [fieldNumber, { kind, ...fieldBase }];        }      }    } else if (statement.type === "oneof") {      yield* iterOneofFields(        statement.oneofBody.statements,        statement.oneofName.text,      );    } else if (statement.type === "map-field") {      yield [evalIntLit(statement.fieldNumber), {        kind: "map",        description: getDescription(statement),        name: statement.mapName.text,        options: getOptions(statement.fieldOptions?.fieldOptionOrCommas),        keyType: stringifyType(statement.keyType),        valueType: stringifyType(statement.valueType),      }];    }  }}
function* iterOneofFields(  statements: ast.OneofBodyStatement[],  oneof: string,): Generator<[number, MessageField]> {  const oneofStatements = filterNodesByType(statements, "oneof-field" as const);  for (const statement of oneofStatements) {    const fieldNumber = evalIntLit(statement.fieldNumber);    yield [fieldNumber, {      kind: "oneof",      description: getDescription(statement),      name: statement.fieldName.text,      options: getOptions(statement.fieldOptions?.fieldOptionOrCommas),      type: stringifyType(statement.fieldType),      oneof,    }];  }}
function getGroup(statement: ast.Group): Group {  const statements = statement.messageBody.statements;  const fields: Message["fields"] = {};  for (const [fieldNumber, field] of iterMessageFields(statements)) {    fields[fieldNumber] = field;  }  return {    kind: statement.groupLabel.text as Group["kind"],    description: getDescription(statement),    fieldNumber: evalIntLit(statement.fieldNumber),    ...getMessageBody(statements),  };}
function getEnum(  statement: ast.Enum,  typePath: string,  filePath: string,): [string, Enum] {  const enumTypePath = typePath + "." + statement.enumName.text;  const _enum: Enum = {    kind: "enum",    filePath,    options: getOptions(statement.enumBody.statements),    description: getDescription(statement),    fields: getEnumFields(statement.enumBody.statements),  };  return [enumTypePath, _enum];}
function getEnumFields(statements: ast.Statement[]): Enum["fields"] {  const fields: Enum["fields"] = {};  const enumFieldStatements = filterNodesByType(    statements,    "enum-field" as const,  );  for (const statement of enumFieldStatements) {    const fieldNumber = evalSignedIntLit(statement.fieldNumber);    fields[fieldNumber] = {      description: getDescription(statement),      name: statement.fieldName.text,      options: getOptions(statement.fieldOptions?.fieldOptionOrCommas),    };  }  return fields;}
function getDescription(statement: ast.StatementBase): Description {  return {    leading: unwrapComments(statement.leadingComments),    trailing: unwrapComments(statement.trailingComments),    leadingDetached: unwrapComments(statement.leadingDetachedComments),  };  function unwrapComments(commentGroups: ast.CommentGroup[]) {    return commentGroups.map(      (commentGroup) =>        commentGroup.comments.map(          (comment) => unwrap(comment.text),        ).join("\n"),    );  }}
export type ResolveTypePathFn = (  type: string,  scope: `.${string}`,) => string | undefined;export function getResolveTypePathFn(  schema: Schema,  filePath: string,): ResolveTypePathFn {  const visibleTypePaths = toPojoSet(getVisibleTypePaths(schema, filePath));  return function resolveTypePath(type, scope) {    if (type in scalarValueTypeSet) return "." + type;    if (type.startsWith(".")) return visibleTypePaths[type];    let currentScope = scope;    while (true) {      const typePath = currentScope + "." + type;      if (typePath in visibleTypePaths) return typePath;      const cut = currentScope.lastIndexOf(".");      if (cut < 0) return undefined;      currentScope = currentScope.slice(0, cut) as `.${string}`;    }  };}const scalarValueTypeSet = toPojoSet(scalarValueTypes);
function getVisibleTypePaths(schema: Schema, filePath: string): string[] {  const file = schema.files[filePath];  if (!file) return [];  return [    ...file.typePaths,    ...file.imports.map((entry) =>      entry.filePath ? getExportedTypePaths(schema, entry.filePath) : []    ).flat(1),  ];}
function getExportedTypePaths(schema: Schema, filePath: string): string[] {  const result: string[] = [];  const done: { [filePath: string]: true } = {};  const queue = [filePath];  while (queue.length) {    const filePath = queue.pop()!;    if (done[filePath]) continue;    done[filePath] = true;    const file = schema.files[filePath];    if (!file) continue;    result.push(...file.typePaths);    for (const entry of file.imports) {      if (entry.kind !== "public") continue;      if (!entry.filePath) continue;      queue.push(entry.filePath);    }  }  return result;}