refactor(ts): moves type patching on message type level

Author: stalniy

ts/script/fix-ts-proto-generated-types.ts

@@ -15,6 +15,13 @@ const helperTypeRegex = new RegExp(
 );
 
 const ROOT_DIR = resolvePath(import.meta.dirname, "..", "src");
+
+const typesToPatch = new Set<string>();
+for await (const patchFile of fs.glob(`${ROOT_DIR}/generated/patches/*CustomTypePatches.ts`)) {
+  const { patches } = await import(patchFile);
+  Object.keys(patches).forEach((key) => typesToPatch.add(key));
+}
+
 for await (const path of fs.glob(`${ROOT_DIR}/generated/protos/**/*.ts`)) {
   const source = await fs.readFile(path, "utf8");
   let newSource = source;
@@ -23,6 +30,8 @@ for await (const path of fs.glob(`${ROOT_DIR}/generated/protos/**/*.ts`)) {
   newSource = newSource.replace(/^\s*create\(base\?:\s*DeepPartial<\w+>\):\s*\w+\s*\{\s*return\s*\w+\.fromPartial\(base \?\? \{\}\);\s*\},?\n?/gm, "");
   newSource = injectOwnHelpers(newSource, path);
 
+  newSource = applyPatching(newSource, path, typesToPatch);
+
   if (newSource !== source) {
     await fs.writeFile(path, newSource);
   }
@@ -50,3 +59,27 @@ function injectOwnHelpers(source: string, path: string) {
 
   return importHelpers + importTypeHelpers + source;
 }
+
+function applyPatching(source: string, filePath: string, typesToPatch: Set<string>) {
+  const imports = new Set<string>();
+  const exports: string[] = [];
+
+  const newSource = source.replace(
+    /^export const (\w+)(:\s*MessageFns<[^>]+,\s*["']([^"']+)["']>\s*=)/gm,
+    (match, symbolName, typeAnnotation, fullName) => {
+      if (!typesToPatch.has(fullName)) return match;
+
+      const namespace = fullName.split(".")[0];
+      const prefix = namespace === "akash" ? "node" : namespace;
+      const importPath = relativePath(filePath, `${ROOT_DIR}/generated/protos/patches/${prefix}PatchMessage.ts`);
+      imports.add(`import { patched } from "${importPath}";`);
+      exports.push(`export const ${symbolName} = patched(_${symbolName});`);
+
+      return `const _${symbolName}${typeAnnotation}`;
+    },
+  );
+
+  if (!exports.length) return source;
+
+  return Array.from(imports).join("\n") + "\n" + newSource + "\n" + exports.join("\n") + "\n";
+}

ts/script/protoc-gen-customtype-patches.ts

@@ -12,22 +12,48 @@ import { basename, normalize as normalizePath } from "path";
 
 import { findPathsToCustomField, getCustomType } from "../src/encoding/customTypes/utils.ts";
 
-runNodeJs(
-  createEcmaScriptPlugin({
-    name: "protoc-gen-customtype-patches",
-    version: "v1",
-    generateTs,
-  }),
-);
+export interface PluginOptions {
+  /**
+   * if true, we will patch the whole tree of the message type, starting from the custom field type and up to the root
+   * in case of patching ts-proto generated types which has self-references, we need to patch only leaf level
+   * @default false
+   */
+  patchWholeTree: boolean;
+}
+
+runNodeJs(createEcmaScriptPlugin<PluginOptions>({ name: "protoc-gen-customtype-patches", version: "v1", parseOptions, generateTs }));
 
 const PROTO_PATH = "../protos";
-function generateTs(schema: Schema): void {
+
+function parseOptions(rawOptions: Array<{
+  key: string;
+  value: string;
+}>): PluginOptions {
+  const options: PluginOptions = {
+    patchWholeTree: false,
+  };
+
+  for (const { key, value } of rawOptions) {
+    if (key === "patch_whole_tree") {
+      options.patchWholeTree = value === "true";
+    }
+  }
+
+  return options;
+}
+
+function generateTs(schema: Schema<PluginOptions>): void {
   const allPaths: DescField[][] = [];
 
   schema.files.forEach((file) => {
     file.messages.forEach((message) => {
       const paths = findPathsToCustomField(message, () => true);
-      allPaths.push(...paths);
+      if (schema.options.patchWholeTree) {
+        allPaths.push(...paths);
+      } else {
+        const leaves = paths.map((path) => path.slice(-1));
+        allPaths.push(...leaves);
+      }
     });
   });
   if (!allPaths.length) {
@@ -100,6 +126,24 @@ function generateTs(schema: Schema): void {
   patchesFile.print(`const p = {\n${indent(patches.join(",\n"))}\n};\n`);
   patchesFile.print(`export const patches = p;`);
 
+  const patchesTypeFileName = fileName.replace("CustomTypePatches", "PatchMessage");
+  const patchTypeFile = schema.generateFile(patchesTypeFileName);
+  patchTypeFile.print(`import { patches } from "./${fileName}";`);
+  patchTypeFile.print(`import type { MessageDesc } from "../../sdk/client/types.ts";`);
+  patchTypeFile.print(`export const patched = <T extends MessageDesc>(messageDesc: T): T => {`);
+  patchTypeFile.print(`  const patchMessage = patches[messageDesc.$type as keyof typeof patches] as any;`);
+  patchTypeFile.print(`  if (!patchMessage) return messageDesc;`);
+  patchTypeFile.print(`  return {`);
+  patchTypeFile.print(`    ...messageDesc,`);
+  patchTypeFile.print(`    encode(message, writer) {`);
+  patchTypeFile.print(`      return messageDesc.encode(patchMessage(message, 'encode'), writer);`);
+  patchTypeFile.print(`    },`);
+  patchTypeFile.print(`    decode(input, length) {`);
+  patchTypeFile.print(`      return patchMessage(messageDesc.decode(input, length), 'decode');`);
+  patchTypeFile.print(`    },`);
+  patchTypeFile.print(`  };`);
+  patchTypeFile.print(`};`);
+
   const testsFile = schema.generateFile(fileName.replace(/\.ts$/, ".spec.ts"));
   generateTests(basename(fileName), testsFile, messageToCustomFields);
 }
@@ -185,13 +229,13 @@ function generateTests(fileName: string, testsFile: GeneratedFile, messageToCust
   testsFile.print(`import { expect, describe, it } from "@jest/globals";`);
   testsFile.print(`import { patches } from "./${basename(fileName)}";`);
   testsFile.print(`import { generateMessage, type MessageSchema } from "@test/helpers/generateMessage";`);
-  testsFile.print(`import type { TypePatches } from "../../sdk/client/applyPatches.ts";`);
+  testsFile.print(`import type { TypePatches } from "../../sdk/client/types.ts";`);
   testsFile.print("");
   testsFile.print(`const messageTypes: Record<string, MessageSchema> = {`);
   for (const [message, fields] of messageToCustomFields.entries()) {
     testsFile.print(`  "${message.typeName}": {`);
     testsFile.print(`    type: `, testsFile.import(message.name, `${PROTO_PATH}/${message.file.name}.ts`), `,`);
-    testsFile.print(`    fields: [`, ...Array.from(fields, f => serializeField(f, testsFile)), `],`);
+    testsFile.print(`    fields: [`, ...Array.from(fields, (f) => serializeField(f, testsFile)), `],`);
     testsFile.print(`  },`);
   }
   testsFile.print(`};`);
@@ -236,7 +280,7 @@ function serializeField(f: DescField, file: GeneratedFile): Printable {
     field.push(`scalarType: ${f.scalar},`);
   }
   if (f.fieldKind === "enum") {
-    field.push(`enum: `, JSON.stringify(f.enum.values.map(v => v.localName)), `,`);
+    field.push(`enum: `, JSON.stringify(f.enum.values.map((v) => v.localName)), `,`);
   }
   if (getCustomType(f)) {
     field.push(`customType: "${getCustomType(f)!.shortName}",`);
@@ -246,10 +290,10 @@ function serializeField(f: DescField, file: GeneratedFile): Printable {
   }
   if (f.message) {
     field.push(`message: {fields: [`,
-      ...f.message.fields.map(nf => serializeField(nf, file)),
+      ...f.message.fields.map((nf) => serializeField(nf, file)),
       `],`,
       `type: `, file.import(f.message.name, `${PROTO_PATH}/${f.message.file.name}.ts`),
-      `},`
+      `},`,
     );
   }
   field.push(`},`);

ts/script/protoc-gen-sdk-object.ts

@@ -90,17 +90,17 @@ function generateTs(schema: Schema): void {
     f.export("const", "serviceLoader"),
     `= `,
     f.import("createServiceLoader", `../sdk/client/createServiceLoader${importExtension}`),
-    `([\n${indent(servicesLoaderDefs.join(",\n"))}\n] as const);`
+    `([\n${indent(servicesLoaderDefs.join(",\n"))}\n] as const);`,
   );
 
   const factoryArgs = hasMsgService
     ? `queryTransport: Transport, txTransport: Transport`
     : `transport: Transport`;
   f.print(
     f.export("function", "createSDK"),
-    `(${factoryArgs}, options?: `, f.import("SDKOptions", `../sdk/types${importExtension}`), `) {\n`,
-    `  const getClient = createClientFactory<CallOptions>(${hasMsgService ? "queryTransport" : "transport"}, options?.clientOptions);\n`,
-    (hasMsgService ? `  const getMsgClient = createClientFactory<TxCallOptions>(txTransport, options?.clientOptions);\n` : ""),
+    `(${factoryArgs}) {\n`,
+    `  const getClient = createClientFactory<CallOptions>(${hasMsgService ? "queryTransport" : "transport"});\n`,
+    (hasMsgService ? `  const getMsgClient = createClientFactory<TxCallOptions>(txTransport);\n` : ""),
     `  return ${indent(stringifyObject(sdkDefs)).trim()};\n`,
     `}`,
   );
@@ -227,7 +227,7 @@ function findExtension(schema: Schema, typeName: string) {
   return extensionsCache[typeName];
 }
 
-const serviceFiles: Record<string, GeneratedFile>  = {};
+const serviceFiles: Record<string, GeneratedFile> = {};
 function generateServiceDefs(service: DescService, schema: Schema) {
   const importExtension = schema.options.importExtension ? `.${schema.options.importExtension}` : "";
   const serviceFilePath = `${service.file.name}_akash.ts`;
@@ -243,10 +243,10 @@ function generateServiceDefs(service: DescService, schema: Schema) {
   service.methods.forEach((method) => {
     file.print(`    ${method.localName}: {`);
     file.print(`      name: "${method.name}",`);
-    if (method.methodKind !== "unary")  file.print(`      kind: "${method.methodKind}",`);
+    if (method.methodKind !== "unary") file.print(`      kind: "${method.methodKind}",`);
     if (httpExtension && hasOption(method, httpExtension)) {
       const httpOption = getOption(method, httpExtension) as {
-        pattern: { case: string, value: string };
+        pattern: { case: string; value: string };
       };
       if (httpOption.pattern.case !== "get") file.print(`      httpMethod: "${httpOption.pattern.case}",`);
 

ts/script/protoc-gen-type-index-files.ts

@@ -1,13 +1,12 @@
 #!/usr/bin/env -S node --experimental-strip-types --no-warnings
 
-import { type DescEnum, type DescField, type DescMessage } from "@bufbuild/protobuf";
+import { type DescEnum, type DescMessage } from "@bufbuild/protobuf";
 import {
   createEcmaScriptPlugin,
   type GeneratedFile,
   runNodeJs,
-  type Schema
+  type Schema,
 } from "@bufbuild/protoplugin";
-import { findPathsToCustomField } from "../src/encoding/customTypes/utils.ts";
 
 runNodeJs(
   createEcmaScriptPlugin({
@@ -18,92 +17,41 @@ runNodeJs(
 );
 
 function generateTs(schema: Schema): void {
-  const allCustomTypeFieldPaths: DescField[][] = [];
-
-  schema.files.forEach((file) => {
-    file.messages.forEach((message) => {
-      const paths = findPathsToCustomField(message, () => true);
-      allCustomTypeFieldPaths.push(...paths);
-    });
-  });
-
-  const typesNamesToPatch = new Set<string>();
-  allCustomTypeFieldPaths.forEach((path) => {
-    path.forEach((field) => {
-      typesNamesToPatch.add(field.parent.typeName);
-    });
-  });
   const protoSource = process.env.PROTO_SOURCE;
   if (!protoSource) {
     throw new Error("PROTO_SOURCE is required and should be set to 'node', 'provider', or 'cosmos'");
   }
-  const patchesFileName = `${protoSource}PatchMessage.ts`;
-
-  if (typesNamesToPatch.size > 0) {
-    const patchesFile = schema.generateFile(patchesFileName);
-    patchesFile.print(`import { patches } from "../patches/${process.env.PROTO_SOURCE}CustomTypePatches.ts";`);
-    patchesFile.print(`import type { MessageDesc } from "../../sdk/client/types.ts";`);
-    patchesFile.print(`export const patched = <T extends MessageDesc>(messageDesc: T): T => {`);
-    patchesFile.print(`  const patchMessage = patches[messageDesc.$type as keyof typeof patches] as any;`);
-    patchesFile.print(`  if (!patchMessage) return messageDesc;`);
-    patchesFile.print(`  return {`);
-    patchesFile.print(`    ...messageDesc,`);
-    patchesFile.print(`    encode(message, writer) {`);
-    patchesFile.print(`      return messageDesc.encode(patchMessage(message, 'encode'), writer);`);
-    patchesFile.print(`    },`);
-    patchesFile.print(`    decode(input, length) {`);
-    patchesFile.print(`      return patchMessage(messageDesc.decode(input, length), 'decode');`);
-    patchesFile.print(`    },`);
-    patchesFile.print(`  };`);
-    patchesFile.print(`};`);
-  }
 
   const indexFiles: Record<string, {
     file: GeneratedFile;
     symbols: Set<string>;
   }> = {};
-  const namespacePrefix = protoSource === 'provider' ? 'provider.' : '';
+  const namespacePrefix = protoSource === "provider" ? "provider." : "";
   schema.files.forEach((file) => {
-    const packageParts = file.proto.package.split('.');
+    const packageParts = file.proto.package.split(".");
     const namespace = namespacePrefix + packageParts[0];
     const version = packageParts.at(-1);
     const path = `index.${namespace}.${version}.ts`;
     indexFiles[path] ??= {
       file: schema.generateFile(path),
       symbols: new Set(),
     };
-    const {file: indexFile, symbols: fileSymbols} = indexFiles[path];
+    const { file: indexFile, symbols: fileSymbols } = indexFiles[path];
 
-    const typesToPatch: Array<{ exportedName: string; name: string }> = [];
     const typesToExport: Array<{ exportedName: string; name: string }> = [];
     for (const type of schema.typesInFile(file)) {
-      if (type.kind === 'service' || type.kind === 'extension') continue;
+      if (type.kind === "service" || type.kind === "extension") continue;
 
       const name = genName(type);
       const exportedName = fileSymbols.has(name) ? genUniqueName(type, fileSymbols) : name;
       fileSymbols.add(exportedName);
-
-      if (type.kind === "message" && typesNamesToPatch.has(type.typeName)) {
-        typesToPatch.push({ exportedName, name });
-      } else {
-        typesToExport.push({ exportedName, name });
-      }
+      typesToExport.push({ exportedName, name });
     }
 
     if (typesToExport.length > 0) {
-      const symbolsToExport = typesToExport.map(type => type.exportedName === type.name ? type.exportedName : `${type.name} as ${type.exportedName}`).join(", ");
+      const symbolsToExport = typesToExport.map((type) => type.exportedName === type.name ? type.exportedName : `${type.name} as ${type.exportedName}`).join(", ");
       indexFile.print(`export { ${symbolsToExport} } from "./${file.name}.ts";`);
     }
-
-    if (typesToPatch.length > 0) {
-      const symbolsToPatch = typesToPatch.map((type) => `${type.name} as _${type.exportedName}`).join(", ");
-      indexFile.print('');
-      indexFile.print(`import { ${symbolsToPatch} } from "./${file.name}.ts";`);
-      for (const type of typesToPatch) {
-        indexFile.print(`export const ${type.exportedName} = `, indexFile.import('patched', `./${patchesFileName}`),`(_${type.exportedName});`);
-        indexFile.print(`export type ${type.exportedName} = _${type.exportedName}`);
-      }
-    }
   });
 }
 
@@ -115,8 +63,8 @@ let uniqueNameCounter = 0;
 function genUniqueName(type: DescMessage | DescEnum, allSymbols: Set<string>, attempt = 0): string {
   const name = genName(type);
   if (allSymbols.has(name)) {
-    const packageParts = type.file.proto.package.split('.');
-    const prefix = packageParts.slice(-2 - attempt, -1).map(capitalize).join('_');
+    const packageParts = type.file.proto.package.split(".");
+    const prefix = packageParts.slice(-2 - attempt, -1).map(capitalize).join("_");
     let newName = `${prefix}_${name}`;
     if (newName === name) {
       newName = `${prefix}_${name}_${uniqueNameCounter++}`;

ts/src/generated/createCosmosSDK.ts

@@ -1,5 +1,4 @@
 import { createServiceLoader } from "../sdk/client/createServiceLoader.ts";
-import { SDKOptions } from "../sdk/types.ts";
 
 import type * as cosmos_app_v1alpha1_query from "./protos/cosmos/app/v1alpha1/query.ts";
 import type * as cosmos_auth_v1beta1_query from "./protos/cosmos/auth/v1beta1/query.ts";
@@ -111,9 +110,9 @@ export const serviceLoader= createServiceLoader([
   () => import("./protos/cosmos/upgrade/v1beta1/tx_akash.ts").then(m => m.Msg),
   () => import("./protos/cosmos/vesting/v1beta1/tx_akash.ts").then(m => m.Msg)
 ] as const);
-export function createSDK(queryTransport: Transport, txTransport: Transport, options?: SDKOptions) {
-  const getClient = createClientFactory<CallOptions>(queryTransport, options?.clientOptions);
-  const getMsgClient = createClientFactory<TxCallOptions>(txTransport, options?.clientOptions);
+export function createSDK(queryTransport: Transport, txTransport: Transport) {
+  const getClient = createClientFactory<CallOptions>(queryTransport);
+  const getMsgClient = createClientFactory<TxCallOptions>(txTransport);
   return {
     cosmos: {
       app: {

ts/src/generated/createIbc-goSDK.ts

@@ -1,5 +1,4 @@
 import { createServiceLoader } from "../sdk/client/createServiceLoader.ts";
-import { SDKOptions } from "../sdk/types.ts";
 
 import type * as ibc_applications_interchain_accounts_controller_v1_query from "./protos/ibc/applications/interchain_accounts/controller/v1/query.ts";
 import type * as ibc_applications_interchain_accounts_controller_v1_tx from "./protos/ibc/applications/interchain_accounts/controller/v1/tx.ts";
@@ -45,9 +44,9 @@ export const serviceLoader= createServiceLoader([
   () => import("./protos/ibc/lightclients/wasm/v1/query_akash.ts").then(m => m.Query),
   () => import("./protos/ibc/lightclients/wasm/v1/tx_akash.ts").then(m => m.Msg)
 ] as const);
-export function createSDK(queryTransport: Transport, txTransport: Transport, options?: SDKOptions) {
-  const getClient = createClientFactory<CallOptions>(queryTransport, options?.clientOptions);
-  const getMsgClient = createClientFactory<TxCallOptions>(txTransport, options?.clientOptions);
+export function createSDK(queryTransport: Transport, txTransport: Transport) {
+  const getClient = createClientFactory<CallOptions>(queryTransport);
+  const getMsgClient = createClientFactory<TxCallOptions>(txTransport);
   return {
     ibc: {
       applications: {

ts/src/generated/createNodeSDK.ts

@@ -1,5 +1,4 @@
 import { createServiceLoader } from "../sdk/client/createServiceLoader.ts";
-import { SDKOptions } from "../sdk/types.ts";
 
 import type * as akash_audit_v1_query from "./protos/akash/audit/v1/query.ts";
 import type * as akash_audit_v1_msg from "./protos/akash/audit/v1/msg.ts";
@@ -41,9 +40,9 @@ export const serviceLoader= createServiceLoader([
   () => import("./protos/akash/take/v1/query_akash.ts").then(m => m.Query),
   () => import("./protos/akash/take/v1/service_akash.ts").then(m => m.Msg)
 ] as const);
-export function createSDK(queryTransport: Transport, txTransport: Transport, options?: SDKOptions) {
-  const getClient = createClientFactory<CallOptions>(queryTransport, options?.clientOptions);
-  const getMsgClient = createClientFactory<TxCallOptions>(txTransport, options?.clientOptions);
+export function createSDK(queryTransport: Transport, txTransport: Transport) {
+  const getClient = createClientFactory<CallOptions>(queryTransport);
+  const getMsgClient = createClientFactory<TxCallOptions>(txTransport);
   return {
     akash: {
       audit: {