From 0812bbfcf520ec325227c0d91f2decf4bd20be1c Mon Sep 17 00:00:00 2001
From: Mark Hansen <markhansen@google.com>
Date: Mon, 17 Jun 2024 16:41:15 -0700
Subject: [PATCH] Add new runtime API that serializes empty extensions without
 allocating

I've made a new name for the interface, ExtensionSerializer, so we can keep ExtensionWriter in the same place for backwards-compatibility for a little while.

PiperOrigin-RevId: 644172922
---
 .../com/google/protobuf/GeneratedMessage.java | 43 +++++++++++++++++--
 1 file changed, 40 insertions(+), 3 deletions(-)

diff --git a/java/core/src/main/java/com/google/protobuf/GeneratedMessage.java b/java/core/src/main/java/com/google/protobuf/GeneratedMessage.java
index 9576051d00..426adb61bf 100644
--- a/java/core/src/main/java/com/google/protobuf/GeneratedMessage.java
+++ b/java/core/src/main/java/com/google/protobuf/GeneratedMessage.java
@@ -1025,10 +1025,28 @@ public abstract class GeneratedMessage extends AbstractMessage implements Serial
 
     /**
      * Used by subclasses to serialize extensions. Extension ranges may be interleaved with field
-     * numbers, but we must write them in canonical (sorted by field number) order. ExtensionWriter
-     * helps us write individual ranges of extensions at once.
+     * numbers, but we must write them in canonical (sorted by field number) order.
+     * ExtensionSerializer helps us write individual ranges of extensions at once.
      */
-    protected class ExtensionWriter {
+    protected interface ExtensionSerializer {
+      public void writeUntil(final int end, final CodedOutputStream output) throws IOException;
+    }
+
+    /** No-op implementation that writes nothing, for messages with no extensions. */
+    private static final class NoOpExtensionSerializer implements ExtensionSerializer {
+      // Singleton instance so we can avoid allocating a new one for each message serialization.
+      private static final NoOpExtensionSerializer INSTANCE = new NoOpExtensionSerializer();
+
+      @Override
+      public void writeUntil(final int end, final CodedOutputStream output) {
+        // no-op
+      }
+    }
+
+    /**
+     * ExtensionSerializer that writes extensions from the FieldSet, for messages with extensions.
+     */
+    protected class ExtensionWriter implements ExtensionSerializer {
       // Imagine how much simpler this code would be if Java iterators had
       // a way to get the next element without advancing the iterator.
 
@@ -1043,6 +1061,7 @@ public abstract class GeneratedMessage extends AbstractMessage implements Serial
         this.messageSetWireFormat = messageSetWireFormat;
       }
 
+      @Override
       public void writeUntil(final int end, final CodedOutputStream output) throws IOException {
         while (next != null && next.getKey().getNumber() < end) {
           FieldDescriptor descriptor = next.getKey();
@@ -1075,14 +1094,32 @@ public abstract class GeneratedMessage extends AbstractMessage implements Serial
       }
     }
 
+    // TODO: Remove, replace with newExtensionSerializer().
     protected ExtensionWriter newExtensionWriter() {
       return new ExtensionWriter(false);
     }
 
+    protected ExtensionSerializer newExtensionSerializer() {
+      // Avoid allocation in the common case of no extensions.
+      if (extensions.isEmpty()) {
+        return NoOpExtensionSerializer.INSTANCE;
+      }
+      return new ExtensionWriter(false);
+    }
+
+    // TODO: Remove, replace with newMessageSetExtensionSerializer().
     protected ExtensionWriter newMessageSetExtensionWriter() {
       return new ExtensionWriter(true);
     }
 
+    protected ExtensionSerializer newMessageSetExtensionSerializer() {
+      // Avoid allocation in the common case of no extensions.
+      if (extensions.isEmpty()) {
+        return NoOpExtensionSerializer.INSTANCE;
+      }
+      return new ExtensionWriter(true);
+    }
+
     /** Called by subclasses to compute the size of extensions. */
     protected int extensionsSerializedSize() {
       return extensions.getSerializedSize();