parsing and serialzation for maps in JavaNano.

4 files changed, 316 insertions, 55 deletions

diff --git a/javanano/src/main/java/com/google/protobuf/nano/CodedOutputByteBufferNano.java b/javanano/src/main/java/com/google/protobuf/nano/CodedOutputByteBufferNano.java
index 37982b57..2777f34c 100644
--- a/javanano/src/main/java/com/google/protobuf/nano/CodedOutputByteBufferNano.java
+++ b/javanano/src/main/java/com/google/protobuf/nano/CodedOutputByteBufferNano.java
@@ -876,4 +876,128 @@ public final class CodedOutputByteBufferNano {
     // Note:  the right-shift must be arithmetic
     return (n << 1) ^ (n >> 63);
   }
+
+  static int computeFieldSize(int number, int type, Object object) {
+    switch (type) {
+      case InternalNano.TYPE_BOOL:
+        return computeBoolSize(number, (Boolean) object);
+      case InternalNano.TYPE_BYTES:
+        return computeBytesSize(number, (byte[]) object);
+      case InternalNano.TYPE_STRING:
+        return computeStringSize(number, (String) object);
+      case InternalNano.TYPE_FLOAT:
+        return computeFloatSize(number, (Float) object);
+      case InternalNano.TYPE_DOUBLE:
+        return computeDoubleSize(number, (Double) object);
+      case InternalNano.TYPE_ENUM:
+        return computeEnumSize(number, (Integer) object);
+      case InternalNano.TYPE_FIXED32:
+        return computeFixed32Size(number, (Integer) object);
+      case InternalNano.TYPE_INT32:
+        return computeInt32Size(number, (Integer) object);
+      case InternalNano.TYPE_UINT32:
+        return computeUInt32Size(number, (Integer) object);
+      case InternalNano.TYPE_SINT32:
+        return computeSInt32Size(number, (Integer) object);
+      case InternalNano.TYPE_SFIXED32:
+        return computeSFixed32Size(number, (Integer) object);
+      case InternalNano.TYPE_INT64:
+        return computeInt64Size(number, (Long) object);
+      case InternalNano.TYPE_UINT64:
+        return computeUInt64Size(number, (Long) object);
+      case InternalNano.TYPE_SINT64:
+        return computeSInt64Size(number, (Long) object);
+      case InternalNano.TYPE_FIXED64:
+        return computeFixed64Size(number, (Long) object);
+      case InternalNano.TYPE_SFIXED64:
+        return computeSFixed64Size(number, (Long) object);
+      case InternalNano.TYPE_MESSAGE:
+        return computeMessageSize(number, (MessageNano) object);
+      case InternalNano.TYPE_GROUP:
+        return computeGroupSize(number, (MessageNano) object);
+      default:
+        throw new IllegalArgumentException("Unknown type: " + type);
+    }
+  }
+
+  void writeField(int number, int type, Object value)
+      throws IOException {
+    switch (type) {
+      case InternalNano.TYPE_DOUBLE:
+        Double doubleValue = (Double) value;
+        writeDouble(number, doubleValue);
+        break;
+      case InternalNano.TYPE_FLOAT:
+        Float floatValue = (Float) value;
+        writeFloat(number, floatValue);
+        break;
+      case InternalNano.TYPE_INT64:
+        Long int64Value = (Long) value;
+        writeInt64(number, int64Value);
+        break;
+      case InternalNano.TYPE_UINT64:
+        Long uint64Value = (Long) value;
+        writeUInt64(number, uint64Value);
+        break;
+      case InternalNano.TYPE_INT32:
+        Integer int32Value = (Integer) value;
+        writeInt32(number, int32Value);
+        break;
+      case InternalNano.TYPE_FIXED64:
+        Long fixed64Value = (Long) value;
+        writeFixed64(number, fixed64Value);
+        break;
+      case InternalNano.TYPE_FIXED32:
+        Integer fixed32Value = (Integer) value;
+        writeFixed32(number, fixed32Value);
+        break;
+      case InternalNano.TYPE_BOOL:
+        Boolean boolValue = (Boolean) value;
+        writeBool(number, boolValue);
+        break;
+      case InternalNano.TYPE_STRING:
+        String stringValue = (String) value;
+        writeString(number, stringValue);
+        break;
+      case InternalNano.TYPE_BYTES:
+        byte[] bytesValue = (byte[]) value;
+        writeBytes(number, bytesValue);
+        break;
+      case InternalNano.TYPE_UINT32:
+        Integer uint32Value = (Integer) value;
+        writeUInt32(number, uint32Value);
+        break;
+      case InternalNano.TYPE_ENUM:
+        Integer enumValue = (Integer) value;
+        writeEnum(number, enumValue);
+        break;
+      case InternalNano.TYPE_SFIXED32:
+        Integer sfixed32Value = (Integer) value;
+        writeSFixed32(number, sfixed32Value);
+        break;
+      case InternalNano.TYPE_SFIXED64:
+        Long sfixed64Value = (Long) value;
+        writeSFixed64(number, sfixed64Value);
+        break;
+      case InternalNano.TYPE_SINT32:
+        Integer sint32Value = (Integer) value;
+        writeSInt32(number, sint32Value);
+        break;
+      case InternalNano.TYPE_SINT64:
+        Long sint64Value = (Long) value;
+        writeSInt64(number, sint64Value);
+        break;
+      case InternalNano.TYPE_MESSAGE:
+        MessageNano messageValue = (MessageNano) value;
+        writeMessage(number, messageValue);
+        break;
+      case InternalNano.TYPE_GROUP:
+        MessageNano groupValue = (MessageNano) value;
+        writeGroup(number, groupValue);
+        break;
+      default:
+        throw new IOException("Unknown type: " + type);
+    }
+  }
+
 }
