fix: Support repeated string enums. (#284)

* fix: Support repeated string enums. * Fix packed encoding. Co-authored-by: kanziw <kanziwoong@gmail.com>
stephenh · Apr 24, 2021 · be9ecf7 · be9ecf7
1 parent cf0b200
commit be9ecf7
Show file tree

Hide file tree

Showing 5 changed files with 56 additions and 2 deletions.
diff --git a/integration/simple-string-enums/simple-test.ts b/integration/simple-string-enums/simple-test.ts
@@ -2,13 +2,17 @@ import { Simple, StateEnum } from './simple';
 
 describe('simple-string-enums', () => {
   it('encodes', () => {
-    const s1: Simple = { name: 'a', state: StateEnum.ON };
+    const s1: Simple = { name: 'a', state: StateEnum.ON, states: [StateEnum.ON, StateEnum.OFF] };
     const b = Simple.encode(s1).finish();
     const s2 = Simple.decode(b);
     expect(s2).toMatchInlineSnapshot(`
       Object {
         "name": "a",
         "state": "ON",
+        "states": Array [
+          "ON",
+          "OFF",
+        ],
       }
     `);
   });

diff --git a/integration/simple-string-enums/simple.bin b/integration/simple-string-enums/simple.bin
diff --git a/integration/simple-string-enums/simple.proto b/integration/simple-string-enums/simple.proto
@@ -4,6 +4,7 @@ package simple;
 message Simple {
   string name = 1;
   StateEnum state = 4;
+  repeated StateEnum states = 5;
 }
 
 enum StateEnum {

diff --git a/integration/simple-string-enums/simple.ts b/integration/simple-string-enums/simple.ts
@@ -58,9 +58,10 @@ export function stateEnumToNumber(object: StateEnum): number {
 export interface Simple {
   name: string;
   state: StateEnum;
+  states: StateEnum[];
 }
 
-const baseSimple: object = { name: '', state: StateEnum.UNKNOWN };
+const baseSimple: object = { name: '', state: StateEnum.UNKNOWN, states: StateEnum.UNKNOWN };
 
 export const Simple = {
   encode(message: Simple, writer: Writer = Writer.create()): Writer {
@@ -70,13 +71,19 @@ export const Simple = {
     if (message.state !== StateEnum.UNKNOWN) {
       writer.uint32(32).int32(stateEnumToNumber(message.state));
     }
+    writer.uint32(42).fork();
+    for (const v of message.states) {
+      writer.int32(stateEnumToNumber(v));
+    }
+    writer.ldelim();
     return writer;
   },
 
   decode(input: Reader | Uint8Array, length?: number): Simple {
     const reader = input instanceof Reader ? input : new Reader(input);
     let end = length === undefined ? reader.len : reader.pos + length;
     const message = { ...baseSimple } as Simple;
+    message.states = [];
     while (reader.pos < end) {
       const tag = reader.uint32();
       switch (tag >>> 3) {
@@ -86,6 +93,16 @@ export const Simple = {
         case 4:
           message.state = stateEnumFromJSON(reader.int32());
           break;
+        case 5:
+          if ((tag & 7) === 2) {
+            const end2 = reader.uint32() + reader.pos;
+            while (reader.pos < end2) {
+              message.states.push(stateEnumFromJSON(reader.int32()));
+            }
+          } else {
+            message.states.push(stateEnumFromJSON(reader.int32()));
+          }
+          break;
         default:
           reader.skipType(tag & 7);
           break;
@@ -96,6 +113,7 @@ export const Simple = {
 
   fromJSON(object: any): Simple {
     const message = { ...baseSimple } as Simple;
+    message.states = [];
     if (object.name !== undefined && object.name !== null) {
       message.name = String(object.name);
     } else {
@@ -106,18 +124,29 @@ export const Simple = {
     } else {
       message.state = StateEnum.UNKNOWN;
     }
+    if (object.states !== undefined && object.states !== null) {
+      for (const e of object.states) {
+        message.states.push(stateEnumFromJSON(e));
+      }
+    }
     return message;
   },
 
   toJSON(message: Simple): unknown {
     const obj: any = {};
     message.name !== undefined && (obj.name = message.name);
     message.state !== undefined && (obj.state = stateEnumToJSON(message.state));
+    if (message.states) {
+      obj.states = message.states.map((e) => stateEnumToJSON(e));
+    } else {
+      obj.states = [];
+    }
     return obj;
   },
 
   fromPartial(object: DeepPartial<Simple>): Simple {
     const message = { ...baseSimple } as Simple;
+    message.states = [];
     if (object.name !== undefined && object.name !== null) {
       message.name = object.name;
     } else {
@@ -128,6 +157,11 @@ export const Simple = {
     } else {
       message.state = StateEnum.UNKNOWN;
     }
+    if (object.states !== undefined && object.states !== null) {
+      for (const e of object.states) {
+        message.states.push(e);
+      }
+    }
     return message;
   },
 };

diff --git a/src/main.ts b/src/main.ts
@@ -793,7 +793,22 @@ function generateEncode(ctx: Context, fullName: string, messageDesc: DescriptorP
             ${writeSnippet('v!')};
           }
         `);
+      } else if (isEnum(field) && options.stringEnums) {
+        // This is a lot like the `else` clause, but we wrap `fooToNumber` around it.
+        // Ideally we'd reuse `writeSnippet` here, but `writeSnippet` has the `writer.uint32(tag)`
+        // embedded inside of it, and we want to drop that so that we can encode it packed
+        // (i.e. just one tag and multiple values).
+        const tag = ((field.number << 3) | 2) >>> 0;
+        const toNumber = getEnumMethod(typeMap, field.typeName, 'ToNumber');
+        chunks.push(code`
+          writer.uint32(${tag}).fork();
+          for (const v of message.${fieldName}) {
+            writer.${toReaderCall(field)}(${toNumber}(v));
+          }
+          writer.ldelim();
+        `);
       } else {
+        // Ideally we'd reuse `writeSnippet` but it has tagging embedded inside of it.
         const tag = ((field.number << 3) | 2) >>> 0;
         chunks.push(code`
           writer.uint32(${tag}).fork();