Use noncontiguous unknown fields API in upb python

PiperOrigin-RevId: 696258062
protocolbuffers · Nov 13, 2024 · ae488a0 · ae488a0
1 parent 6aaf5c9
commit ae488a0
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 24 deletions.
diff --git a/python/message.c b/python/message.c
@@ -12,8 +12,10 @@
 #include "python/extension_dict.h"
 #include "python/map.h"
 #include "python/repeated.h"
+#include "upb/base/string_view.h"
 #include "upb/message/compare.h"
 #include "upb/message/copy.h"
+#include "upb/message/message.h"
 #include "upb/reflection/def.h"
 #include "upb/reflection/message.h"
 #include "upb/text/encode.h"
@@ -575,9 +577,7 @@ static bool PyUpb_Message_IsEmpty(const upb_Message* msg,
   upb_MessageValue val;
   if (upb_Message_Next(msg, m, ext_pool, &f, &val, &iter)) return false;
 
-  size_t len;
-  (void)upb_Message_GetUnknown(msg, &len);
-  return len == 0;
+  return !upb_Message_HasUnknown(msg);
 }
 
 static bool PyUpb_Message_IsEqual(PyUpb_Message* m1, PyObject* _m2) {

diff --git a/python/unknown_fields.c b/python/unknown_fields.c
@@ -9,6 +9,7 @@
 
 #include "python/message.h"
 #include "python/protobuf.h"
+#include "upb/message/message.h"
 #include "upb/wire/eps_copy_input_stream.h"
 #include "upb/wire/reader.h"
 #include "upb/wire/types.h"
@@ -32,6 +33,7 @@ PyUpb_UnknownFieldSet* PyUpb_UnknownFieldSet_NewBare(void) {
   PyUpb_ModuleState* s = PyUpb_ModuleState_Get();
   PyUpb_UnknownFieldSet* self =
       (void*)PyType_GenericAlloc(s->unknown_fields_type, 0);
+  self->fields = PyList_New(0);
   return self;
 }
 
@@ -118,7 +120,6 @@ static const char* PyUpb_UnknownFieldSet_BuildMessageSetItem(
 static const char* PyUpb_UnknownFieldSet_BuildMessageSet(
     PyUpb_UnknownFieldSet* self, upb_EpsCopyInputStream* stream,
     const char* ptr) {
-  self->fields = PyList_New(0);
   while (!upb_EpsCopyInputStream_IsDone(stream, &ptr)) {
     uint32_t tag;
     ptr = upb_WireReader_ReadTag(ptr, &tag);
@@ -135,7 +136,6 @@ static const char* PyUpb_UnknownFieldSet_BuildMessageSet(
 
 err:
   Py_DECREF(self->fields);
-  self->fields = NULL;
   return NULL;
 }
 
@@ -199,7 +199,6 @@ static const char* PyUpb_UnknownFieldSet_Build(PyUpb_UnknownFieldSet* self,
                                                const char* ptr,
                                                int group_number) {
   PyUpb_ModuleState* s = PyUpb_ModuleState_Get();
-  self->fields = PyList_New(0);
   while (!upb_EpsCopyInputStream_IsDone(stream, &ptr)) {
     uint32_t tag;
     ptr = upb_WireReader_ReadTag(ptr, &tag);
@@ -228,7 +227,6 @@ static const char* PyUpb_UnknownFieldSet_Build(PyUpb_UnknownFieldSet* self,
 
 err:
   Py_DECREF(self->fields);
-  self->fields = NULL;
   return NULL;
 }
 
@@ -246,24 +244,25 @@ static PyObject* PyUpb_UnknownFieldSet_New(PyTypeObject* type, PyObject* args,
   upb_Message* msg = PyUpb_Message_GetIfReified(py_msg);
   if (!msg) return &self->ob_base;
 
-  size_t size;
-  const char* ptr = upb_Message_GetUnknown(msg, &size);
-  if (size == 0) return &self->ob_base;
-
-  upb_EpsCopyInputStream stream;
-  upb_EpsCopyInputStream_Init(&stream, &ptr, size, true);
-  const upb_MessageDef* msgdef = PyUpb_Message_GetMsgdef(py_msg);
-
-  bool ok;
-  if (upb_MessageDef_IsMessageSet(msgdef)) {
-    ok = PyUpb_UnknownFieldSet_BuildMessageSet(self, &stream, ptr) != NULL;
-  } else {
-    ok = PyUpb_UnknownFieldSet_Build(self, &stream, ptr, -1) != NULL;
-  }
+  uintptr_t iter = kUpb_Message_UnknownBegin;
+  upb_StringView view;
+  while (upb_Message_NextUnknown(msg, &view, &iter)) {
+    const char* ptr = view.data;
+    upb_EpsCopyInputStream stream;
+    upb_EpsCopyInputStream_Init(&stream, &ptr, view.size, true);
+    const upb_MessageDef* msgdef = PyUpb_Message_GetMsgdef(py_msg);
+
+    bool ok;
+    if (upb_MessageDef_IsMessageSet(msgdef)) {
+      ok = PyUpb_UnknownFieldSet_BuildMessageSet(self, &stream, ptr) != NULL;
+    } else {
+      ok = PyUpb_UnknownFieldSet_Build(self, &stream, ptr, -1) != NULL;
+    }
 
-  if (!ok) {
-    Py_DECREF(&self->ob_base);
-    return NULL;
+    if (!ok) {
+      Py_DECREF(&self->ob_base);
+      return NULL;
+    }
   }
 
   return &self->ob_base;

diff --git a/upb/message/message.c b/upb/message/message.c
@@ -70,6 +70,14 @@ bool upb_Message_NextUnknown(const upb_Message* msg, upb_StringView* data,
   return false;
 }
 
+bool upb_Message_HasUnknown(const upb_Message* msg) {
+  const upb_Message_Internal* in = UPB_PRIVATE(_upb_Message_GetInternal)(msg);
+  if (in) {
+    return in->unknown_end > message_overhead;
+  }
+  return false;
+}
+
 const char* upb_Message_GetUnknown(const upb_Message* msg, size_t* len) {
   upb_Message_Internal* in = UPB_PRIVATE(_upb_Message_GetInternal)(msg);
   if (in) {

diff --git a/upb/message/message.h b/upb/message/message.h
@@ -48,6 +48,8 @@ UPB_API upb_Message* upb_Message_New(const upb_MiniTable* m, upb_Arena* arena);
 bool upb_Message_NextUnknown(const upb_Message* msg, upb_StringView* data,
                              uintptr_t* iter);
 
+bool upb_Message_HasUnknown(const upb_Message* msg);
+
 // Returns a reference to the message's unknown data.
 const char* upb_Message_GetUnknown(const upb_Message* msg, size_t* len);