Support printing non-canonical extensions in the text encoder.

PiperOrigin-RevId: 967565448
diff --git a/upb/message/BUILD b/upb/message/BUILD
index c4aa476..c3a7004 100644
--- a/upb/message/BUILD
+++ b/upb/message/BUILD
@@ -76,6 +76,7 @@
         "//third_party/upb/wasm_abi/upb:__pkg__",
         "//upb:__pkg__",
         "//upb/test:__pkg__",
+        "//upb/text:__pkg__",
         "//upb/wire:__pkg__",
         "//upb/wire/decode_fast:__pkg__",
         "//video/youtube/utils/elements/client/component_type:__pkg__",
diff --git a/upb/text/BUILD b/upb/text/BUILD
index be52852..25e3e02 100644
--- a/upb/text/BUILD
+++ b/upb/text/BUILD
@@ -81,6 +81,11 @@
         "//upb/lex",
         "//upb/message",
         "//upb/message:internal",
+        "//upb/message:iterator",
+        "//upb/message:message_unknowns",
+        "//upb/message:types",
+        "//upb/mini_table",
+        "//upb/mini_table:internal",
         "//upb/port",
         "//upb/wire:eps_copy_input_stream",
         "//upb/wire:reader",
diff --git a/upb/text/debug_string.c b/upb/text/debug_string.c
index 8584f57..89edee0 100644
--- a/upb/text/debug_string.c
+++ b/upb/text/debug_string.c
@@ -7,202 +7,16 @@
 
 #include "upb/text/debug_string.h"
 
-#include <inttypes.h>
-#include <stdarg.h>
 #include <stddef.h>
-#include <stdint.h>
-#include <stdio.h>
-#include <string.h>
 
-#include "upb/base/descriptor_constants.h"
-#include "upb/message/array.h"
-#include "upb/message/internal/iterator.h"
-#include "upb/message/internal/map_entry.h"
 #include "upb/message/internal/map_sorter.h"
-#include "upb/message/map.h"
 #include "upb/message/message.h"
-#include "upb/message/value.h"
-#include "upb/mini_table/extension.h"
-#include "upb/mini_table/field.h"
-#include "upb/mini_table/internal/field.h"
-#include "upb/mini_table/internal/message.h"
 #include "upb/mini_table/message.h"
 #include "upb/text/internal/encode.h"
 
 // Must be last.
 #include "upb/port/def.inc"
 