diff --git a/javanano/src/main/java/com/google/protobuf/nano/MapUtil.java b/javanano/src/main/java/com/google/protobuf/nano/MapUtil.java
index 8e7647dd..bc544081 100644
--- a/javanano/src/main/java/com/google/protobuf/nano/MapUtil.java
+++ b/javanano/src/main/java/com/google/protobuf/nano/MapUtil.java
@@ -33,6 +33,7 @@ package com.google.protobuf.nano;
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Map.Entry;
 
 /**
  * Utility class for maps support.
@@ -55,42 +56,178 @@ public final class MapUtil {
   }
   private static volatile MapFactory mapFactory = new DefaultMapFactory();
 
-  @SuppressWarnings("unchecked")
-  public static final <K, V> Map<K, V> mergeEntry(
-      Map<K, V> target, CodedInputByteBufferNano input,
-      int keyType, int valueType, V value,
-      int keyTag, int valueTag)
-      throws IOException {
-    target = mapFactory.forMap(target);
-    final int length = input.readRawVarint32();
-    final int oldLimit = input.pushLimit(length);
-    K key = null;
-    while (true) {
-      int tag = input.readTag();
-      if (tag == 0) {
-        break;
+  /**
+   * Internal utilities to implement maps for generated messages.
+   * Do NOT use it explicitly.
+   */
+  public static class Internal {
+    private static final byte[] emptyBytes = new byte[0];
+    private static Object primitiveDefaultValue(int type) {
+      switch (type) {
+        case InternalNano.TYPE_BOOL:
+          return Boolean.FALSE;
+        case InternalNano.TYPE_BYTES:
+          return emptyBytes;
+        case InternalNano.TYPE_STRING:
+          return "";
+        case InternalNano.TYPE_FLOAT:
+          return Float.valueOf(0);
+        case InternalNano.TYPE_DOUBLE:
+          return Double.valueOf(0);
+        case InternalNano.TYPE_ENUM:
+        case InternalNano.TYPE_FIXED32:
+        case InternalNano.TYPE_INT32:
+        case InternalNano.TYPE_UINT32:
+        case InternalNano.TYPE_SINT32:
+        case InternalNano.TYPE_SFIXED32:
+          return Integer.valueOf(0);
+        case InternalNano.TYPE_INT64:
+        case InternalNano.TYPE_UINT64:
+        case InternalNano.TYPE_SINT64:
+        case InternalNano.TYPE_FIXED64:
+        case InternalNano.TYPE_SFIXED64:
+          return Long.valueOf(0L);
+        case InternalNano.TYPE_MESSAGE:
+        case InternalNano.TYPE_GROUP:
+        default:
+          throw new IllegalArgumentException(
+              "Type: " + type + " is not a primitive type.");
       }
-      if (tag == keyTag) {
-        key = (K) input.readData(keyType);
-      } else if (tag == valueTag) {
-        if (valueType == InternalNano.TYPE_MESSAGE) {
-          input.readMessage((MessageNano) value);
+    }
+
+    /**
+     * Merges the map entry into the map field. Note this is only supposed to
+     * 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}
+     * @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 keyTag wire tag for the key
+     * @param valueTag wire tag for the value
+     * @return the map field
+     * @throws IOException
+     */
+    @SuppressWarnings("unchecked")
+    public static final <K, V> Map<K, V> mergeEntry(
+        CodedInputByteBufferNano input,
+        Map<K, V> map,
+        int keyType,
+        int valueType,
+        Class<V> valueClazz,
+        int keyTag,
+        int valueTag) throws IOException {
+      map = mapFactory.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) {
+          break;
+        }
+        if (tag == keyTag) {
+          key = (K) input.readData(keyType);
+        } else if (tag == valueTag) {
+          if (valueType == InternalNano.TYPE_MESSAGE) {
+            payload = input.readBytes();
+          } else {
+            value = (V) input.readData(valueType);
+          }
         } else {
-          value = (V) input.readData(valueType);
+          if (!input.skipField(tag)) {
+            break;
+          }
         }
-      } else {
-        if (!input.skipField(tag)) {
-          break;
+      }
+      input.checkLastTagWas(0);
+      input.popLimit(oldLimit);
+
+      if (key == null) {
+        key = (K) primitiveDefaultValue(keyType);
+      }
+
+      // Special case: merge the value when the value is a message.
+      if (valueType == InternalNano.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) {
+        value = (V) primitiveDefaultValue(valueType);
+      }
+
+      map.put(key, value);
+      return map;
+    }
+
+    public static <K, V> void serializeMapField(
+        CodedOutputByteBufferNano output,
+        Map<K, V> map, int number, int keyType, int valueType)
+            throws IOException {
+      for (Entry<K, V> entry: map.entrySet()) {
+        K key = entry.getKey();
+        V value = entry.getValue();
+        if (key == null || value == null) {
+          throw new IllegalStateException(
+              "keys and values in maps cannot be null");
+        }
+        int entrySize =
+            CodedOutputByteBufferNano.computeFieldSize(1, keyType, key) +
+            CodedOutputByteBufferNano.computeFieldSize(2, valueType, value);
+        output.writeTag(number, WireFormatNano.WIRETYPE_LENGTH_DELIMITED);
+        output.writeRawVarint32(entrySize);
+        output.writeField(1, keyType, key);
+        output.writeField(2, valueType, value);
+      }
     }
-    input.checkLastTagWas(0);
-    input.popLimit(oldLimit);
 
-    if (key != null) {
-      target.put(key, value);
+    public static <K, V> int computeMapFieldSize(
+        Map<K, V> map, int number, int keyType, int valueType) {
+      int size = 0;
+      int tagSize = CodedOutputByteBufferNano.computeTagSize(number);
+      for (Entry<K, V> entry: map.entrySet()) {
+        K key = entry.getKey();
+        V value = entry.getValue();
+        if (key == null || value == null) {
+          throw new IllegalStateException(
+              "keys and values in maps cannot be null");
+        }
+        int entrySize =
+            CodedOutputByteBufferNano.computeFieldSize(1, keyType, key) +
+            CodedOutputByteBufferNano.computeFieldSize(2, valueType, value);
+        size += tagSize + entrySize
+            + CodedOutputByteBufferNano.computeRawVarint32Size(entrySize);
+      }
+      return size;
     }
-    return target;
+
+    private Internal() {}
   }
 
   private MapUtil() {}
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 442f0b74..d9428432 100644
--- a/javanano/src/test/java/com/google/protobuf/nano/NanoTest.java
+++ b/javanano/src/test/java/com/google/protobuf/nano/NanoTest.java
@@ -30,31 +30,9 @@
 
 package com.google.protobuf.nano;
 
-import com.google.protobuf.nano.CodedInputByteBufferNano;
-import com.google.protobuf.nano.EnumClassNanoMultiple;
-import com.google.protobuf.nano.EnumClassNanos;
-import com.google.protobuf.nano.EnumValidity;
-import com.google.protobuf.nano.EnumValidityAccessors;
-import com.google.protobuf.nano.FileScopeEnumMultiple;
-import com.google.protobuf.nano.FileScopeEnumRefNano;
-import com.google.protobuf.nano.InternalNano;
-import com.google.protobuf.nano.InvalidProtocolBufferNanoException;
-import com.google.protobuf.nano.MessageNano;
-import com.google.protobuf.nano.MessageScopeEnumRefNano;
-import com.google.protobuf.nano.MultipleImportingNonMultipleNano1;
-import com.google.protobuf.nano.MultipleImportingNonMultipleNano2;
-import com.google.protobuf.nano.MultipleNameClashNano;
 import com.google.protobuf.nano.NanoAccessorsOuterClass.TestNanoAccessors;
 import com.google.protobuf.nano.NanoHasOuterClass.TestAllTypesNanoHas;
-import com.google.protobuf.nano.NanoOuterClass;
 import com.google.protobuf.nano.NanoOuterClass.TestAllTypesNano;
-import com.google.protobuf.nano.NanoReferenceTypes;
-import com.google.protobuf.nano.NanoRepeatedPackables;
-import com.google.protobuf.nano.PackedExtensions;
-import com.google.protobuf.nano.RepeatedExtensions;
-import com.google.protobuf.nano.SingularExtensions;
-import com.google.protobuf.nano.TestRepeatedMergeNano;
-import com.google.protobuf.nano.UnittestMultipleNano;
 import com.google.protobuf.nano.UnittestRecursiveNano.RecursiveMessageNano;
 import com.google.protobuf.nano.UnittestSimpleNano.SimpleMessageNano;
 import com.google.protobuf.nano.UnittestSingleNano.SingleMessageNano;
@@ -3754,6 +3732,13 @@ public class NanoTest extends TestCase {
     assertTrue(Arrays.equals(new boolean[] {false, true, false, true}, nonPacked.bools));
   }
 
+  public void testMapsSerializeAndParse() throws Exception {
+    // TODO(liujisi): Test basic serialization/parsing roundtrip.
+    // TODO(liujisi): Test null values in serialization.
+    // TODO(liujisi): Test merging message type values.
+    // TODO(liujisi): Test missing key/value in parsing.
+  }
+
   private void assertRepeatedPackablesEqual(
       NanoRepeatedPackables.NonPacked nonPacked, NanoRepeatedPackables.Packed packed) {
     // Not using MessageNano.equals() -- that belongs to a separate test.
diff --git a/src/google/protobuf/compiler/javanano/javanano_map_field.cc b/src/google/protobuf/compiler/javanano/javanano_map_field.cc
index 4453cdfa..dead3685 100644
--- a/src/google/protobuf/compiler/javanano/javanano_map_field.cc
+++ b/src/google/protobuf/compiler/javanano/javanano_map_field.cc
@@ -89,6 +89,7 @@ void SetMapVariables(const Params& params,
   const FieldDescriptor* value = ValueField(descriptor);
   (*variables)["name"] =
     RenameJavaKeywords(UnderscoresToCamelCase(descriptor));
+  (*variables)["number"] = SimpleItoa(descriptor->number());
   (*variables)["key_type"] = TypeName(params, key, false);
   (*variables)["boxed_key_type"] = TypeName(params,key, true);
   (*variables)["key_desc_type"] =
@@ -101,9 +102,9 @@ void SetMapVariables(const Params& params,
   (*variables)["value_tag"] = SimpleItoa(internal::WireFormat::MakeTag(value));
   (*variables)["type_parameters"] =
       (*variables)["boxed_key_type"] + ", " + (*variables)["boxed_value_type"];
-  (*variables)["value_default"] =
+  (*variables)["value_class"] =
       value->type() == FieldDescriptor::TYPE_MESSAGE
-          ? "new " + (*variables)["value_type"] + "()"
+          ? (*variables)["value_type"] + ".class"
           : "null";
 }
 }  // namespace
@@ -132,21 +133,35 @@ GenerateClearCode(io::Printer* printer) const {
 void MapFieldGenerator::
 GenerateMergingCode(io::Printer* printer) const {
   printer->Print(variables_,
-    "$name$ = com.google.protobuf.nano.MapUtil.mergeEntry(\n"
-    "  $name$, input,\n"
+    "this.$name$ = com.google.protobuf.nano.MapUtil.Internal.mergeEntry(\n"
+    "  input, this.$name$,\n"
     "  com.google.protobuf.nano.InternalNano.$key_desc_type$,\n"
     "  com.google.protobuf.nano.InternalNano.$value_desc_type$,\n"
-    "  $value_default$,\n"
+    "  $value_class$,\n"
     "  $key_tag$, $value_tag$);\n"
     "\n");
 }
 
 void MapFieldGenerator::
 GenerateSerializationCode(io::Printer* printer) const {
+  printer->Print(variables_,
+    "if (this.$name$ != null) {\n"
+    "  com.google.protobuf.nano.MapUtil.Internal.serializeMapField(\n"
+    "    output, this.$name$, $number$,\n"
+    "  com.google.protobuf.nano.InternalNano.$key_desc_type$,\n"
+    "  com.google.protobuf.nano.InternalNano.$value_desc_type$);\n"
+    "}\n");
 }
 
 void MapFieldGenerator::
 GenerateSerializedSizeCode(io::Printer* printer) const {
+  printer->Print(variables_,
+    "if (this.$name$ != null) {\n"
+    "  size += com.google.protobuf.nano.MapUtil.Internal.computeMapFieldSize(\n"
+    "    this.$name$, $number$,\n"
+    "  com.google.protobuf.nano.InternalNano.$key_desc_type$,\n"
+    "  com.google.protobuf.nano.InternalNano.$value_desc_type$);\n"
+    "}\n");
 }
 
 void MapFieldGenerator::