Fixed the initialization sequence for MiniTables.

Prior to this CL we were allocating a MiniTable for each message and then overwriting it later.  This could lead to an inconsistent state, and is unnecessary.  This CL adds an extra phase to initialization so that the MiniTable is assigned only one time for each message.

PiperOrigin-RevId: 507617479

upb/reflection/def_pool.c

@@ -86,7 +86,7 @@ err:
 }
 
 bool _upb_DefPool_InsertExt(upb_DefPool* s, const upb_MiniTableExtension* ext,
-                            upb_FieldDef* f) {
+                            const upb_FieldDef* f) {
   return upb_inttable_insert(&s->exts, (uintptr_t)ext, upb_value_constptr(f),
                              s->arena);
 }

upb/reflection/def_pool_internal.h

@@ -43,7 +43,7 @@ size_t _upb_DefPool_BytesLoaded(const upb_DefPool* s);
 upb_ExtensionRegistry* _upb_DefPool_ExtReg(const upb_DefPool* s);
 
 bool _upb_DefPool_InsertExt(upb_DefPool* s, const upb_MiniTableExtension* ext,
-                            upb_FieldDef* f);
+                            const upb_FieldDef* f);
 bool _upb_DefPool_InsertSym(upb_DefPool* s, upb_StringView sym, upb_value v,
                             upb_Status* status);
 bool _upb_DefPool_LookupSym(const upb_DefPool* s, const char* sym, size_t size,

upb/reflection/field_def.c

@@ -697,19 +697,6 @@ static void _upb_FieldDef_CreateNotExt(upb_DefBuilder* ctx, const char* prefix,
   }
 
   _upb_MessageDef_InsertField(ctx, m, f);
-
-  if (!ctx->layout) return;
-
-  const upb_MiniTable* mt = upb_MessageDef_MiniTable(m);
-  const upb_MiniTableField* fields = mt->fields;
-  for (int i = 0; i < mt->field_count; i++) {
-    if (fields[i].number == f->number_) {
-      f->layout_index = i;
-      return;
-    }
-  }
-
-  UPB_ASSERT(false);  // It should be impossible to reach this point.
 }
 
 upb_FieldDef* _upb_Extensions_New(
@@ -744,7 +731,13 @@ upb_FieldDef* _upb_FieldDefs_New(
 
     _upb_FieldDef_CreateNotExt(ctx, prefix, protos[i], m, f);
     f->index_ = i;
-    if (!ctx->layout) f->layout_index = i;
+    if (!ctx->layout) {
+      // Speculate that the def fields are sorted.  We will always sort the
+      // MiniTable fields, so if defs are sorted then indices will match.
+      //
+      // If this is incorrect, we will overwrite later.
+      f->layout_index = i;
+    }
 
     const uint32_t current = f->number_;
     if (previous > current) *is_sorted = false;
@@ -804,6 +797,9 @@ static int _upb_FieldDef_Compare(const void* p1, const void* p2) {
   return (v1 < v2) ? -1 : (v1 > v2);
 }
 
+// _upb_FieldDefs_Sorted() is mostly a pure function of its inputs, but has one
+// critical side effect that we depend on: it sets layout_index appropriately
+// for non-sorted lists of fields.
 const upb_FieldDef** _upb_FieldDefs_Sorted(const upb_FieldDef* f, int n,
                                            upb_Arena* a) {
   // TODO(salo): Replace this arena alloc with a persistent scratch buffer.
@@ -861,7 +857,10 @@ static void resolve_extension(upb_DefBuilder* ctx, const char* prefix,
         "field number %u in extension %s has no extension range in message %s",
         (unsigned)f->number_, f->full_name, upb_MessageDef_FullName(m));
   }
+}
 
+void _upb_FieldDef_BuildMiniTableExtension(upb_DefBuilder* ctx,
+                                           const upb_FieldDef* f) {
   const upb_MiniTableExtension* ext = _upb_FieldDef_ExtensionMiniTable(f);
 
   if (ctx->layout) {
@@ -880,8 +879,8 @@ static void resolve_extension(upb_DefBuilder* ctx, const char* prefix,
       sub.subenum = _upb_EnumDef_MiniTable(f->sub.enumdef);
     }
     bool ok2 = upb_MiniTableExtension_Build(desc.data, desc.size, mut_ext,
-                                            upb_MessageDef_MiniTable(m), sub,
-                                            ctx->status);
+                                            upb_MessageDef_MiniTable(f->msgdef),
+                                            sub, ctx->status);
     if (!ok2) _upb_DefBuilder_Errf(ctx, "Could not build extension mini table");
   }
 

upb/reflection/field_def_internal.h

@@ -47,6 +47,8 @@ int _upb_FieldDef_LayoutIndex(const upb_FieldDef* f);
 uint64_t _upb_FieldDef_Modifiers(const upb_FieldDef* f);
 void _upb_FieldDef_Resolve(upb_DefBuilder* ctx, const char* prefix,
                            upb_FieldDef* f);
+void _upb_FieldDef_BuildMiniTableExtension(upb_DefBuilder* ctx,
+                                           const upb_FieldDef* f);
 
 // Allocate and initialize an array of |n| extensions (field defs).
 upb_FieldDef* _upb_Extensions_New(

upb/reflection/file_def.c

@@ -347,11 +347,19 @@ void _upb_FileDef_Create(upb_DefBuilder* ctx,
     _upb_FieldDef_Resolve(ctx, file->package, f);
   }
 
-  if (!ctx->layout) {
-    for (int i = 0; i < file->top_lvl_msg_count; i++) {
-      upb_MessageDef* m = (upb_MessageDef*)upb_FileDef_TopLevelMessage(file, i);
-      _upb_MessageDef_LinkMiniTable(ctx, m);
-    }
+  for (int i = 0; i < file->top_lvl_msg_count; i++) {
+    upb_MessageDef* m = (upb_MessageDef*)upb_FileDef_TopLevelMessage(file, i);
+    _upb_MessageDef_CreateMiniTable(ctx, (upb_MessageDef*)m);
+  }
+
+  for (int i = 0; i < file->top_lvl_ext_count; i++) {
+    upb_FieldDef* f = (upb_FieldDef*)upb_FileDef_TopLevelExtension(file, i);
+    _upb_FieldDef_BuildMiniTableExtension(ctx, f);
+  }
+
+  for (int i = 0; i < file->top_lvl_msg_count; i++) {
+    upb_MessageDef* m = (upb_MessageDef*)upb_FileDef_TopLevelMessage(file, i);
+    _upb_MessageDef_LinkMiniTable(ctx, m);
   }
 
   if (file->ext_count) {

upb/reflection/message_def.c

@@ -349,6 +349,8 @@ bool upb_MessageDef_IsMessageSet(const upb_MessageDef* m) {
 static upb_MiniTable* _upb_MessageDef_MakeMiniTable(upb_DefBuilder* ctx,
                                                     const upb_MessageDef* m) {
   upb_StringView desc;
+  // Note: this will assign layout_index for fields, so upb_FieldDef_MiniTable()
+  // is safe to call only after this call.
   bool ok = upb_MessageDef_MiniDescriptorEncode(m, ctx->tmp_arena, &desc);
   if (!ok) _upb_DefBuilder_OomErr(ctx);
 
@@ -368,22 +370,6 @@ void _upb_MessageDef_Resolve(upb_DefBuilder* ctx, upb_MessageDef* m) {
     _upb_FieldDef_Resolve(ctx, m->full_name, f);
   }
 
-  if (!ctx->layout) {
-    m->layout = _upb_MessageDef_MakeMiniTable(ctx, m);
-  }
-
-#ifndef NDEBUG
-  for (int i = 0; i < m->field_count; i++) {
-    const upb_FieldDef* f = upb_MessageDef_Field(m, i);
-    const int layout_index = _upb_FieldDef_LayoutIndex(f);
-    UPB_ASSERT(layout_index < m->layout->field_count);
-    const upb_MiniTableField* mt_f = &m->layout->fields[layout_index];
-    UPB_ASSERT(upb_FieldDef_Type(f) == upb_MiniTableField_Type(mt_f));
-    UPB_ASSERT(upb_FieldDef_HasPresence(f) ==
-               upb_MiniTableField_HasPresence(mt_f));
-  }
-#endif
-
   m->in_message_set = false;
   for (int i = 0; i < upb_MessageDef_NestedExtensionCount(m); i++) {
     upb_FieldDef* ext = (upb_FieldDef*)upb_MessageDef_NestedExtension(m, i);
@@ -446,8 +432,40 @@ void _upb_MessageDef_InsertField(upb_DefBuilder* ctx, upb_MessageDef* m,
   if (!ok) _upb_DefBuilder_OomErr(ctx);
 }
 
+void _upb_MessageDef_CreateMiniTable(upb_DefBuilder* ctx, upb_MessageDef* m) {
+  if (ctx->layout) {
+    /* create_fielddef() below depends on this being set. */
+    UPB_ASSERT(ctx->msg_count < ctx->layout->msg_count);
+    m->layout = ctx->layout->msgs[ctx->msg_count++];
+    UPB_ASSERT(m->field_count == m->layout->field_count);
+    // We don't need the result of this call, but it will assign layout_index
+    // for all the fields in O(n lg n) time.
+    _upb_FieldDefs_Sorted(m->fields, m->field_count, ctx->tmp_arena);
+  } else {
+    /* Allocate now (to allow cross-linking), populate later. */
+    m->layout = _upb_MessageDef_MakeMiniTable(ctx, m);
+  }
+
+  for (int i = 0; i < m->nested_msg_count; i++) {
+    upb_MessageDef* nested =
+        (upb_MessageDef*)upb_MessageDef_NestedMessage(m, i);
+    _upb_MessageDef_CreateMiniTable(ctx, nested);
+  }
+}
+
 void _upb_MessageDef_LinkMiniTable(upb_DefBuilder* ctx,
                                    const upb_MessageDef* m) {
+  for (int i = 0; i < upb_MessageDef_NestedExtensionCount(m); i++) {
+    const upb_FieldDef* ext = upb_MessageDef_NestedExtension(m, i);
+    _upb_FieldDef_BuildMiniTableExtension(ctx, ext);
+  }
+
+  for (int i = 0; i < m->nested_msg_count; i++) {
+    _upb_MessageDef_LinkMiniTable(ctx, upb_MessageDef_NestedMessage(m, i));
+  }
+
+  if (ctx->layout) return;
+
   for (int i = 0; i < m->field_count; i++) {
     const upb_FieldDef* f = upb_MessageDef_Field(m, i);
     const upb_MessageDef* sub_m = upb_FieldDef_MessageSubDef(f);
@@ -475,9 +493,17 @@ void _upb_MessageDef_LinkMiniTable(upb_DefBuilder* ctx,
     }
   }
 
-  for (int i = 0; i < m->nested_msg_count; i++) {
-    _upb_MessageDef_LinkMiniTable(ctx, upb_MessageDef_NestedMessage(m, i));
+#ifndef NDEBUG
+  for (int i = 0; i < m->field_count; i++) {
+    const upb_FieldDef* f = upb_MessageDef_Field(m, i);
+    const int layout_index = _upb_FieldDef_LayoutIndex(f);
+    UPB_ASSERT(layout_index < m->layout->field_count);
+    const upb_MiniTableField* mt_f = &m->layout->fields[layout_index];
+    UPB_ASSERT(upb_FieldDef_Type(f) == upb_MiniTableField_Type(mt_f));
+    UPB_ASSERT(upb_FieldDef_HasPresence(f) ==
+               upb_MiniTableField_HasPresence(mt_f));
   }
+#endif
 }
 
 static uint64_t _upb_MessageDef_Modifiers(const upb_MessageDef* m) {
@@ -628,16 +654,6 @@ static void create_msgdef(upb_DefBuilder* ctx, const char* prefix,
   ok = upb_strtable_init(&m->ntof, n_oneof + n_field, ctx->arena);
   if (!ok) _upb_DefBuilder_OomErr(ctx);
 
-  if (ctx->layout) {
-    /* create_fielddef() below depends on this being set. */
-    UPB_ASSERT(ctx->msg_count < ctx->layout->msg_count);
-    m->layout = ctx->layout->msgs[ctx->msg_count++];
-    UPB_ASSERT(n_field == m->layout->field_count);
-  } else {
-    /* Allocate now (to allow cross-linking), populate later. */
-    m->layout = _upb_DefBuilder_Alloc(ctx, sizeof(*m->layout));
-  }
-
   UPB_DEF_SET_OPTIONS(m->opts, DescriptorProto, MessageOptions, msg_proto);
 
   m->oneof_count = n_oneof;

upb/reflection/message_def_internal.h

@@ -44,6 +44,7 @@ bool _upb_MessageDef_Insert(upb_MessageDef* m, const char* name, size_t size,
 void _upb_MessageDef_InsertField(upb_DefBuilder* ctx, upb_MessageDef* m,
                                  const upb_FieldDef* f);
 bool _upb_MessageDef_IsValidExtensionNumber(const upb_MessageDef* m, int n);
+void _upb_MessageDef_CreateMiniTable(upb_DefBuilder* ctx, upb_MessageDef* m);
 void _upb_MessageDef_LinkMiniTable(upb_DefBuilder* ctx,
                                    const upb_MessageDef* m);
 void _upb_MessageDef_Resolve(upb_DefBuilder* ctx, upb_MessageDef* m);