-static void _upb_MessageDebugString(txtenc* e, const upb_Message* msg,
-                                    const upb_MiniTable* mt);
-
-static void _upb_FieldDebugString(txtenc* e, upb_MessageValue val,
-                                  const upb_MiniTableField* f,
-                                  const upb_MiniTable* mt, const char* label,
-                                  const upb_MiniTableExtension* ext) {
-  UPB_PRIVATE(_upb_TextEncode_Indent)(e);
-  const upb_CType ctype = upb_MiniTableField_CType(f);
-  const bool is_ext = upb_MiniTableField_IsExtension(f);
-  char number[10];  // A 32-bit integer can hold up to 10 digits.
-  snprintf(number, sizeof(number), "%" PRIu32, upb_MiniTableField_Number(f));
-  // label is to pass down whether we're dealing with a "key" of a map or
-  // a "value" of a map.
-  if (!label) label = number;
-
-  if (is_ext) {
-    UPB_PRIVATE(_upb_TextEncode_Printf)(e, "[%s]", label);
-  } else {
-    UPB_PRIVATE(_upb_TextEncode_Printf)(e, "%s", label);
-  }
-
-  if (ctype == kUpb_CType_Message) {
-    UPB_PRIVATE(_upb_TextEncode_Printf)(e, " {");
-    UPB_PRIVATE(_upb_TextEncode_EndField)(e);
-    e->indent_depth++;
-    const upb_MiniTable* subm = ext ? upb_MiniTableExtension_GetSubMessage(ext)
-                                    : upb_MiniTable_SubMessage(f);
-    _upb_MessageDebugString(e, val.msg_val, subm);
-    e->indent_depth--;
-    UPB_PRIVATE(_upb_TextEncode_Indent)(e);
-    UPB_PRIVATE(_upb_TextEncode_PutStr)(e, "}");
-    UPB_PRIVATE(_upb_TextEncode_EndField)(e);
-    return;
-  }
-
-  UPB_PRIVATE(_upb_TextEncode_Printf)(e, ": ");
-
-  if (ctype ==
-      kUpb_CType_Enum) {  // Enum has to be processed separately because of
-                          // divergent behavior between encoders
-    UPB_PRIVATE(_upb_TextEncode_Printf)(e, "%" PRId32, val.int32_val);
-  } else {
-    UPB_PRIVATE(_upb_TextEncode_Scalar)(e, val, ctype);
-  }
-
-  UPB_PRIVATE(_upb_TextEncode_EndField)(e);
-}
-
-/*
- * Arrays print as simple repeated elements, eg.
- *
- *    5: 1
- *    5: 2
- *    5: 3
- */
-static void _upb_ArrayDebugString(txtenc* e, const upb_Array* arr,
-                                  const upb_MiniTableField* f,
-                                  const upb_MiniTable* mt,
-                                  const upb_MiniTableExtension* ext) {
-  for (size_t i = 0, n = upb_Array_Size(arr); i < n; i++) {
-    _upb_FieldDebugString(e, upb_Array_Get(arr, i), f, mt, NULL, ext);
-  }
-}
-
-static void _upb_MapEntryDebugString(txtenc* e, upb_MessageValue key,
-                                     upb_MessageValue val,
-                                     const upb_MiniTableField* f,
-                                     const upb_MiniTable* mt) {
-  const upb_MiniTable* entry = upb_MiniTable_SubMessage(f);
-  const upb_MiniTableField* key_f = upb_MiniTable_MapKey(entry);
-  const upb_MiniTableField* val_f = upb_MiniTable_MapValue(entry);
-
-  UPB_PRIVATE(_upb_TextEncode_Indent)(e);
-  UPB_PRIVATE(_upb_TextEncode_Printf)(e, "%u {", upb_MiniTableField_Number(f));
-  UPB_PRIVATE(_upb_TextEncode_EndField)(e);
-  e->indent_depth++;
-
-  _upb_FieldDebugString(e, key, key_f, entry, "key", NULL);
-  _upb_FieldDebugString(e, val, val_f, entry, "value", NULL);
-
-  e->indent_depth--;
-  UPB_PRIVATE(_upb_TextEncode_Indent)(e);
-  UPB_PRIVATE(_upb_TextEncode_PutStr)(e, "}");
-  UPB_PRIVATE(_upb_TextEncode_EndField)(e);
-}
-
-/*
- * Maps print as messages of key/value, etc.
- *
- *    1 {
- *      key: "abc"
- *      value: 123
- *    }
- *    2 {
- *      key: "def"
- *      value: 456
- *    }
- */
-static void _upb_MapDebugString(txtenc* e, const upb_Map* map,
-                                const upb_MiniTableField* f,
-                                const upb_MiniTable* mt) {
-  if (e->options & UPB_TXTENC_NOSORT) {
-    size_t iter = kUpb_Map_Begin;
-    upb_MessageValue key, val;
-    while (upb_Map_Next(map, &key, &val, &iter)) {
-      _upb_MapEntryDebugString(e, key, val, f, mt);
-    }
-  } else {
-    if (upb_Map_Size(map) == 0) return;
-
-    const upb_MiniTable* entry = upb_MiniTable_SubMessage(f);
-    const upb_MiniTableField* key_f = upb_MiniTable_GetFieldByIndex(entry, 0);
-    _upb_sortedmap sorted;
-    upb_MapEntry ent;
-
-    _upb_mapsorter_pushmap(&e->sorter, upb_MiniTableField_Type(key_f), map,
-                           &sorted);
-    while (_upb_sortedmap_next(&e->sorter, map, &sorted, &ent)) {
-      upb_MessageValue key, val;
-      memcpy(&key, &ent.k, sizeof(key));
-      memcpy(&val, &ent.v, sizeof(val));
-      _upb_MapEntryDebugString(e, key, val, f, mt);
-    }
-    _upb_mapsorter_popmap(&e->sorter, &sorted);
-  }
-}
-
-static void _upb_MessageDebugString(txtenc* e, const upb_Message* msg,
-                                    const upb_MiniTable* mt) {
-  size_t iter = kUpb_BaseField_Begin;
-  const upb_MiniTableField* f;
-  upb_MessageValue val;
-
-  // Base fields will be printed out first, followed by extension fields, and
-  // finally unknown fields.
-
-  while (UPB_PRIVATE(_upb_Message_NextBaseField)(msg, mt, &f, &val, &iter)) {
-    if (upb_MiniTableField_IsMap(f)) {
-      _upb_MapDebugString(e, val.map_val, f, mt);
-    } else if (upb_MiniTableField_IsArray(f)) {
-      // ext set to NULL as we're not dealing with extensions yet
-      _upb_ArrayDebugString(e, val.array_val, f, mt, NULL);
-    } else {
-      // ext set to NULL as we're not dealing with extensions yet
-      // label set to NULL as we're not currently working with a MapEntry
-      _upb_FieldDebugString(e, val, f, mt, NULL, NULL);
-    }
-  }
-
-  const upb_MiniTableExtension* ext;
-  upb_MessageValue val_ext;
-  iter = kUpb_Message_ExtensionBegin;
-  while (upb_Message_NextExtension(msg, &ext, &val_ext, &iter)) {
-    const upb_MiniTableField* f = &ext->UPB_PRIVATE(field);
-    // It is not sufficient to only pass |f| as we lose valuable information
-    // about sub-messages. It is required that we pass |ext|.
-    if (upb_MiniTableField_IsMap(f)) {
-      UPB_UNREACHABLE();  // Maps cannot be extensions.
-      break;
-    } else if (upb_MiniTableField_IsArray(f)) {
-      _upb_ArrayDebugString(e, val_ext.array_val, f, mt, ext);
-    } else {
-      // label set to NULL as we're not currently working with a MapEntry
-      _upb_FieldDebugString(e, val_ext, f, mt, NULL, ext);
-    }
-  }
-
-  UPB_PRIVATE(_upb_TextEncode_ParseUnknown)(e, msg);
-}
-
 size_t upb_DebugString(const upb_Message* msg, const upb_MiniTable* mt,
                        int options, char* buf, size_t size) {
   txtenc e;
@@ -216,7 +30,7 @@
   e.ext_pool = NULL;
   _upb_mapsorter_init(&e.sorter);
 
-  _upb_MessageDebugString(&e, msg, mt);
+  UPB_PRIVATE(_upb_MessageDebugString)(&e, msg, mt);
   _upb_mapsorter_destroy(&e.sorter);
   return UPB_PRIVATE(_upb_TextEncode_Nullz)(&e, size);
 }
