Add MapFieldBuilder and change codegen to generate it and the put{field}BuilderIfAbsent method.

PiperOrigin-RevId: 563738402

java/core/src/main/java/com/google/protobuf/MapFieldBuilder.java

@@ -58,8 +58,10 @@ public class MapFieldBuilder<
   /** nullable */
   Map<KeyT, MessageT> messageMap = null;
 
-  // messageList elements are always MapEntry<KeyT, MessageT>, but we need a List<Message> for
+  // We need a List<Message> for reflection.
-  // reflection.
+  //
+  // messageList elements are always MapEntry<KeyT, SomeT extends Message>, where SomeT and MessageT
+  // have the same descriptor (i.e. SomeT can be DynamicMessage)
   /** nullable */
   List<Message> messageList = null;
 
@@ -80,8 +82,15 @@ public class MapFieldBuilder<
   @SuppressWarnings("unchecked")
   private List<MapEntry<KeyT, MessageT>> getMapEntryList() {
     ArrayList<MapEntry<KeyT, MessageT>> list = new ArrayList<>(messageList.size());
+    Class<?> valueClass = converter.defaultEntry().getValue().getClass();
     for (Message entry : messageList) {
-      list.add((MapEntry<KeyT, MessageT>) entry);
+      MapEntry<KeyT, ?> typedEntry = (MapEntry<KeyT, ?>) entry;
+      if (valueClass.isInstance(typedEntry.getValue())) {
+        list.add((MapEntry<KeyT, MessageT>) typedEntry);
+      } else {
+        // This needs to use mergeFrom to allow MapEntry<KeyT, DynamicMessage> to be used.
+        list.add(converter.defaultEntry().toBuilder().mergeFrom(entry).build());
+      }
     }
     return list;
   }

java/core/src/test/java/com/google/protobuf/MapForProto2Test.java

