From 40998462d655249212a19c92eb3c58eacc1e626b Mon Sep 17 00:00:00 2001
From: Eric Salo <salo@google.com>
Date: Wed, 12 Oct 2022 11:30:38 -0700
Subject: [PATCH] add upb_MtDataEncoder_EncodeExtension()

We need to sharpen the distinction between messages and extensions in the mini
descriptor encoder, so split the code paths for each.

PiperOrigin-RevId: 480675339
---
 upb/mini_table.c           | 10 ++++++++++
 upb/mini_table.h           |  5 +++++
 upb/mini_table.hpp         |  8 ++++++++
 upb/reflection/field_def.c |  8 ++------
 upbc/file_layout.cc        |  5 ++---
 5 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/upb/mini_table.c b/upb/mini_table.c
index a1d03e7a27..13dbf873ec 100644
--- a/upb/mini_table.c
+++ b/upb/mini_table.c
@@ -188,6 +188,16 @@ char* upb_MtDataEncoder_PutModifier(upb_MtDataEncoder* e, char* ptr,
   return ptr;
 }
 
+char* upb_MtDataEncoder_EncodeExtension(upb_MtDataEncoder* e, char* ptr,
+                                        upb_FieldType type, uint32_t field_num,
+                                        uint64_t field_mod) {
+  upb_MtDataEncoderInternal* in = upb_MtDataEncoder_GetInternal(e, ptr);
+  in->state.msg_state.msg_modifiers = 0;
+  in->state.msg_state.last_field_num = 0;
+  in->state.msg_state.oneof_state = kUpb_OneofState_NotStarted;
+  return upb_MtDataEncoder_PutField(e, ptr, type, field_num, field_mod);
+}
+
 char* upb_MtDataEncoder_StartMessage(upb_MtDataEncoder* e, char* ptr,
                                      uint64_t msg_mod) {
   upb_MtDataEncoderInternal* in = upb_MtDataEncoder_GetInternal(e, ptr);
diff --git a/upb/mini_table.h b/upb/mini_table.h
index 219fa895f2..543a5c1449 100644
--- a/upb/mini_table.h
+++ b/upb/mini_table.h
@@ -121,6 +121,11 @@ char* upb_MtDataEncoder_PutEnumValue(upb_MtDataEncoder* e, char* ptr,
                                      uint32_t val);
 char* upb_MtDataEncoder_EndEnum(upb_MtDataEncoder* e, char* ptr);
 
+// Encodes an entire mini descriptor for one extension.
+char* upb_MtDataEncoder_EncodeExtension(upb_MtDataEncoder* e, char* ptr,
+                                        upb_FieldType type, uint32_t field_num,
+                                        uint64_t field_mod);
+
 /** upb_MiniTable *************************************************************/
 
 typedef enum {
diff --git a/upb/mini_table.hpp b/upb/mini_table.hpp
index 57de9e9e3e..5ba814dcc8 100644
--- a/upb/mini_table.hpp
+++ b/upb/mini_table.hpp
@@ -80,6 +80,14 @@ class MtDataEncoder {
 
   const std::string& data() const { return appender_.data(); }
 
+  bool EncodeExtension(upb_FieldType type, uint32_t field_num,
+                       uint64_t field_mod) {
+    return appender_([=](char* buf) {
+      return upb_MtDataEncoder_EncodeExtension(&encoder_, buf, type, field_num,
+                                               field_mod);
+    });
+  }
+
  private:
   class StringAppender {
    public:
diff --git a/upb/reflection/field_def.c b/upb/reflection/field_def.c
index 7ebee56de9..7484af6d33 100644
--- a/upb/reflection/field_def.c
+++ b/upb/reflection/field_def.c
@@ -818,16 +818,12 @@ bool upb_FieldDef_MiniDescriptorEncode(const upb_FieldDef* f, upb_Arena* a,
   upb_DescState s;
   _upb_DescState_Init(&s);
 
-  if (!_upb_DescState_Grow(&s, a)) return false;
-  s.ptr = upb_MtDataEncoder_StartMessage(&s.e, s.ptr, 0);
-
   const int number = upb_FieldDef_Number(f);
   const uint64_t modifiers = _upb_FieldDef_Modifiers(f);
 
   if (!_upb_DescState_Grow(&s, a)) return false;
-  s.ptr = upb_MtDataEncoder_PutField(&s.e, s.ptr, f->type_, number, modifiers);
-
-  if (!_upb_DescState_Grow(&s, a)) return false;
+  s.ptr = upb_MtDataEncoder_EncodeExtension(&s.e, s.ptr, f->type_, number,
+                                            modifiers);
   *s.ptr = '\0';
 
   out->data = s.buf;
diff --git a/upbc/file_layout.cc b/upbc/file_layout.cc
index f9dbac39f0..f9b2a56196 100644
--- a/upbc/file_layout.cc
+++ b/upbc/file_layout.cc
@@ -260,9 +260,8 @@ void FilePlatformLayout::BuildExtensions(const protobuf::FileDescriptor* fd) {
   upb::Status status;
   for (const auto* f : sorted) {
     upb::MtDataEncoder e;
-    e.StartMessage(0);
-    e.PutField(static_cast<upb_FieldType>(f->type()), f->number(),
-               GetFieldModifiers(f));
+    e.EncodeExtension(static_cast<upb_FieldType>(f->type()), f->number(),
+                      GetFieldModifiers(f));
     upb_MiniTable_Extension& ext = extension_map_[f];
     upb_MiniTable_Sub sub;
     // The extendee may be from another file, so we build a temporary MiniTable