diff --git a/upb/text/internal/encode.c b/upb/text/internal/encode.c
index ddcf6ba..bd74836 100644
--- a/upb/text/internal/encode.c
+++ b/upb/text/internal/encode.c
@@ -8,14 +8,28 @@
 #include "upb/text/internal/encode.h"
 
 #include <inttypes.h>
+#include <stdarg.h>
 #include <stddef.h>
 #include <stdint.h>
+#include <stdio.h>
+#include <string.h>
 
 #include "upb/base/descriptor_constants.h"
 #include "upb/base/string_view.h"
 #include "upb/lex/round_trip.h"
 #include "upb/message/array.h"
+#include "upb/message/internal/iterator.h"
+#include "upb/message/internal/map_entry.h"
+#include "upb/message/internal/map_sorter.h"
+#include "upb/message/map.h"
 #include "upb/message/message.h"
+#include "upb/message/unknown_fields.h"
+#include "upb/message/value.h"
+#include "upb/mini_table/extension.h"
+#include "upb/mini_table/field.h"
+#include "upb/mini_table/internal/field.h"
+#include "upb/mini_table/internal/message.h"
+#include "upb/mini_table/message.h"
 #include "upb/text/options.h"
 #include "upb/wire/eps_copy_input_stream.h"
 #include "upb/wire/reader.h"
