1 files changed, 52 insertions, 19 deletions
diff --git a/java/src/main/java/com/google/protobuf/nano/ExtendableMessageNano.java b/java/src/main/java/com/google/protobuf/nano/ExtendableMessageNano.java
index 63c8afc..a0c2731 100644
--- a/java/src/main/java/com/google/protobuf/nano/ExtendableMessageNano.java
+++ b/java/src/main/java/com/google/protobuf/nano/ExtendableMessageNano.java
@@ -31,8 +31,6 @@
 package com.google.protobuf.nano;
 
 import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
 
 /**
  * Base class of those Protocol Buffer messages that need to store unknown fields,
@@ -44,27 +42,28 @@ public abstract class ExtendableMessageNano<M extends ExtendableMessageNano<M>>
      * A container for fields unknown to the message, including extensions. Extension fields can
      * can be accessed through the {@link #getExtension} and {@link #setExtension} methods.
      */
-    protected List<UnknownFieldData> unknownFieldData;
+    protected FieldArray unknownFieldData;
 
     @Override
     protected int computeSerializedSize() {
         int size = 0;
-        int unknownFieldCount = unknownFieldData == null ? 0 : unknownFieldData.size();
-        for (int i = 0; i < unknownFieldCount; i++) {
-            UnknownFieldData unknownField = unknownFieldData.get(i);
-            size += CodedOutputByteBufferNano.computeRawVarint32Size(unknownField.tag);
-            size += unknownField.bytes.length;
+        if (unknownFieldData != null) {
+            for (int i = 0; i < unknownFieldData.size(); i++) {
+                FieldData field = unknownFieldData.dataAt(i);
+                size += field.computeSerializedSize();
+            }
         }
         return size;
     }
 
     @Override
     public void writeTo(CodedOutputByteBufferNano output) throws IOException {
-        int unknownFieldCount = unknownFieldData == null ? 0 : unknownFieldData.size();
-        for (int i = 0; i < unknownFieldCount; i++) {
-            UnknownFieldData unknownField = unknownFieldData.get(i);
-            output.writeRawVarint32(unknownField.tag);
-            output.writeRawBytes(unknownField.bytes);
+        if (unknownFieldData == null) {
+            return;
+        }
+        for (int i = 0; i < unknownFieldData.size(); i++) {
+            FieldData field = unknownFieldData.dataAt(i);
+            field.writeTo(output);
         }
     }
 
@@ -72,14 +71,38 @@ public abstract class ExtendableMessageNano<M extends ExtendableMessageNano<M>>
      * Gets the value stored in the specified extension of this message.
      */
     public final <T> T getExtension(Extension<M, T> extension) {
-        return extension.getValueFrom(unknownFieldData);
+        if (unknownFieldData == null) {
+            return null;
+        }
+        FieldData field = unknownFieldData.get(WireFormatNano.getTagFieldNumber(extension.tag));
+        return field == null ? null : field.getValue(extension);
     }
 
     /**
      * Sets the value of the specified extension of this message.
      */
     public final <T> M setExtension(Extension<M, T> extension, T value) {
-        unknownFieldData = extension.setValueTo(value, unknownFieldData);
+        int fieldNumber = WireFormatNano.getTagFieldNumber(extension.tag);
+        if (value == null) {
+            if (unknownFieldData != null) {
+                unknownFieldData.remove(fieldNumber);
+                if (unknownFieldData.isEmpty()) {
+                    unknownFieldData = null;
+                }
+            }
+        } else {
+            FieldData field = null;
+            if (unknownFieldData == null) {
+                unknownFieldData = new FieldArray();
+            } else {
+                field = unknownFieldData.get(fieldNumber);
+            }
+            if (field == null) {
+                unknownFieldData.put(fieldNumber, new FieldData(extension, value));
+            } else {
+                field.setValue(extension, value);
+            }
+        }
 
         @SuppressWarnings("unchecked") // Generated code should guarantee type safety
         M typedThis = (M) this;
@@ -106,12 +129,22 @@ public abstract class ExtendableMessageNano<M extends ExtendableMessageNano<M>>
         if (!input.skipField(tag)) {
             return false;  // This wasn't an unknown field, it's an end-group tag.
         }
-        if (unknownFieldData == null) {
-            unknownFieldData = new ArrayList<UnknownFieldData>();
-        }
+        int fieldNumber = WireFormatNano.getTagFieldNumber(tag);
         int endPos = input.getPosition();
         byte[] bytes = input.getData(startPos, endPos - startPos);
-        unknownFieldData.add(new UnknownFieldData(tag, bytes));
+        UnknownFieldData unknownField = new UnknownFieldData(tag, bytes);
+
+        FieldData field = null;
+        if (unknownFieldData == null) {
+            unknownFieldData = new FieldArray();
+        } else {
+            field = unknownFieldData.get(fieldNumber);
+        }
+        if (field == null) {
+            field = new FieldData();
+            unknownFieldData.put(fieldNumber, field);
+        }
+        field.addUnknownField(unknownField);
         return true;
     }
 }