Save state before using BinaryWriter

packages/protobuf-test/src/msg-message.test.ts

@@ -14,9 +14,11 @@
 
 import { expect, test } from "@jest/globals";
 import type { JsonValue, PlainMessage } from "@bufbuild/protobuf";
+import { BinaryWriter } from "@bufbuild/protobuf";
 import { MessageFieldMessage as TS_MessageFieldMessage } from "./gen/ts/extra/msg-message_pb.js";
 import { MessageFieldMessage as JS_MessageFieldMessage } from "./gen/js/extra/msg-message_pb.js";
 import { describeMT } from "./helpers.js";
+import assert from "node:assert";
 
 describeMT(
   { ts: TS_MessageFieldMessage, js: JS_MessageFieldMessage },
@@ -102,5 +104,33 @@ describeMT(
         }),
       );
     });
+    test("example binary encodes and decodes using same writer instance", () => {
+      const msg = new messageType(exampleFields);
+      const writer = new BinaryWriter();
+
+      let runs = 10;
+      let enc;
+      while (runs > 0) {
+        enc = msg.toBinary({
+          writerFactory: () => writer,
+        });
+        runs--;
+      }
+      assert(enc !== undefined);
+
+      const got = messageType.fromBinary(enc);
+      expect(got.messageField?.name).toStrictEqual(
+        exampleFields.messageField.name,
+      );
+      expect(got.repeatedMessageField.length).toStrictEqual(
+        exampleFields.repeatedMessageField.length,
+      );
+      expect(got.repeatedMessageField[0].name).toStrictEqual(
+        exampleFields.repeatedMessageField[0].name,
+      );
+      expect(got.repeatedMessageField[1].name).toStrictEqual(
+        exampleFields.repeatedMessageField[1].name,
+      );
+    });
   },
 );

packages/protobuf/src/binary-encoding.ts

@@ -172,7 +172,7 @@ export interface IBinaryReader {
 
 export interface IBinaryWriter {
   /**
-   * Return all bytes written and reset this writer.
+   * Return all bytes written and return to the previous state.
    */
   finish(): Uint8Array;
 
@@ -314,6 +314,8 @@ export class BinaryWriter implements IBinaryWriter {
     this.textEncoder = textEncoder ?? new TextEncoder();
     this.chunks = [];
     this.buf = [];
+
+    this.fork();
   }
 
   /**
@@ -329,15 +331,28 @@ export class BinaryWriter implements IBinaryWriter {
       bytes.set(this.chunks[i], offset);
       offset += this.chunks[i].length;
     }
-    this.chunks = [];
+
+    this.restore();
+
     return bytes;
   }
 
+  /*
+   * Restore the previous stack entry.
+   */
+  restore(): IBinaryWriter {
+    let prev = this.stack.pop();
+    if (!prev) throw new Error("invalid state, fork stack empty");
+    this.chunks = prev.chunks;
+    this.buf = prev.buf;
+    return this;
+  }
+
   /**
    * Start a new fork for length-delimited data like a message
    * or a packed repeated field.
    *
-   * Must be joined later with `join()`.
+   * Must be joined later with `join()` or flushed with `finish()`.
    */
   fork(): IBinaryWriter {
     this.stack.push({ chunks: this.chunks, buf: this.buf });
@@ -353,13 +368,6 @@ export class BinaryWriter implements IBinaryWriter {
   join(): IBinaryWriter {
     // get chunk of fork
     let chunk = this.finish();
-
-    // restore previous state
-    let prev = this.stack.pop();
-    if (!prev) throw new Error("invalid state, fork stack empty");
-    this.chunks = prev.chunks;
-    this.buf = prev.buf;
-
     // write length of chunk as varint
     this.uint32(chunk.byteLength);
     return this.raw(chunk);

packages/protobuf/src/extension-accessor.ts

@@ -87,6 +87,7 @@ export function setExtension<E extends Message<E>, V>(
     }
   }
   const writer = writeOpt.writerFactory();
+  writer.fork();
   extension.runtime.bin.writeField(extension.field, value, writer, writeOpt);
   const reader = readOpt.readerFactory(writer.finish());
   while (reader.pos < reader.len) {

packages/protobuf/src/message.ts

@@ -109,6 +109,7 @@ export class Message<T extends Message<T> = AnyMessage> {
       bin = type.runtime.bin,
       opt = bin.makeWriteOptions(options),
       writer = opt.writerFactory();
+    writer.fork();
     bin.writeMessage(this, writer, opt);
     return writer.finish();
   }

packages/protobuf/src/proto-delimited.ts

@@ -34,7 +34,7 @@ export const protoDelimited = {
    */
   enc(message: Message, options?: BinaryWriteOptions): Uint8Array {
     const opt = makeBinaryFormatCommon().makeWriteOptions(options);
-    return opt.writerFactory().bytes(message.toBinary(opt)).finish();
+    return opt.writerFactory().fork().bytes(message.toBinary(opt)).finish();
   },
 
   /**