@@ -31,6 +45,16 @@
     }              \
   } while (0)
 
+static void _upb_FieldDebugString(txtenc* e, upb_MessageValue val,
+                                  const upb_MiniTableField* f,
+                                  const upb_MiniTable* mt, const char* label,
+                                  const upb_MiniTableExtension* ext);
+
+static void _upb_ArrayDebugString(txtenc* e, const upb_Array* arr,
+                                  const upb_MiniTableField* f,
+                                  const upb_MiniTable* mt,
+                                  const upb_MiniTableExtension* ext);
+
 /*
  * Unknown fields are printed by number.
  *
@@ -138,14 +162,30 @@
   if ((e->options & UPB_TXTENC_SKIPUNKNOWN) != 0) return;
 
   uintptr_t iter = kUpb_Message_UnknownBegin;
-  upb_StringView view;
-  while (upb_Message_NextUnknown(msg, &view, &iter)) {
-    char* start = e->ptr;
-    upb_EpsCopyInputStream stream;
-    upb_EpsCopyInputStream_Init(&stream, &view.data, view.size);
-    if (!UPB_PRIVATE(_upb_TextEncode_Unknown)(e, view.data, &stream, -1)) {
-      /* Unknown failed to parse, back up and don't print it at all. */
-      e->ptr = start;
+  upb_MessageUnknown unknown;
+  while (upb_Message_NextUnknown2(msg, &unknown, &iter)) {
+    if (unknown.type == kUpb_MessageUnknownType_StringView) {
+      upb_StringView view = unknown.value.bytes;
+      char* start = e->ptr;
+      upb_EpsCopyInputStream stream;
+      upb_EpsCopyInputStream_Init(&stream, &view.data, view.size);
+      if (!UPB_PRIVATE(_upb_TextEncode_Unknown)(e, view.data, &stream, -1)) {
+        /* Unknown failed to parse, back up and don't print it at all. */
+        e->ptr = start;
+      }
+    } else {
+      UPB_ASSERT(unknown.type == kUpb_MessageUnknownType_NonCanonicalExtension);
+      const struct upb_Extension* ext_struct = unknown.value.extension;
+      const upb_MiniTableExtension* ext = ext_struct->ext;
+      upb_MessageValue val_ext = ext_struct->data;
+      const upb_MiniTableField* f = upb_MiniTableExtension_ToField(ext);
+      const upb_MiniTable* mt = upb_MiniTableExtension_Extendee(ext);
+      UPB_ASSERT(!upb_MiniTableField_IsMap(f));
+      if (upb_MiniTableField_IsArray(f)) {
+        _upb_ArrayDebugString(e, val_ext.array_val, f, mt, ext);
+      } else {
+        _upb_FieldDebugString(e, val_ext, f, mt, NULL, ext);
+      }
     }
   }
 }
@@ -194,3 +234,171 @@
       UPB_UNREACHABLE();
   }
 }