@@ -333,7 +333,7 @@ public class MapForProto2Test {
     assertThat(builder.build().getInt32ToInt32Field()).isEqualTo(newMap(1, 2));
     try {
       intMap.put(2, 3);
-      assertWithMessage("expected exception").fail();
+      assertWithMessage("expected exception intMap").fail();
     } catch (UnsupportedOperationException e) {
       // expected
     }
@@ -347,7 +347,7 @@ public class MapForProto2Test {
         .isEqualTo(newMap(1, TestMap.EnumValue.BAR));
     try {
       enumMap.put(2, TestMap.EnumValue.FOO);
-      assertWithMessage("expected exception").fail();
+      assertWithMessage("expected exception enumMap").fail();
     } catch (UnsupportedOperationException e) {
       // expected
     }
@@ -361,7 +361,7 @@ public class MapForProto2Test {
     assertThat(builder.build().getInt32ToStringField()).isEqualTo(newMap(1, "1"));
     try {
       stringMap.put(2, "2");
-      assertWithMessage("expected exception").fail();
+      assertWithMessage("expected exception stringMap").fail();
     } catch (UnsupportedOperationException e) {
       // expected
     }
@@ -369,19 +369,16 @@ public class MapForProto2Test {
     builder.getMutableInt32ToStringField().put(2, "2");
     assertThat(builder.getInt32ToStringField()).isEqualTo(newMap(1, "1", 2, "2"));
 
+    // Message maps are handled differently, and don't freeze old mutable collections.
     Map<Integer, TestMap.MessageValue> messageMap = builder.getMutableInt32ToMessageField();
     messageMap.put(1, TestMap.MessageValue.getDefaultInstance());
     assertThat(builder.build().getInt32ToMessageField())
         .isEqualTo(newMap(1, TestMap.MessageValue.getDefaultInstance()));
-    try {
+    // Mutations on old mutable maps don't affect the builder state.
-      messageMap.put(2, TestMap.MessageValue.getDefaultInstance());
+    messageMap.put(2, TestMap.MessageValue.getDefaultInstance());
-      assertWithMessage("expected exception").fail();
+    assertThat(builder.getInt32ToMessageField()).isEqualTo(
-    } catch (UnsupportedOperationException e) {
+        newMap(1, TestMap.MessageValue.getDefaultInstance()));
-      // expected
+    builder.putInt32ToMessageField(2, TestMap.MessageValue.getDefaultInstance());
-    }
-    assertThat(builder.getInt32ToMessageField())
-        .isEqualTo(newMap(1, TestMap.MessageValue.getDefaultInstance()));
-    builder.getMutableInt32ToMessageField().put(2, TestMap.MessageValue.getDefaultInstance());
     assertThat(builder.getInt32ToMessageField()).isEqualTo(
         newMap(1, TestMap.MessageValue.getDefaultInstance(),
             2, TestMap.MessageValue.getDefaultInstance()));

java/core/src/test/java/com/google/protobuf/MapTest.java

@@ -332,7 +332,7 @@ public class MapTest {
     assertThat(builder.build().getInt32ToInt32Field()).isEqualTo(newMap(1, 2));
     try {
       intMap.put(2, 3);
-      assertWithMessage("expected exception").fail();
+      assertWithMessage("expected exception intMap").fail();
     } catch (UnsupportedOperationException e) {
       // expected
     }
@@ -346,7 +346,7 @@ public class MapTest {
         .isEqualTo(newMap(1, TestMap.EnumValue.BAR));
     try {
       enumMap.put(2, TestMap.EnumValue.FOO);
-      assertWithMessage("expected exception").fail();
+      assertWithMessage("expected exception enumMap").fail();
     } catch (UnsupportedOperationException e) {
       // expected
     }
@@ -360,7 +360,7 @@ public class MapTest {
     assertThat(builder.build().getInt32ToStringField()).isEqualTo(newMap(1, "1"));
     try {
       stringMap.put(2, "2");
-      assertWithMessage("expected exception").fail();
+      assertWithMessage("expected exception stringMap").fail();
     } catch (UnsupportedOperationException e) {
       // expected
     }
@@ -368,18 +368,15 @@ public class MapTest {
     builder.putInt32ToStringField(2, "2");
     assertThat(builder.getInt32ToStringField()).isEqualTo(newMap(1, "1", 2, "2"));
 
+    // Message maps are handled differently, and don't freeze old mutable collections.
     Map<Integer, TestMap.MessageValue> messageMap = builder.getMutableInt32ToMessageField();
     messageMap.put(1, TestMap.MessageValue.getDefaultInstance());
-    assertThat( builder.build().getInt32ToMessageField())
+    assertThat(builder.build().getInt32ToMessageField())
-        .isEqualTo(newMap(1, TestMap.MessageValue.getDefaultInstance()));
-    try {
-      messageMap.put(2, TestMap.MessageValue.getDefaultInstance());
-      assertWithMessage("expected exception").fail();
-    } catch (UnsupportedOperationException e) {
-      // expected
-    }
-    assertThat(builder.getInt32ToMessageField())
         .isEqualTo(newMap(1, TestMap.MessageValue.getDefaultInstance()));
+    // Mutations on old mutable maps don't affect the builder state.
+    messageMap.put(2, TestMap.MessageValue.getDefaultInstance());
+    assertThat(builder.getInt32ToMessageField()).isEqualTo(
+        newMap(1, TestMap.MessageValue.getDefaultInstance()));
     builder.putInt32ToMessageField(2, TestMap.MessageValue.getDefaultInstance());
     assertThat(builder.getInt32ToMessageField()).isEqualTo(
         newMap(1, TestMap.MessageValue.getDefaultInstance(),

src/google/protobuf/compiler/java/map_field.cc

@@ -30,9 +30,13 @@
 
 #include "google/protobuf/compiler/java/map_field.h"
 
+#include <string>
+
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
 #include "google/protobuf/compiler/java/context.h"
 #include "google/protobuf/compiler/java/doc_comment.h"
+#include "google/protobuf/compiler/java/field.h"
 #include "google/protobuf/compiler/java/helpers.h"
 #include "google/protobuf/compiler/java/name_resolver.h"
 #include "google/protobuf/io/printer.h"
@@ -76,138 +80,144 @@ std::string WireType(const FieldDescriptor* field) {
                       FieldTypeName(field->type()));
 }
 
-void SetMessageVariables(
+}  // namespace
+
+ImmutableMapFieldGenerator::ImmutableMapFieldGenerator(
     const FieldDescriptor* descriptor, int messageBitIndex, int builderBitIndex,
-    const FieldGeneratorInfo* info, Context* context,
+    Context* context)
-    absl::flat_hash_map<absl::string_view, std::string>* variables) {
+    : descriptor_(descriptor),
-  SetCommonFieldVariables(descriptor, info, variables);
+      message_bit_index_(messageBitIndex),
-  ClassNameResolver* name_resolver = context->GetNameResolver();
+      builder_bit_index_(builderBitIndex),
-
+      name_resolver_(context->GetNameResolver()),
-  (*variables)["type"] =
+      context_(context) {
-      name_resolver->GetImmutableClassName(descriptor->message_type());
+  SetMessageVariables(context->GetFieldGeneratorInfo(descriptor));
-  const FieldDescriptor* key = MapKeyField(descriptor);
+}
-  const FieldDescriptor* value = MapValueField(descriptor);
+
+void ImmutableMapFieldGenerator::SetMessageVariables(
+    const FieldGeneratorInfo* info) {
+  SetCommonFieldVariables(descriptor_, info, &variables_);
+  ClassNameResolver* name_resolver = context_->GetNameResolver();
+
+  variables_["type"] =
+      name_resolver->GetImmutableClassName(descriptor_->message_type());
+  const FieldDescriptor* key = MapKeyField(descriptor_);
+  const FieldDescriptor* value = MapValueField(descriptor_);
   const JavaType keyJavaType = GetJavaType(key);
   const JavaType valueJavaType = GetJavaType(value);
 
   // The code that generates the open-source version appears not to understand
   // #else, so we have an #ifndef instead.
   std::string pass_through_nullness =
-      context->options().opensource_runtime
+      context_->options().opensource_runtime
           ? "/* nullable */\n"
           : "@com.google.protobuf.Internal.ProtoPassThroughNullness ";
 
-  (*variables)["key_type"] = TypeName(key, name_resolver, false);
+  variables_["key_type"] = TypeName(key, name_resolver, false);
   std::string boxed_key_type = TypeName(key, name_resolver, true);
-  (*variables)["boxed_key_type"] = boxed_key_type;
+  variables_["boxed_key_type"] = boxed_key_type;
-  (*variables)["kt_key_type"] = KotlinTypeName(key, name_resolver);
+  variables_["kt_key_type"] = KotlinTypeName(key, name_resolver);
-  (*variables)["kt_value_type"] = KotlinTypeName(value, name_resolver);
+  variables_["kt_value_type"] = KotlinTypeName(value, name_resolver);
   // Used for calling the serialization function.
-  (*variables)["short_key_type"] =
+  variables_["short_key_type"] =
       boxed_key_type.substr(boxed_key_type.rfind('.') + 1);
-  (*variables)["key_wire_type"] = WireType(key);
+  variables_["key_wire_type"] = WireType(key);
-  (*variables)["key_default_value"] =
+  variables_["key_default_value"] =
-      DefaultValue(key, true, name_resolver, context->options());
+      DefaultValue(key, true, name_resolver, context_->options());
-  (*variables)["key_null_check"] =
+  variables_["key_null_check"] =
       IsReferenceType(keyJavaType)
           ? "if (key == null) { throw new NullPointerException(\"map key\"); }"
           : "";
-  (*variables)["value_null_check"] =
+  variables_["value_null_check"] =
       valueJavaType != JAVATYPE_ENUM && IsReferenceType(valueJavaType)
           ? "if (value == null) { "
             "throw new NullPointerException(\"map value\"); }"
           : "";
   if (valueJavaType == JAVATYPE_ENUM) {
     // We store enums as Integers internally.
-    (*variables)["value_type"] = "int";
+    variables_["value_type"] = "int";
-    variables->insert(
+    variables_.insert(
-        {"value_type_pass_through_nullness", (*variables)["value_type"]});
+        {"value_type_pass_through_nullness", variables_["value_type"]});
-    (*variables)["boxed_value_type"] = "java.lang.Integer";
+    variables_["boxed_value_type"] = "java.lang.Integer";
-    (*variables)["value_wire_type"] = WireType(value);
+    variables_["value_wire_type"] = WireType(value);
-    (*variables)["value_default_value"] =
+    variables_["value_default_value"] =
-        DefaultValue(value, true, name_resolver, context->options()) +
+        DefaultValue(value, true, name_resolver, context_->options()) +
         ".getNumber()";
 
-    (*variables)["value_enum_type"] = TypeName(value, name_resolver, false);
+    variables_["value_enum_type"] = TypeName(value, name_resolver, false);
 
-    variables->insert(
+    variables_.insert(
         {"value_enum_type_pass_through_nullness",
-         absl::StrCat(pass_through_nullness, (*variables)["value_enum_type"])});
+         absl::StrCat(pass_through_nullness, variables_["value_enum_type"])});
 
     if (SupportUnknownEnumValue(value)) {
       // Map unknown values to a special UNRECOGNIZED value if supported.
-      variables->insert(
+      variables_.insert(
           {"unrecognized_value",
-           absl::StrCat((*variables)["value_enum_type"], ".UNRECOGNIZED")});
+           absl::StrCat(variables_["value_enum_type"], ".UNRECOGNIZED")});
     } else {
       // Map unknown values to the default value if we don't have UNRECOGNIZED.
-      (*variables)["unrecognized_value"] =
+      variables_["unrecognized_value"] =
-          DefaultValue(value, true, name_resolver, context->options());
+          DefaultValue(value, true, name_resolver, context_->options());
     }
   } else {
-    (*variables)["value_type"] = TypeName(value, name_resolver, false);
+    variables_["value_type"] = TypeName(value, name_resolver, false);
 
-    variables->insert(
+    variables_.insert(
         {"value_type_pass_through_nullness",
          absl::StrCat(
              (IsReferenceType(valueJavaType) ? pass_through_nullness : ""),
-             (*variables)["value_type"])});
+             variables_["value_type"])});
 
-    (*variables)["boxed_value_type"] = TypeName(value, name_resolver, true);
+    variables_["boxed_value_type"] = TypeName(value, name_resolver, true);
-    (*variables)["value_wire_type"] = WireType(value);
+    variables_["value_wire_type"] = WireType(value);
-    (*variables)["value_default_value"] =
+    variables_["value_default_value"] =
-        DefaultValue(value, true, name_resolver, context->options());
+        DefaultValue(value, true, name_resolver, context_->options());
+  }
+
+  variables_.insert(
+      {"type_parameters", absl::StrCat(variables_["boxed_key_type"], ", ",
+                                       variables_["boxed_value_type"])});
+
+  if (valueJavaType == JAVATYPE_MESSAGE) {
+    variables_["value_interface_type"] =
+        absl::StrCat(variables_["boxed_value_type"], "OrBuilder");
+    variables_["value_builder_type"] =
+        absl::StrCat(variables_["boxed_value_type"], ".Builder");
+    variables_["builder_type_parameters"] = absl::StrJoin(
+        {variables_["boxed_key_type"], variables_["value_interface_type"],
+         variables_["boxed_value_type"], variables_["value_builder_type"]},
+        ", ");
   }
-  variables->insert(
-      {"type_parameters", absl::StrCat((*variables)["boxed_key_type"], ", ",
-                                       (*variables)["boxed_value_type"])});
   // TODO(birdo): Add @deprecated javadoc when generating javadoc is supported
   // by the proto compiler
-  (*variables)["deprecation"] =
+  variables_["deprecation"] =
-      descriptor->options().deprecated() ? "@java.lang.Deprecated " : "";
+      descriptor_->options().deprecated() ? "@java.lang.Deprecated " : "";
-  variables->insert(
+  variables_.insert(
       {"kt_deprecation",
-       descriptor->options().deprecated()
+       descriptor_->options().deprecated()
            ? absl::StrCat("@kotlin.Deprecated(message = \"Field ",
-                          (*variables)["name"], " is deprecated\") ")
+                          variables_["name"], " is deprecated\") ")
            : ""});
-  (*variables)["on_changed"] = "onChanged();";
+  variables_["on_changed"] = "onChanged();";
 
-  variables->insert(
+  variables_.insert(
-      {"default_entry", absl::StrCat((*variables)["capitalized_name"],
+      {"default_entry", absl::StrCat(variables_["capitalized_name"],
                                      "DefaultEntryHolder.defaultEntry")});
-  variables->insert({"map_field_parameter", (*variables)["default_entry"]});
+  variables_.insert({"map_field_parameter", variables_["default_entry"]});
-  (*variables)["descriptor"] = absl::StrCat(
+  variables_["descriptor"] = absl::StrCat(
-      name_resolver->GetImmutableClassName(descriptor->file()), ".internal_",
+      name_resolver->GetImmutableClassName(descriptor_->file()), ".internal_",
-      UniqueFileScopeIdentifier(descriptor->message_type()), "_descriptor, ");
+      UniqueFileScopeIdentifier(descriptor_->message_type()), "_descriptor, ");
-  (*variables)["ver"] = GeneratedCodeVersionSuffix();
+  variables_["ver"] = GeneratedCodeVersionSuffix();
-
+
-  (*variables)["get_has_field_bit_builder"] = GenerateGetBit(builderBitIndex);
+  variables_["get_has_field_bit_builder"] = GenerateGetBit(builder_bit_index_);
-  (*variables)["get_has_field_bit_from_local"] =
+  variables_["get_has_field_bit_from_local"] =
-      GenerateGetBitFromLocal(builderBitIndex);
+      GenerateGetBitFromLocal(builder_bit_index_);
-  (*variables)["set_has_field_bit_builder"] =
+  variables_["set_has_field_bit_builder"] =
-      absl::StrCat(GenerateSetBit(builderBitIndex), ";");
+      absl::StrCat(GenerateSetBit(builder_bit_index_), ";");
-  (*variables)["clear_has_field_bit_builder"] =
+  variables_["clear_has_field_bit_builder"] =
-      absl::StrCat(GenerateClearBit(builderBitIndex), ";");
+      absl::StrCat(GenerateClearBit(builder_bit_index_), ";");
-}
-
-}  // namespace
-
-ImmutableMapFieldGenerator::ImmutableMapFieldGenerator(
-    const FieldDescriptor* descriptor, int messageBitIndex, int builderBitIndex,
-    Context* context)
-    : descriptor_(descriptor),
-      message_bit_index_(messageBitIndex),
-      builder_bit_index_(builderBitIndex),
-      name_resolver_(context->GetNameResolver()),
-      context_(context) {
-  SetMessageVariables(descriptor, messageBitIndex, builderBitIndex,
-                      context->GetFieldGeneratorInfo(descriptor), context,
-                      &variables_);
 }
 
-ImmutableMapFieldGenerator::~ImmutableMapFieldGenerator() {}
-
 int ImmutableMapFieldGenerator::GetMessageBitIndex() const {
   return message_bit_index_;
 }
@@ -376,6 +386,9 @@ void ImmutableMapFieldGenerator::GenerateMembers(io::Printer* printer) const {
 
 void ImmutableMapFieldGenerator::GenerateBuilderMembers(
     io::Printer* printer) const {
+  if (GetJavaType(MapValueField(descriptor_)) == JAVATYPE_MESSAGE) {
+    return GenerateMessageMapBuilderMembers(printer);
+  }
   printer->Print(
       variables_,
       "private com.google.protobuf.MapField<\n"
@@ -743,6 +756,229 @@ void ImmutableMapFieldGenerator::GenerateMapGetters(
   }
 }
 
+void ImmutableMapFieldGenerator::GenerateMessageMapBuilderMembers(
+    io::Printer* printer) const {
+  printer->Print(
+      variables_,
+      "private static final class $capitalized_name$Converter implements "
+      "com.google.protobuf.MapFieldBuilder.Converter<$boxed_key_type$, "
+      "$value_interface_type$, $boxed_value_type$> "
+      "{\n");
+  {
+    auto i1 = printer->WithIndent();
+    printer->Print("@java.lang.Override\n");
+    printer->Print(
+        variables_,
+        "public $boxed_value_type$ build($value_interface_type$ val) {\n");
+    {
+      auto i2 = printer->WithIndent();
+      printer->Print(variables_,
+                     "if (val instanceof $boxed_value_type$) {"
+                     " return ($boxed_value_type$) val; }\n");
+      printer->Print(variables_,
+                     "return (($value_builder_type$) val).build();\n");
+    }
+    printer->Print("}\n\n");
+
+    printer->Print("@java.lang.Override\n");
+    printer->Print(variables_,
+                   "public com.google.protobuf.MapEntry<$boxed_key_type$, "
+                   "$boxed_value_type$> defaultEntry() {\n");
+    {
+      auto i2 = printer->WithIndent();
+      printer->Print(
+          variables_,
+          "return $capitalized_name$DefaultEntryHolder.defaultEntry;\n");
+    }
+    printer->Print("}\n");
+  }
+  printer->Print("};\n");
+  printer->Print(variables_,
+                 "private static final $capitalized_name$Converter "
+                 "$name$Converter = new $capitalized_name$Converter();\n\n");
+
+  printer->Print(
+      variables_,
+      "private com.google.protobuf.MapFieldBuilder<\n"
+      "    $builder_type_parameters$> $name$_;\n"
+      "$deprecation$private "
+      "com.google.protobuf.MapFieldBuilder<$builder_type_parameters$>\n"
+      "    internalGet$capitalized_name$() {\n"
+      "  if ($name$_ == null) {\n"
+      "    return new com.google.protobuf.MapFieldBuilder<>($name$Converter);\n"
+      "  }\n"
+      "  return $name$_;\n"
+      "}\n"
+      "$deprecation$private "
+      "com.google.protobuf.MapFieldBuilder<$builder_type_parameters$>\n"
+      "    internalGetMutable$capitalized_name$() {\n"
+      "  if ($name$_ == null) {\n"
+      "    $name$_ = new "
+      "com.google.protobuf.MapFieldBuilder<>($name$Converter);\n"
+      "  }\n"
+      "  $set_has_field_bit_builder$\n"
+      "  $on_changed$\n"
+      "  return $name$_;\n"
+      "}\n");
+  GenerateMessageMapGetters(printer);
+  printer->Print(
+      variables_,
+      "$deprecation$public Builder ${$clear$capitalized_name$$}$() {\n"
+      "  $clear_has_field_bit_builder$\n"
+      "  internalGetMutable$capitalized_name$().clear();\n"
+      "  return this;\n"
+      "}\n");
+  printer->Annotate("{", "}", descriptor_, Semantic::kSet);
+
+  WriteFieldDocComment(printer, descriptor_);
+  printer->Print(variables_,
+                 "$deprecation$public Builder ${$remove$capitalized_name$$}$(\n"
+                 "    $key_type$ key) {\n"
+                 "  $key_null_check$\n"
+                 "  internalGetMutable$capitalized_name$().ensureBuilderMap()\n"
+                 "      .remove(key);\n"
+                 "  return this;\n"
+                 "}\n");
+  printer->Annotate("{", "}", descriptor_, Semantic::kSet);
+
+  if (context_->options().opensource_runtime) {
+    printer->Print(
+        variables_,
+        "/**\n"
+        " * Use alternate mutation accessors instead.\n"
+        " */\n"
+        "@java.lang.Deprecated\n"
+        "public java.util.Map<$type_parameters$>\n"
+        "    ${$getMutable$capitalized_name$$}$() {\n"
+        "  $set_has_field_bit_builder$\n"
+        "  return internalGetMutable$capitalized_name$().ensureMessageMap();\n"
+        "}\n");
+    printer->Annotate("{", "}", descriptor_);
+  }
+
+  WriteFieldDocComment(printer, descriptor_);
+  printer->Print(variables_,
+                 "$deprecation$public Builder ${$put$capitalized_name$$}$(\n"
+                 "    $key_type$ key,\n"
+                 "    $value_type$ value) {\n"
+                 "  $key_null_check$\n"
+                 "  $value_null_check$\n"
+                 "  internalGetMutable$capitalized_name$().ensureBuilderMap()\n"
+                 "      .put(key, value);\n"
+                 "  $set_has_field_bit_builder$\n"
+                 "  return this;\n"
+                 "}\n");
+  printer->Annotate("{", "}", descriptor_, Semantic::kSet);
+
+  WriteFieldDocComment(printer, descriptor_);
+  printer->Print(
+      variables_,
+      "$deprecation$public Builder ${$putAll$capitalized_name$$}$(\n"
+      "    java.util.Map<$type_parameters$> values) {\n"
+      "  for (java.util.Map.Entry<$type_parameters$> e : values.entrySet()) {\n"
+      "    if (e.getKey() == null || e.getValue() == null) {\n"
+      "      throw new NullPointerException();\n"
+      "    }\n"
+      "  }\n"
+      "  internalGetMutable$capitalized_name$().ensureBuilderMap()\n"
+      "      .putAll(values);\n"
+      "  $set_has_field_bit_builder$\n"
+      "  return this;\n"
+      "}\n");
+  printer->Annotate("{", "}", descriptor_, Semantic::kSet);
+
+  WriteFieldDocComment(printer, descriptor_);
+  printer->Print(
+      variables_,
+      "$deprecation$public $value_builder_type$ "
+      "put$capitalized_name$BuilderIfAbsent(\n"
+      "    $key_type$ key) {\n"
+      "  java.util.Map<$boxed_key_type$, $value_interface_type$> builderMap = "
+      "internalGetMutable$capitalized_name$().ensureBuilderMap();\n"
+      "  $value_interface_type$ entry = builderMap.get(key);\n"
+      "  if (entry == null) { entry = "
+      "builderMap.put(key, $value_type$.newBuilder()); }\n"
+      "  if (entry instanceof $value_builder_type$) { return "
+      "($value_builder_type$) entry; }\n"
+      "  return ($value_builder_type$) builderMap.put(key, "
+      "(($value_type$) entry).toBuilder());\n"
+      "}\n");
+  printer->Annotate("{", "}", descriptor_, Semantic::kSet);
+}
+
+void ImmutableMapFieldGenerator::GenerateMessageMapGetters(
+    io::Printer* printer) const {
+  printer->Print(
+      variables_,
+      "$deprecation$public int ${$get$capitalized_name$Count$}$() {\n"
+      "  return internalGet$capitalized_name$().ensureBuilderMap().size();\n"
+      "}\n");
+  printer->Annotate("{", "}", descriptor_);
+
+  WriteFieldDocComment(printer, descriptor_);
+  printer->Print(
+      variables_,
+      "@java.lang.Override\n"
+      "$deprecation$public boolean ${$contains$capitalized_name$$}$(\n"
+      "    $key_type$ key) {\n"
+      "  $key_null_check$\n"
+      "  return "
+      "internalGet$capitalized_name$().ensureBuilderMap().containsKey(key);\n"
+      "}\n");
+  printer->Annotate("{", "}", descriptor_);
+  if (context_->options().opensource_runtime) {
+    printer->Print(variables_,
+                   "/**\n"
+                   " * Use {@link #get$capitalized_name$Map()} instead.\n"
+                   " */\n"
+                   "@java.lang.Override\n"
+                   "@java.lang.Deprecated\n"
+                   "public java.util.Map<$type_parameters$> "
+                   "${$get$capitalized_name$$}$() {\n"
+                   "  return get$capitalized_name$Map();\n"
+                   "}\n");
+    printer->Annotate("{", "}", descriptor_);
+  }
+  WriteFieldDocComment(printer, descriptor_);
+  printer->Print(variables_,
+                 "@java.lang.Override\n"
+                 "$deprecation$public java.util.Map<$type_parameters$> "
+                 "${$get$capitalized_name$Map$}$() {\n"
+                 "  return internalGet$capitalized_name$().getImmutableMap();\n"
+                 "}\n");
+  printer->Annotate("{", "}", descriptor_);
+  WriteFieldDocComment(printer, descriptor_);
+  printer->Print(
+      variables_,
+      "@java.lang.Override\n"
+      "$deprecation$public $value_type_pass_through_nullness$ "
+      "${$get$capitalized_name$OrDefault$}$(\n"
+      "    $key_type$ key,\n"
+      "    $value_type_pass_through_nullness$ defaultValue) {\n"
+      "  $key_null_check$\n"
+      "  java.util.Map<$boxed_key_type$, $value_interface_type$> map = "
+      "internalGetMutable$capitalized_name$().ensureBuilderMap();\n"
+      "  return map.containsKey(key) ? $name$Converter.build(map.get(key)) : "
+      "defaultValue;\n"
+      "}\n");
+  printer->Annotate("{", "}", descriptor_);
+  WriteFieldDocComment(printer, descriptor_);
+  printer->Print(
+      variables_,
+      "@java.lang.Override\n"
+      "$deprecation$public $value_type$ ${$get$capitalized_name$OrThrow$}$(\n"
+      "    $key_type$ key) {\n"
+      "  $key_null_check$\n"
+      "  java.util.Map<$boxed_key_type$, $value_interface_type$> map = "
+      "internalGetMutable$capitalized_name$().ensureBuilderMap();\n"
+      "  if (!map.containsKey(key)) {\n"
+      "    throw new java.lang.IllegalArgumentException();\n"
+      "  }\n"
+      "  return $name$Converter.build(map.get(key));\n"
+      "}\n");
+  printer->Annotate("{", "}", descriptor_);
+}
+
 void ImmutableMapFieldGenerator::GenerateKotlinDslMembers(
     io::Printer* printer) const {
   printer->Print(
@@ -808,7 +1044,8 @@ void ImmutableMapFieldGenerator::GenerateKotlinDslMembers(
       "@JvmName(\"putAll$kt_capitalized_name$\")\n"
       "public fun com.google.protobuf.kotlin.DslMap"
       "<$kt_key_type$, $kt_value_type$, ${$$kt_capitalized_name$Proxy$}$>\n"
-      "  .putAll(map: kotlin.collections.Map<$kt_key_type$, $kt_value_type$>) "
+      "  .putAll(map: kotlin.collections.Map<$kt_key_type$, "
+      "$kt_value_type$>) "
       "{\n"
       "     $kt_dsl_builder$.${$putAll$capitalized_name$$}$(map)\n"
       "   }\n");
@@ -852,6 +1089,15 @@ void ImmutableMapFieldGenerator::GenerateMergingCode(
 
 void ImmutableMapFieldGenerator::GenerateBuildingCode(
     io::Printer* printer) const {
+  if (GetJavaType(MapValueField(descriptor_)) == JAVATYPE_MESSAGE) {
+    printer->Print(
+        variables_,
+        "if ($get_has_field_bit_from_local$) {\n"
+        "  result.$name$_ = "
+        "internalGet$capitalized_name$().build($map_field_parameter$);\n"
+        "}\n");
+    return;
+  }
   printer->Print(variables_,
                  "if ($get_has_field_bit_from_local$) {\n"
                  "  result.$name$_ = internalGet$capitalized_name$();\n"
@@ -862,7 +1108,19 @@ void ImmutableMapFieldGenerator::GenerateBuildingCode(
 void ImmutableMapFieldGenerator::GenerateBuilderParsingCode(
     io::Printer* printer) const {
   const FieldDescriptor* value = MapValueField(descriptor_);
-  if (!SupportUnknownEnumValue(value) && GetJavaType(value) == JAVATYPE_ENUM) {
+  const JavaType type = GetJavaType(value);
+  if (type == JAVATYPE_MESSAGE) {
+    printer->Print(
+        variables_,
+        "com.google.protobuf.MapEntry<$type_parameters$>\n"
+        "$name$__ = input.readMessage(\n"
+        "    $default_entry$.getParserForType(), extensionRegistry);\n"
+        "internalGetMutable$capitalized_name$().ensureBuilderMap().put(\n"
+        "    $name$__.getKey(), $name$__.getValue());\n"
+        "$set_has_field_bit_builder$\n");
+    return;
+  }
+  if (!SupportUnknownEnumValue(value) && type == JAVATYPE_ENUM) {
     printer->Print(
         variables_,
         "com.google.protobuf.ByteString bytes = input.readBytes();\n"
@@ -875,16 +1133,15 @@ void ImmutableMapFieldGenerator::GenerateBuilderParsingCode(
         "      $name$__.getKey(), $name$__.getValue());\n"
         "  $set_has_field_bit_builder$\n"
         "}\n");
-  } else {
+    return;
-    printer->Print(
-        variables_,
-        "com.google.protobuf.MapEntry<$type_parameters$>\n"
-        "$name$__ = input.readMessage(\n"
-        "    $default_entry$.getParserForType(), extensionRegistry);\n"
-        "internalGetMutable$capitalized_name$().getMutableMap().put(\n"
-        "    $name$__.getKey(), $name$__.getValue());\n"
-        "$set_has_field_bit_builder$\n");
   }
+  printer->Print(variables_,
+                 "com.google.protobuf.MapEntry<$type_parameters$>\n"
+                 "$name$__ = input.readMessage(\n"
+                 "    $default_entry$.getParserForType(), extensionRegistry);\n"
+                 "internalGetMutable$capitalized_name$().getMutableMap().put(\n"
+                 "    $name$__.getKey(), $name$__.getValue());\n"
+                 "$set_has_field_bit_builder$\n");
 }
 void ImmutableMapFieldGenerator::GenerateSerializationCode(
     io::Printer* printer) const {

src/google/protobuf/compiler/java/map_field.h

@@ -43,7 +43,7 @@ class ImmutableMapFieldGenerator : public ImmutableFieldGenerator {
   explicit ImmutableMapFieldGenerator(const FieldDescriptor* descriptor,
                                       int messageBitIndex, int builderBitIndex,
                                       Context* context);
-  ~ImmutableMapFieldGenerator() override;
+  ~ImmutableMapFieldGenerator() override = default;
 
   // implements ImmutableFieldGenerator ---------------------------------------
   int GetMessageBitIndex() const override;
@@ -75,7 +75,10 @@ class ImmutableMapFieldGenerator : public ImmutableFieldGenerator {
   absl::flat_hash_map<absl::string_view, std::string> variables_;
   ClassNameResolver* name_resolver_;
   Context* context_;
+  void SetMessageVariables(const FieldGeneratorInfo* info);
   void GenerateMapGetters(io::Printer* printer) const;
+  void GenerateMessageMapBuilderMembers(io::Printer* printer) const;
+  void GenerateMessageMapGetters(io::Printer* printer) const;
 };
 
 }  // namespace java

src/google/protobuf/compiler/java/message.cc

@@ -825,7 +825,8 @@ void ImmutableMessageGenerator::GenerateDescriptorMethods(
     printer->Print(
         "@SuppressWarnings({\"rawtypes\"})\n"
         "@java.lang.Override\n"
-        "protected com.google.protobuf.MapField internalGetMapField(\n"
+        "protected com.google.protobuf.MapFieldReflectionAccessor "
+        "internalGetMapFieldReflection(\n"
         "    int number) {\n"
         "  switch (number) {\n");
     printer->Indent();

src/google/protobuf/compiler/java/message_builder.cc

@@ -247,7 +247,8 @@ void MessageBuilderGenerator::GenerateDescriptorMethods(io::Printer* printer) {
   if (!map_fields.empty()) {
     printer->Print(
         "@SuppressWarnings({\"rawtypes\"})\n"
-        "protected com.google.protobuf.MapField internalGetMapField(\n"
+        "protected com.google.protobuf.MapFieldReflectionAccessor "
+        "internalGetMapFieldReflection(\n"
         "    int number) {\n"
         "  switch (number) {\n");
     printer->Indent();
@@ -272,7 +273,8 @@ void MessageBuilderGenerator::GenerateDescriptorMethods(io::Printer* printer) {
         "}\n");
     printer->Print(
         "@SuppressWarnings({\"rawtypes\"})\n"
-        "protected com.google.protobuf.MapField internalGetMutableMapField(\n"
+        "protected com.google.protobuf.MapFieldReflectionAccessor "
+        "internalGetMutableMapFieldReflection(\n"
         "    int number) {\n"
         "  switch (number) {\n");
     printer->Indent();