Refactor TextFormat.Printer.MapEntryAdapter to handle more entry type…

…s. This will properly sort map keys for DynamicMessage maps. PiperOrigin-RevId: 718402359
protocolbuffers · Jan 22, 2025 · 84769d3 · 84769d3
1 parent cb77b0e
commit 84769d3
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 40 deletions.
diff --git a/java/core/src/main/java/com/google/protobuf/TextFormat.java b/java/core/src/main/java/com/google/protobuf/TextFormat.java
@@ -459,68 +459,57 @@ private void printField(
     }
 
     /** An adapter class that can take a {@link MapEntry} and returns its key and entry. */
-    private static class MapEntryAdapter implements Comparable<MapEntryAdapter> {
+    static class MapEntryAdapter implements Comparable<MapEntryAdapter> {
       private Object entry;
-
-      @SuppressWarnings({"rawtypes"})
-      private MapEntry mapEntry;
-
-      private final FieldDescriptor.JavaType fieldType;
+      private Message messageEntry;
+      private final FieldDescriptor keyField;
 
       MapEntryAdapter(Object entry, FieldDescriptor fieldDescriptor) {
-        if (entry instanceof MapEntry) {
-          this.mapEntry = (MapEntry) entry;
+        if (entry instanceof Message) {
+          this.messageEntry = (Message) entry;
         } else {
           this.entry = entry;
         }
-        this.fieldType = extractFieldType(fieldDescriptor);
-      }
-
-      private static FieldDescriptor.JavaType extractFieldType(FieldDescriptor fieldDescriptor) {
-        return fieldDescriptor.getMessageType().getFields().get(0).getJavaType();
+        this.keyField = fieldDescriptor.getMessageType().findFieldByName("key");
       }
 
       Object getKey() {
-        if (mapEntry != null) {
-          return mapEntry.getKey();
+        if (messageEntry != null && keyField != null) {
+          return messageEntry.getField(keyField);
         }
         return null;
       }
 
       Object getEntry() {
-        if (mapEntry != null) {
-          return mapEntry;
+        if (messageEntry != null) {
+          return messageEntry;
         }
         return entry;
       }
 
       @Override
       public int compareTo(MapEntryAdapter b) {
-        if (getKey() == null || b.getKey() == null) {
-          logger.info("Invalid key for map field.");
+        Object aKey = getKey();
+        Object bKey = b.getKey();
+        if (aKey == null && bKey == null) {
+          return 0;
+        } else if (aKey == null) {
           return -1;
-        }
-        switch (fieldType) {
-          case BOOLEAN:
-            return ((Boolean) getKey()).compareTo((Boolean) b.getKey());
-          case LONG:
-            return ((Long) getKey()).compareTo((Long) b.getKey());
-          case INT:
-            return ((Integer) getKey()).compareTo((Integer) b.getKey());
-          case STRING:
-            String aString = (String) getKey();
-            String bString = (String) b.getKey();
-            if (aString == null && bString == null) {
+        } else if (bKey == null) {
+          return 1;
+        } else {
+          switch (keyField.getJavaType()) {
+            case BOOLEAN:
+              return ((Boolean) aKey).compareTo((Boolean) bKey);
+            case LONG:
+              return ((Long) aKey).compareTo((Long) bKey);
+            case INT:
+              return ((Integer) aKey).compareTo((Integer) bKey);
+            case STRING:
+              return ((String) aKey).compareTo((String) bKey);
+            default:
               return 0;
-            } else if (aString == null && bString != null) {
-              return -1;
-            } else if (aString != null && bString == null) {
-              return 1;
-            } else {
-              return aString.compareTo(bString);
-            }
-          default:
-            return 0;
+          }
         }
       }
     }

diff --git a/java/core/src/test/java/com/google/protobuf/TextFormatTest.java b/java/core/src/test/java/com/google/protobuf/TextFormatTest.java
@@ -1976,6 +1976,58 @@ public void testMapOverwrite() throws Exception {
     }
   }
 
+  @Test
+  public void testMapDynamicMessage() throws Exception {
+    TestMap message =
+        TestMap.newBuilder()
+            .putInt32ToStringField(30, "cherry")
+            .putInt32ToStringField(10, "apple")
+            .putInt32ToStringField(20, "banana")
+            .build();
+    DynamicMessage dynamic =
+        DynamicMessage.parseFrom(
+            TestMap.getDescriptor(), message.toByteString(), ExtensionRegistry.getEmptyRegistry());
+    assertThat(TextFormat.printer().printToString(dynamic))
+        .isEqualTo(TextFormat.printer().printToString(message));
+  }
+
+  @Test
+  public void testMapKeyAdapterComparison() throws Exception {
+    FieldDescriptor int32ToStringField =
+        TestMap.getDescriptor().findFieldByNumber(TestMap.INT32_TO_STRING_FIELD_FIELD_NUMBER);
+    TestMap map =
+        TestMap.newBuilder()
+            .putInt32ToStringField(10, "apple")
+            .putInt32ToStringField(20, "banana")
+            .build();
+    TextFormat.Printer.MapEntryAdapter nullEntry =
+        new TextFormat.Printer.MapEntryAdapter(null, int32ToStringField);
+    TextFormat.Printer.MapEntryAdapter entry1 =
+        new TextFormat.Printer.MapEntryAdapter(
+            map.getRepeatedField(int32ToStringField, 0), int32ToStringField);
+    TextFormat.Printer.MapEntryAdapter entry2 =
+        new TextFormat.Printer.MapEntryAdapter(
+            map.getRepeatedField(int32ToStringField, 1), int32ToStringField);
+    assertThat(nullEntry).isEquivalentAccordingToCompareTo(nullEntry);
+    assertThat(entry1).isEquivalentAccordingToCompareTo(entry1);
+    assertThat(entry2).isEquivalentAccordingToCompareTo(entry2);
+    assertThat(nullEntry).isLessThan(entry1);
+    assertThat(entry1).isGreaterThan(nullEntry);
+    assertThat(nullEntry).isLessThan(entry2);
+    assertThat(entry2).isGreaterThan(nullEntry);
+    assertThat(entry1).isLessThan(entry2);
+    assertThat(entry2).isGreaterThan(entry1);
+  }
+
+  @Test
+  public void testMapKeyAdapterComparisonInvalidType() throws Exception {
+    TextFormat.Printer.MapEntryAdapter invalidEntry =
+        new TextFormat.Printer.MapEntryAdapter(
+            TestOneof2.NestedMessage.getDefaultInstance(),
+            TestOneof2.getDescriptor().findFieldByNumber(TestOneof2.FOO_MESSAGE_FIELD_NUMBER));
+    assertThat(invalidEntry).isEquivalentAccordingToCompareTo(invalidEntry);
+  }
+
   // =======================================================================
   // test location information