+
+static void _upb_FieldDebugString(txtenc* e, upb_MessageValue val,
+                                  const upb_MiniTableField* f,
+                                  const upb_MiniTable* mt, const char* label,
+                                  const upb_MiniTableExtension* ext) {
+  UPB_PRIVATE(_upb_TextEncode_Indent)(e);
+  const upb_CType ctype = upb_MiniTableField_CType(f);
+  const bool is_ext = upb_MiniTableField_IsExtension(f);
+  char number[10];  // A 32-bit integer can hold up to 10 digits.
+  snprintf(number, sizeof(number), "%" PRIu32, upb_MiniTableField_Number(f));
+  // label is to pass down whether we're dealing with a "key" of a map or
+  // a "value" of a map.
+  if (!label) label = number;
+
+  if (is_ext) {
+    UPB_PRIVATE(_upb_TextEncode_Printf)(e, "[%s]", label);
+  } else {
+    UPB_PRIVATE(_upb_TextEncode_Printf)(e, "%s", label);
+  }
+
+  if (ctype == kUpb_CType_Message) {
+    UPB_PRIVATE(_upb_TextEncode_Printf)(e, " {");
+    UPB_PRIVATE(_upb_TextEncode_EndField)(e);
+    e->indent_depth++;
+    const upb_MiniTable* subm = ext ? upb_MiniTableExtension_GetSubMessage(ext)
+                                    : upb_MiniTable_SubMessage(f);
+    UPB_PRIVATE(_upb_MessageDebugString)(e, val.msg_val, subm);
+    e->indent_depth--;
+    UPB_PRIVATE(_upb_TextEncode_Indent)(e);
+    UPB_PRIVATE(_upb_TextEncode_PutStr)(e, "}");
+    UPB_PRIVATE(_upb_TextEncode_EndField)(e);
+    return;
+  }
+
+  UPB_PRIVATE(_upb_TextEncode_Printf)(e, ": ");
+
+  if (ctype ==
+      kUpb_CType_Enum) {  // Enum has to be processed separately because of
+                          // divergent behavior between encoders
+    UPB_PRIVATE(_upb_TextEncode_Printf)(e, "%" PRId32, val.int32_val);
+  } else {
+    UPB_PRIVATE(_upb_TextEncode_Scalar)(e, val, ctype);
+  }
+
+  UPB_PRIVATE(_upb_TextEncode_EndField)(e);
+}
+
+/*
+ * Arrays print as simple repeated elements, eg.
+ *
+ *    5: 1
+ *    5: 2
+ *    5: 3
+ */
+static void _upb_ArrayDebugString(txtenc* e, const upb_Array* arr,
+                                  const upb_MiniTableField* f,
+                                  const upb_MiniTable* mt,
+                                  const upb_MiniTableExtension* ext) {
+  for (size_t i = 0, n = upb_Array_Size(arr); i < n; i++) {
+    _upb_FieldDebugString(e, upb_Array_Get(arr, i), f, mt, NULL, ext);
+  }
+}
+
+static void _upb_MapEntryDebugString(txtenc* e, upb_MessageValue key,
+                                     upb_MessageValue val,
+                                     const upb_MiniTableField* f,
+                                     const upb_MiniTable* mt) {
+  const upb_MiniTable* entry = upb_MiniTable_SubMessage(f);
+  const upb_MiniTableField* key_f = upb_MiniTable_MapKey(entry);
+  const upb_MiniTableField* val_f = upb_MiniTable_MapValue(entry);
+
+  UPB_PRIVATE(_upb_TextEncode_Indent)(e);
+  UPB_PRIVATE(_upb_TextEncode_Printf)(e, "%u {", upb_MiniTableField_Number(f));
+  UPB_PRIVATE(_upb_TextEncode_EndField)(e);
+  e->indent_depth++;
+
+  _upb_FieldDebugString(e, key, key_f, entry, "key", NULL);
+  _upb_FieldDebugString(e, val, val_f, entry, "value", NULL);
+
+  e->indent_depth--;
+  UPB_PRIVATE(_upb_TextEncode_Indent)(e);
+  UPB_PRIVATE(_upb_TextEncode_PutStr)(e, "}");
+  UPB_PRIVATE(_upb_TextEncode_EndField)(e);
+}
+
+/*
+ * Maps print as messages of key/value, etc.
+ *
+ *    1 {
+ *      key: "abc"
+ *      value: 123
+ *    }
+ *    2 {
+ *      key: "def"
+ *      value: 456
+ *    }
+ */
+static void _upb_MapDebugString(txtenc* e, const upb_Map* map,
+                                const upb_MiniTableField* f,
+                                const upb_MiniTable* mt) {
+  if (e->options & UPB_TXTENC_NOSORT) {
+    size_t iter = kUpb_Map_Begin;
+    upb_MessageValue key, val;
+    while (upb_Map_Next(map, &key, &val, &iter)) {
+      _upb_MapEntryDebugString(e, key, val, f, mt);
+    }
+  } else {
+    if (upb_Map_Size(map) == 0) return;
+
+    const upb_MiniTable* entry = upb_MiniTable_SubMessage(f);
+    const upb_MiniTableField* key_f = upb_MiniTable_GetFieldByIndex(entry, 0);
+    _upb_sortedmap sorted;
+    upb_MapEntry ent;
+
+    _upb_mapsorter_pushmap(&e->sorter, upb_MiniTableField_Type(key_f), map,
+                           &sorted);
+    while (_upb_sortedmap_next(&e->sorter, map, &sorted, &ent)) {
+      upb_MessageValue key, val;
+      memcpy(&key, &ent.k, sizeof(key));
+      memcpy(&val, &ent.v, sizeof(val));
+      _upb_MapEntryDebugString(e, key, val, f, mt);
+    }
+    _upb_mapsorter_popmap(&e->sorter, &sorted);
+  }
+}
+
+void UPB_PRIVATE(_upb_MessageDebugString)(txtenc* e, const upb_Message* msg,
+                                          const upb_MiniTable* mt) {
+  size_t iter = kUpb_BaseField_Begin;
+  const upb_MiniTableField* f;
+  upb_MessageValue val;
+
+  // Base fields will be printed out first, followed by extension fields, and
+  // finally unknown fields.
+
+  while (UPB_PRIVATE(_upb_Message_NextBaseField)(msg, mt, &f, &val, &iter)) {
+    if (upb_MiniTableField_IsMap(f)) {
+      _upb_MapDebugString(e, val.map_val, f, mt);
+    } else if (upb_MiniTableField_IsArray(f)) {
+      // ext set to NULL as we're not dealing with extensions yet
+      _upb_ArrayDebugString(e, val.array_val, f, mt, NULL);
+    } else {
+      // ext set to NULL as we're not dealing with extensions yet
+      // label set to NULL as we're not currently working with a MapEntry
+      _upb_FieldDebugString(e, val, f, mt, NULL, NULL);
+    }
+  }
+
+  const upb_MiniTableExtension* ext;
+  upb_MessageValue val_ext;
+  iter = kUpb_Message_ExtensionBegin;
+  while (upb_Message_NextExtension(msg, &ext, &val_ext, &iter)) {
+    const upb_MiniTableField* f = &ext->UPB_PRIVATE(field);
+    // It is not sufficient to only pass |f| as we lose valuable information
+    // about sub-messages. It is required that we pass |ext|.
+    if (upb_MiniTableField_IsMap(f)) {
+      UPB_UNREACHABLE();  // Maps cannot be extensions.
+      break;
+    } else if (upb_MiniTableField_IsArray(f)) {
+      _upb_ArrayDebugString(e, val_ext.array_val, f, mt, ext);
+    } else {
+      // label set to NULL as we're not currently working with a MapEntry
+      _upb_FieldDebugString(e, val_ext, f, mt, NULL, ext);
+    }
+  }
+
+  UPB_PRIVATE(_upb_TextEncode_ParseUnknown)(e, msg);
+}
diff --git a/upb/text/internal/encode.h b/upb/text/internal/encode.h
index a204d67..6de9662 100644
--- a/upb/text/internal/encode.h
+++ b/upb/text/internal/encode.h
@@ -16,6 +16,7 @@
 #include "upb/message/array.h"
 #include "upb/message/internal/map_sorter.h"
 #include "upb/message/message.h"
+#include "upb/mini_table/internal/message.h"
 #include "upb/port/vsnprintf_compat.h"
 #include "upb/text/options.h"
 #include "upb/wire/eps_copy_input_stream.h"
@@ -235,6 +236,9 @@
 void UPB_PRIVATE(_upb_TextEncode_ParseUnknown)(txtenc* e,
                                                const upb_Message* msg);
 
+void UPB_PRIVATE(_upb_MessageDebugString)(txtenc* e, const upb_Message* msg,
+                                          const struct upb_MiniTable* mt);
+
 // Must not be called for ctype = kUpb_CType_Enum, as they require different
 // handling depending on whether or not we're doing reflection-based encoding.
 void UPB_PRIVATE(_upb_TextEncode_Scalar)(txtenc* e, upb_MessageValue val,