Fix the behavior when merging conflicting keys, the new value always

override the existing one even for message types.

3 files changed, 35 insertions, 34 deletions

diff --git a/javanano/src/main/java/com/google/protobuf/nano/InternalNano.java b/javanano/src/main/java/com/google/protobuf/nano/InternalNano.java
index 0705b0c0..fc0a36cd 100644
--- a/javanano/src/main/java/com/google/protobuf/nano/InternalNano.java
+++ b/javanano/src/main/java/com/google/protobuf/nano/InternalNano.java
@@ -391,12 +391,12 @@ public final class InternalNano {
    * be called by generated messages.
    *
    * @param map the map field; may be null, in which case a map will be
-   *        instantiated using the {@link MapUtil.MapFactory}
+   *        instantiated using the {@link MapFactories.MapFactory}
    * @param input the input byte buffer
    * @param keyType key type, as defined in InternalNano.TYPE_*
    * @param valueType value type, as defined in InternalNano.TYPE_*
-   * @param valueClazz class of the value field if the valueType is
-   *        TYPE_MESSAGE; otherwise the parameter is ignored and can be null.
+   * @param value an new instance of the value, if the value is a TYPE_MESSAGE;
+   *        otherwise this parameter can be null and will be ignored.
    * @param keyTag wire tag for the key
    * @param valueTag wire tag for the value
    * @return the map field
@@ -408,15 +408,13 @@ public final class InternalNano {
       Map<K, V> map,
       int keyType,
       int valueType,
-      Class<V> valueClazz,
+      V value,
       int keyTag,
       int valueTag) throws IOException {
     map = MapFactories.getMapFactory().forMap(map);
     final int length = input.readRawVarint32();
     final int oldLimit = input.pushLimit(length);
-    byte[] payload = null;
     K key = null;
-    V value = null;
     while (true) {
       int tag = input.readTag();
       if (tag == 0) {
@@ -426,7 +424,7 @@ public final class InternalNano {
         key = (K) input.readData(keyType);
       } else if (tag == valueTag) {
         if (valueType == TYPE_MESSAGE) {
-          payload = input.readBytes();
+          input.readMessage((MessageNano) value);
         } else {
           value = (V) input.readData(valueType);
         }
@@ -440,36 +438,12 @@ public final class InternalNano {
     input.popLimit(oldLimit);
 
     if (key == null) {
+      // key can only be primitive types.
       key = (K) primitiveDefaultValue(keyType);
     }
 
-    // Special case: merge the value when the value is a message.
-    if (valueType == TYPE_MESSAGE) {
-      MessageNano oldMessageValue = (MessageNano) map.get(key);
-      if (oldMessageValue != null) {
-        if (payload != null) {
-          MessageNano.mergeFrom(oldMessageValue, payload);
-        }
-        return map;
-      }
-      // Otherwise, create a new value message.
-      try {
-        value = valueClazz.newInstance();
-      } catch (InstantiationException e) {
-        throw new IOException(
-            "Unable to create value message " + valueClazz.getName()
-            + " in maps.");
-      } catch (IllegalAccessException e) {
-        throw new IOException(
-            "Unable to create value message " + valueClazz.getName()
-            + " in maps.");
-      }
-      if (payload != null) {
-        MessageNano.mergeFrom((MessageNano) value, payload);
-      }
-    }
-
     if (value == null) {
+      // message type
       value = (V) primitiveDefaultValue(valueType);
     }
 
diff --git a/javanano/src/test/java/com/google/protobuf/nano/NanoTest.java b/javanano/src/test/java/com/google/protobuf/nano/NanoTest.java
index a7383cb4..bf8a391a 100644
--- a/javanano/src/test/java/com/google/protobuf/nano/NanoTest.java
+++ b/javanano/src/test/java/com/google/protobuf/nano/NanoTest.java
@@ -3742,7 +3742,6 @@ public class NanoTest extends TestCase {
     byte[] output = MessageNano.toByteArray(origin);
     TestMap parsed = new TestMap();
     MessageNano.mergeFrom(parsed, output);
-    // TODO(liujisi): Test merging message type values.
     // TODO(liujisi): Test missing key/value in parsing.
   }
 
@@ -3769,6 +3768,33 @@ public class NanoTest extends TestCase {
     }
   }
 
+  /**
+   * Tests that merging bytes containing conflicting keys with override the
+   * message value instead of merging the message value into the existing entry.
+   */
+  public void testMapMergeOverrideMessageValues() throws Exception {
+    TestMap.MessageValue origValue = new TestMap.MessageValue();
+    origValue.value = 1;
+    origValue.value2 = 2;
+    TestMap.MessageValue newValue = new TestMap.MessageValue();
+    newValue.value = 3;
+
+    TestMap origMessage = new TestMap();
+    origMessage.int32ToMessageField =
+        new HashMap<Integer, MapTestProto.TestMap.MessageValue>();
+    origMessage.int32ToMessageField.put(1, origValue);
+
+    TestMap newMessage = new TestMap();
+    newMessage.int32ToMessageField =
+        new HashMap<Integer, MapTestProto.TestMap.MessageValue>();
+    newMessage.int32ToMessageField.put(1, newValue);
+    MessageNano.mergeFrom(origMessage,
+        MessageNano.toByteArray(newMessage));
+    TestMap.MessageValue mergedValue = origMessage.int32ToMessageField.get(1);
+    assertEquals(3, mergedValue.value);
+    assertEquals(0, mergedValue.value2);
+  }
+
   private static final Integer[] int32Values = new Integer[] {
     0, 1, -1, Integer.MAX_VALUE, Integer.MIN_VALUE,
   };
diff --git a/javanano/src/test/java/com/google/protobuf/nano/map_test.proto b/javanano/src/test/java/com/google/protobuf/nano/map_test.proto
index 5ea86717..f72833ad 100644
--- a/javanano/src/test/java/com/google/protobuf/nano/map_test.proto
+++ b/javanano/src/test/java/com/google/protobuf/nano/map_test.proto
@@ -38,6 +38,7 @@ option java_outer_classname = "MapTestProto";
 message TestMap {
   message MessageValue {
     int32 value = 1;
+    int32 value2 = 2;
   }
   enum EnumValue {
     FOO = 0;