From cbfb6991f3bee8a61f7d99de5c48407d5a90f5d9 Mon Sep 17 00:00:00 2001
From: Protobuf Team Bot <protobuf-github-bot@google.com>
Date: Thu, 7 Mar 2024 12:18:55 -0800
Subject: [PATCH] Add api to trace message allocations. Initialize minitable
 tracing field for MtDecoder Add api for non C language runtimes to initialize
 dynamic mini table name.

PiperOrigin-RevId: 613664457
---
 upb/message/internal/message.c    | 17 +++++++++++++++++
 upb/message/internal/message.h    | 10 ++++++++++
 upb/mini_descriptor/decode.c      |  7 ++++++-
 upb/mini_table/internal/message.h |  8 ++++++++
 4 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/upb/message/internal/message.c b/upb/message/internal/message.c
index 594629a8ad..e5eb16e6e9 100644
--- a/upb/message/internal/message.c
+++ b/upb/message/internal/message.c
@@ -57,3 +57,20 @@ bool UPB_PRIVATE(_upb_Message_Realloc)(struct upb_Message* msg, size_t need,
   UPB_ASSERT(in->ext_begin - in->unknown_end >= need);
   return true;
 }
+
+#if UPB_TRACING_ENABLED
+static void (*_new_message_trace_handler)(const upb_MiniTable*,
+                                          const upb_Arena*);
+
+void upb_Message_SetNewMessageTraceHandler(
+    void (*new_message_trace_handler)(const upb_MiniTable*, const upb_Arena*)) {
+  _new_message_trace_handler = new_message_trace_handler;
+}
+
+void upb_Message_LogNewMessage(const upb_MiniTable* mini_table,
+                               const upb_Arena* arena) {
+  if (_new_message_trace_handler) {
+    _new_message_trace_handler(mini_table, arena);
+  }
+}
+#endif
diff --git a/upb/message/internal/message.h b/upb/message/internal/message.h
index dcabf56c05..7f3686a28a 100644
--- a/upb/message/internal/message.h
+++ b/upb/message/internal/message.h
@@ -59,9 +59,19 @@ typedef struct upb_Message_Internal {
   //   char data[size - sizeof(upb_Message_Internal)];
 } upb_Message_Internal;
 
+#ifdef UPB_TRACING_ENABLED
+void upb_Message_SetNewMessageTraceHandler(
+    void (*newMessageTraceHandler)(const upb_MiniTable*, const upb_Arena*));
+void upb_Message_LogNewMessage(const upb_MiniTable* mini_table,
+                               const upb_Arena* arena);
+#endif
+
 // Inline version upb_Message_New(), for internal use.
 UPB_INLINE struct upb_Message* _upb_Message_New(const upb_MiniTable* m,
                                                 upb_Arena* a) {
+#ifdef UPB_TRACING_ENABLED
+  upb_Message_LogNewMessage(m, a);
+#endif
   const int size = m->UPB_PRIVATE(size);
   struct upb_Message* msg = (struct upb_Message*)upb_Arena_Malloc(a, size);
   if (UPB_UNLIKELY(!msg)) return NULL;
diff --git a/upb/mini_descriptor/decode.c b/upb/mini_descriptor/decode.c
index 2029aff69a..d355112f57 100644
--- a/upb/mini_descriptor/decode.c
+++ b/upb/mini_descriptor/decode.c
@@ -48,7 +48,7 @@ typedef enum {
   kUpb_LayoutItemType_Max = kUpb_LayoutItemType_Field,
 } upb_LayoutItemType;
 
-#define kUpb_LayoutItem_IndexSentinel ((uint16_t)-1)
+#define kUpb_LayoutItem_IndexSentinel ((uint16_t) - 1)
 
 typedef struct {
   // Index of the corresponding field.  When this is a oneof field, the field's
@@ -737,6 +737,11 @@ static upb_MiniTable* upb_MtDecoder_DoBuildMiniTableWithBuf(
   decoder->table->UPB_PRIVATE(dense_below) = 0;
   decoder->table->UPB_PRIVATE(table_mask) = -1;
   decoder->table->UPB_PRIVATE(required_count) = 0;
+#if UPB_TRACING_ENABLED
+  // MiniTables built from MiniDescriptors will not be able to vend the message
+  // name unless it is explicitly set with upb_MiniTable_SetFullName().
+  decoder->table->UPB_PRIVATE(full_name) = 0;
+#endif
 
   // Strip off and verify the version tag.
   if (!len--) goto done;
diff --git a/upb/mini_table/internal/message.h b/upb/mini_table/internal/message.h
index cac7a68793..d605189483 100644
--- a/upb/mini_table/internal/message.h
+++ b/upb/mini_table/internal/message.h
@@ -57,6 +57,7 @@ struct upb_MiniTable {
   uint8_t UPB_PRIVATE(dense_below);
   uint8_t UPB_PRIVATE(table_mask);
   uint8_t UPB_PRIVATE(required_count);  // Required fields have the low hasbits.
+
 #ifdef UPB_TRACING_ENABLED
   const char* UPB_PRIVATE(full_name);
 #endif
@@ -164,6 +165,13 @@ UPB_INLINE const char* upb_MiniTable_FullName(
     const struct upb_MiniTable* mini_table) {
   return mini_table->UPB_PRIVATE(full_name);
 }
+// Initializes tracing proto name from language runtimes that construct
+// mini tables dynamically at runtime. The runtime is responsible for passing
+// controlling lifetime of name such as storing in same arena as mini_table.
+UPB_INLINE const char* upb_MiniTable_SetFullName(
+    struct upb_MiniTable* mini_table, char* full_name) {
+  mini_table->UPB_PRIVATE(full_name) = full_name;
+}
 #endif
 
 #ifdef __cplusplus