Add mutable registry (#899)

bufbuild · Jun 21, 2024 · 5e0b35c · 5e0b35c
1 parent 14c0018
commit 5e0b35c
Show file tree

Hide file tree

Showing 6 changed files with 75 additions and 37 deletions.
diff --git a/packages/bundle-size/README.md b/packages/bundle-size/README.md
@@ -15,7 +15,7 @@ usually do. We repeat this for an increasing number of files.
 <!--- TABLE-START -->
 | code generator  | files    | bundle size             | minified               | compressed         |
 |-----------------|----------|------------------------:|-----------------------:|-------------------:|
-| protobuf-es | 1 | 79,463 b | 34,300 b | 9,781 b |
+| protobuf-es | 1 | 79,463 b | 34,300 b | 9,786 b |
 | protobuf-es | 4 | 92,560 b | 37,274 b | 10,114 b |
 | protobuf-es | 8 | 101,901 b | 41,772 b | 10,808 b |
 | protobuf-es | 16 | 165,581 b | 67,017 b | 13,320 b |

diff --git a/packages/bundle-size/chart.svg b/packages/bundle-size/chart.svg
diff --git a/packages/protobuf-test/src/create-registry.test.ts b/packages/protobuf-test/src/create-registry.test.ts
@@ -15,6 +15,7 @@
 import { describe, expect, test } from "@jest/globals";
 import {
   createRegistry,
+  createMutableRegistry,
   Int32Value,
   MethodKind,
   proto3,
@@ -155,3 +156,11 @@ describe("createRegistry()", () => {
     });
   });
 });
+
+describe("createMutableRegistry()", () => {
+  test("add() adds message", () => {
+    const reg = createMutableRegistry();
+    reg.add(MessageFieldMessage);
+    expect(reg.findMessage(MessageFieldMessage.typeName)).toBeDefined();
+  });
+});
diff --git a/packages/protobuf/src/create-registry.ts b/packages/protobuf/src/create-registry.ts
@@ -19,6 +19,7 @@ import type {
   IEnumTypeRegistry,
   IExtensionRegistry,
   IMessageTypeRegistry,
+  IMutableRegistry,
   IServiceTypeRegistry,
 } from "./type-registry.js";
 import type { Extension } from "./extension.js";
@@ -33,12 +34,24 @@ export function createRegistry(
   IEnumTypeRegistry &
   IExtensionRegistry &
   IServiceTypeRegistry {
+  const mutable = createMutableRegistry(...types);
+  delete (mutable as Partial<IMutableRegistry>).add;
+  return mutable;
+}
+
+/**
+ * Create a mutable registry from the given types.
+ */
+export function createMutableRegistry(
+  ...types: Array<MessageType | EnumType | ServiceType | Extension>
+): IMutableRegistry {
   const messages: Record<string, MessageType> = {};
   const enums: Record<string, EnumType> = {};
   const services: Record<string, ServiceType> = {};
   const extensionsByName = new Map<string, Extension>();
   const extensionsByExtendee = new Map<string, Map<number, Extension>>();
-  const registry = {
+
+  const registry: IMutableRegistry = {
     findMessage(typeName: string): MessageType | undefined {
       return messages[typeName];
     },
@@ -54,49 +67,51 @@ export function createRegistry(
     findExtension(typeName: string): Extension | undefined {
       return extensionsByName.get(typeName) ?? undefined;
     },
-  };
-
-  function addType(type: MessageType | EnumType | ServiceType | Extension) {
-    if ("fields" in type) {
-      if (!registry.findMessage(type.typeName)) {
-        messages[type.typeName] = type;
-        type.fields.list().forEach(addField);
-      }
-    } else if ("methods" in type) {
-      if (!registry.findService(type.typeName)) {
-        services[type.typeName] = type;
-        for (const method of Object.values(type.methods)) {
-          addType(method.I);
-          addType(method.O);
+    add(type: MessageType | EnumType | ServiceType | Extension) {
+      if ("fields" in type) {
+        if (!this.findMessage(type.typeName)) {
+          messages[type.typeName] = type;
+          type.fields.list().forEach(addField);
         }
-      }
-    } else if ("extendee" in type) {
-      if (!extensionsByName.has(type.typeName)) {
-        extensionsByName.set(type.typeName, type);
-        const extendeeName = type.extendee.typeName;
-        if (!extensionsByExtendee.has(extendeeName)) {
-          extensionsByExtendee.set(extendeeName, new Map<number, Extension>());
+      } else if ("methods" in type) {
+        if (!this.findService(type.typeName)) {
+          services[type.typeName] = type;
+          for (const method of Object.values(type.methods)) {
+            this.add(method.I);
+            this.add(method.O);
+          }
         }
-        extensionsByExtendee.get(extendeeName)?.set(type.field.no, type);
-        addType(type.extendee);
-        addField(type.field);
+      } else if ("extendee" in type) {
+        if (!extensionsByName.has(type.typeName)) {
+          extensionsByName.set(type.typeName, type);
+          const extendeeName = type.extendee.typeName;
+          if (!extensionsByExtendee.has(extendeeName)) {
+            extensionsByExtendee.set(
+              extendeeName,
+              new Map<number, Extension>(),
+            );
+          }
+          extensionsByExtendee.get(extendeeName)?.set(type.field.no, type);
+          this.add(type.extendee);
+          addField(type.field);
+        }
+      } else {
+        enums[type.typeName] = type;
       }
-    } else {
-      enums[type.typeName] = type;
-    }
-  }
+    },
+  };
 
   function addField(field: FieldInfo) {
     if (field.kind == "message") {
-      addType(field.T);
+      registry.add(field.T);
     } else if (field.kind == "map" && field.V.kind == "message") {
-      addType(field.V.T);
+      registry.add(field.V.T);
     } else if (field.kind == "enum") {
-      addType(field.T);
+      registry.add(field.T);
     }
   }
   for (const type of types) {
-    addType(type);
+    registry.add(type);
   }
   return registry;
 }
diff --git a/packages/protobuf/src/index.ts b/packages/protobuf/src/index.ts
@@ -84,7 +84,7 @@ export type {
   IMessageTypeRegistry,
   IExtensionRegistry,
 } from "./type-registry.js";
-export { createRegistry } from "./create-registry.js";
+export { createRegistry, createMutableRegistry } from "./create-registry.js";
 export { createRegistryFromDescriptors } from "./create-registry-from-desc.js";
 export { toPlainMessage } from "./to-plain-message.js";
 

diff --git a/packages/protobuf/src/type-registry.ts b/packages/protobuf/src/type-registry.ts
@@ -74,3 +74,17 @@ export interface IExtensionRegistry {
    */
   findExtension(typeName: string): Extension | undefined;
 }
+
+/**
+ * A registry that allows
+ */
+export interface IMutableRegistry
+  extends IMessageTypeRegistry,
+    IEnumTypeRegistry,
+    IServiceTypeRegistry,
+    IExtensionRegistry {
+  /**
+   * Adds the type to the registry.
+   */
+  add(type: MessageType | EnumType | ServiceType | Extension): void;
+}