diff --git a/java/core/src/main/java/com/google/protobuf/ExtensionSchemaLite.java b/java/core/src/main/java/com/google/protobuf/ExtensionSchemaLite.java
index ca0095123f..19cad5f212 100644
--- a/java/core/src/main/java/com/google/protobuf/ExtensionSchemaLite.java
+++ b/java/core/src/main/java/com/google/protobuf/ExtensionSchemaLite.java
@@ -255,6 +255,9 @@ final class ExtensionSchemaLite extends ExtensionSchema<ExtensionDescriptor> {
             value = reader.readString();
             break;
           case GROUP:
+            // Special case handling for non-repeated sub-messages: merge in-place rather than
+            // building up new sub-messages and merging those, which is too slow.
+            // TODO(b/249368670): clean this up
             if (!extension.isRepeated()) {
               Object oldValue = extensions.getField(extension.descriptor);
               if (oldValue instanceof GeneratedMessageLite) {
@@ -275,6 +278,9 @@ final class ExtensionSchemaLite extends ExtensionSchema<ExtensionDescriptor> {
             break;
 
           case MESSAGE:
+            // Special case handling for non-repeated sub-messages: merge in-place rather than
+            // building up new sub-messages and merging those, which is too slow.
+            // TODO(b/249368670): clean this up
             if (!extension.isRepeated()) {
               Object oldValue = extensions.getField(extension.descriptor);
               if (oldValue instanceof GeneratedMessageLite) {
@@ -304,6 +310,7 @@ final class ExtensionSchemaLite extends ExtensionSchema<ExtensionDescriptor> {
         switch (extension.getLiteType()) {
           case MESSAGE:
           case GROUP:
+            // TODO(b/249368670): this shouldn't be reachable, clean this up
             Object oldValue = extensions.getField(extension.descriptor);
             if (oldValue != null) {
               value = Internal.mergeMessage(oldValue, value);
diff --git a/java/core/src/main/java/com/google/protobuf/FieldSet.java b/java/core/src/main/java/com/google/protobuf/FieldSet.java
index 4f514c224e..b64f63b351 100644
--- a/java/core/src/main/java/com/google/protobuf/FieldSet.java
+++ b/java/core/src/main/java/com/google/protobuf/FieldSet.java
@@ -39,6 +39,7 @@ import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 
 /**
  * A class which represents an arbitrary set of fields of some message type. This is used to
@@ -123,9 +124,10 @@ final class FieldSet<T extends FieldSet.FieldDescriptorLite<T>> {
     if (isImmutable) {
       return;
     }
-    for (Object value : fields.values()) {
-      if (value instanceof GeneratedMessageLite) {
-        ((GeneratedMessageLite<?, ?>) value).makeImmutable();
+    for (int i = 0; i < fields.getNumArrayEntries(); ++i) {
+      Entry<T, Object> entry = fields.getArrayEntryAt(i);
+      if (entry.getValue() instanceof GeneratedMessageLite) {
+        ((GeneratedMessageLite<?, ?>) entry.getValue()).makeImmutable();
       }
     }
     fields.makeImmutable();