implement upb_Map_Next() as the new upb_Map iterator

update upb to use upb_Map_Next()
remove upb_MapIterator_SetValue(), which was declared but not actually implemented
remove upb_MapIterator_Done(), which was implemented but not actually used

PiperOrigin-RevId: 489989481

Author: ericsalo

lua/msg.c

@@ -596,9 +596,8 @@ static int lupb_MapIterator_Next(lua_State* L) {
   size_t* iter = lua_touserdata(L, lua_upvalueindex(1));
   lupb_map* lmap = lupb_map_check(L, map);
 
-  if (upb_MapIterator_Next(lmap->map, iter)) {
-    upb_MessageValue key = upb_MapIterator_Key(lmap->map, *iter);
-    upb_MessageValue val = upb_MapIterator_Value(lmap->map, *iter);
+  upb_MessageValue key, val;
+  if (upb_Map_Next(lmap->map, &key, &val, iter)) {
     lupb_pushmsgval(L, map, lmap->key_type, key);
     lupb_pushmsgval(L, map, lmap->value_type, val);
     return 2;

python/convert.c

@@ -334,9 +334,8 @@ bool PyUpb_Map_IsEqual(const upb_Map* map1, const upb_Map* map2,
   const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
   size_t iter = kUpb_Map_Begin;
 
-  while (upb_MapIterator_Next(map1, &iter)) {
-    upb_MessageValue key = upb_MapIterator_Key(map1, iter);
-    upb_MessageValue val1 = upb_MapIterator_Value(map1, iter);
+  upb_MessageValue key, val1;
+  while (upb_Map_Next(map1, &key, &val1, &iter)) {
     upb_MessageValue val2;
     if (!upb_Map_Get(map2, key, &val2)) return false;
     if (!PyUpb_ValueEq(val1, val2, val_f)) return false;

python/map.c

@@ -318,11 +318,10 @@ static PyObject* PyUpb_MapContainer_Repr(PyObject* _self) {
     const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
     const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
     size_t iter = kUpb_Map_Begin;
-    while (upb_MapIterator_Next(map, &iter)) {
-      PyObject* key =
-          PyUpb_UpbToPy(upb_MapIterator_Key(map, iter), key_f, self->arena);
-      PyObject* val =
-          PyUpb_UpbToPy(upb_MapIterator_Value(map, iter), val_f, self->arena);
+    upb_MessageValue map_key, map_val;
+    while (upb_Map_Next(map, &map_key, &map_val, &iter)) {
+      PyObject* key = PyUpb_UpbToPy(map_key, key_f, self->arena);
+      PyObject* val = PyUpb_UpbToPy(map_val, val_f, self->arena);
       if (!key || !val) {
         Py_XDECREF(key);
         Py_XDECREF(val);
@@ -474,8 +473,8 @@ PyObject* PyUpb_MapIterator_IterNext(PyObject* _self) {
   }
   upb_Map* map = PyUpb_MapContainer_GetIfReified(self->map);
   if (!map) return NULL;
-  if (!upb_MapIterator_Next(map, &self->iter)) return NULL;
-  upb_MessageValue key = upb_MapIterator_Key(map, self->iter);
+  upb_MessageValue key, val;
+  if (!upb_Map_Next(map, &key, &val, &self->iter)) return NULL;
   const upb_FieldDef* f = PyUpb_MapContainer_GetField(self->map);
   const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
   const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);

upb/collections/map.c

@@ -76,16 +76,20 @@ bool upb_Map_Delete(upb_Map* map, upb_MessageValue key) {
   return _upb_Map_Delete(map, &key, map->key_size);
 }
 
-bool upb_MapIterator_Next(const upb_Map* map, size_t* iter) {
-  return _upb_map_next(map, iter);
+bool upb_Map_Next(const upb_Map* map, upb_MessageValue* key,
+                  upb_MessageValue* val, size_t* iter) {
+  upb_StringView k;
+  upb_value v;
+  const bool ok = upb_strtable_next2(&map->table, &k, &v, iter);
+  if (ok) {
+    _upb_map_fromkey(k, key, map->key_size);
+    _upb_map_fromvalue(v, val, map->val_size);
+  }
+  return ok;
 }
 
-bool upb_MapIterator_Done(const upb_Map* map, size_t iter) {
-  upb_strtable_iter i;
-  UPB_ASSERT(iter != kUpb_Map_Begin);
-  i.t = &map->table;
-  i.index = iter;
-  return upb_strtable_done(&i);
+bool upb_MapIterator_Next(const upb_Map* map, size_t* iter) {
+  return _upb_map_next(map, iter);
 }
 
 // Returns the key and value for this entry of the map.

upb/collections/map.h

@@ -78,37 +78,39 @@ UPB_INLINE bool upb_Map_Set(upb_Map* map, upb_MessageValue key,
 // Deletes this key from the table. Returns true if the key was present.
 bool upb_Map_Delete(upb_Map* map, upb_MessageValue key);
 
+// Map iteration:
+//
+// size_t iter = kUpb_Map_Begin;
+// upb_MessageValue key, val;
+// while (upb_Map_Next(map, &key, &val, &iter)) {
+//   ...
+// }
+
+#define kUpb_Map_Begin ((size_t)-1)
+
+// Advances to the next entry. Returns false if no more entries are present.
+// Otherwise returns true and populates both *key and *value.
+bool upb_Map_Next(const upb_Map* map, upb_MessageValue* key,
+                  upb_MessageValue* val, size_t* iter);
+
+// DEPRECATED iterator, slated for removal.
+
 /* Map iteration:
  *
  * size_t iter = kUpb_Map_Begin;
  * while (upb_MapIterator_Next(map, &iter)) {
  *   upb_MessageValue key = upb_MapIterator_Key(map, iter);
  *   upb_MessageValue val = upb_MapIterator_Value(map, iter);
- *
- *   // If mutating is desired.
- *   upb_MapIterator_SetValue(map, iter, value2);
  * }
  */
 
-#define kUpb_Map_Begin ((size_t)-1)
-
 // Advances to the next entry. Returns false if no more entries are present.
 bool upb_MapIterator_Next(const upb_Map* map, size_t* iter);
 
-/* Returns true if the iterator still points to a valid entry, or false if the
- * iterator is past the last element. It is an error to call this function with
- * kUpb_Map_Begin (you must call next() at least once first). */
-bool upb_MapIterator_Done(const upb_Map* map, size_t iter);
-
 /* Returns the key and value for this entry of the map. */
 upb_MessageValue upb_MapIterator_Key(const upb_Map* map, size_t iter);
 upb_MessageValue upb_MapIterator_Value(const upb_Map* map, size_t iter);
 
-/* Sets the value for this entry.  The iterator must not be done, and the
- * iterator must not have been initialized const. */
-void upb_MapIterator_SetValue(upb_Map* map, size_t iter,
-                              upb_MessageValue value);
-
 #ifdef __cplusplus
 } /* extern "C" */
 #endif

upb/json/encode.c

@@ -470,20 +470,20 @@ static void jsonenc_fieldmask(jsonenc* e, const upb_Message* msg,
 
 static void jsonenc_struct(jsonenc* e, const upb_Message* msg,
                            const upb_MessageDef* m) {
+  jsonenc_putstr(e, "{");
+
   const upb_FieldDef* fields_f = upb_MessageDef_FindFieldByNumber(m, 1);
   const upb_Map* fields = upb_Message_Get(msg, fields_f).map_val;
-  const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(fields_f);
-  const upb_FieldDef* value_f = upb_MessageDef_FindFieldByNumber(entry_m, 2);
-  size_t iter = kUpb_Map_Begin;
-  bool first = true;
-
-  jsonenc_putstr(e, "{");
 
   if (fields) {
-    while (upb_MapIterator_Next(fields, &iter)) {
-      upb_MessageValue key = upb_MapIterator_Key(fields, iter);
-      upb_MessageValue val = upb_MapIterator_Value(fields, iter);
+    const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(fields_f);
+    const upb_FieldDef* value_f = upb_MessageDef_FindFieldByNumber(entry_m, 2);
+
+    size_t iter = kUpb_Map_Begin;
+    bool first = true;
 
+    upb_MessageValue key, val;
+    while (upb_Map_Next(fields, &key, &val, &iter)) {
       jsonenc_putsep(e, ",", &first);
       jsonenc_string(e, key.str_val);
       jsonenc_putstr(e, ":");
@@ -677,19 +677,21 @@ static void jsonenc_array(jsonenc* e, const upb_Array* arr,
 }
 
 static void jsonenc_map(jsonenc* e, const upb_Map* map, const upb_FieldDef* f) {
+  jsonenc_putstr(e, "{");
+
   const upb_MessageDef* entry = upb_FieldDef_MessageSubDef(f);
   const upb_FieldDef* key_f = upb_MessageDef_FindFieldByNumber(entry, 1);
   const upb_FieldDef* val_f = upb_MessageDef_FindFieldByNumber(entry, 2);
-  size_t iter = kUpb_Map_Begin;
-  bool first = true;
-
-  jsonenc_putstr(e, "{");
 
   if (map) {
-    while (upb_MapIterator_Next(map, &iter)) {
+    size_t iter = kUpb_Map_Begin;
+    bool first = true;
+
+    upb_MessageValue key, val;
+    while (upb_Map_Next(map, &key, &val, &iter)) {
       jsonenc_putsep(e, ",", &first);
-      jsonenc_mapkey(e, upb_MapIterator_Key(map, iter), key_f);
-      jsonenc_scalar(e, upb_MapIterator_Value(map, iter), val_f);
+      jsonenc_mapkey(e, key, key_f);
+      jsonenc_scalar(e, val, val_f);
     }
   }
 

upb/reflection/message.c

@@ -299,8 +299,8 @@ bool _upb_Message_DiscardUnknown(upb_Message* msg, const upb_MessageDef* m,
 
       if (!val_m) continue;
 
-      while (upb_MapIterator_Next(map, &iter)) {
-        upb_MessageValue map_val = upb_MapIterator_Value(map, iter);
+      upb_MessageValue map_key, map_val;
+      while (upb_Map_Next(map, &map_key, &map_val, &iter)) {
         if (!_upb_Message_DiscardUnknown((upb_Message*)map_val.msg_val, val_m,
                                          depth)) {
           ret = false;

upb/text/encode.c

@@ -282,9 +282,8 @@ static void txtenc_mapentry(txtenc* e, upb_MessageValue key,
 static void txtenc_map(txtenc* e, const upb_Map* map, const upb_FieldDef* f) {
   if (e->options & UPB_TXTENC_NOSORT) {
     size_t iter = kUpb_Map_Begin;
-    while (upb_MapIterator_Next(map, &iter)) {
-      upb_MessageValue key = upb_MapIterator_Key(map, iter);
-      upb_MessageValue val = upb_MapIterator_Value(map, iter);
+    upb_MessageValue key, val;
+    while (upb_Map_Next(map, &key, &val, &iter)) {
       txtenc_mapentry(e, key, val, f);
     }
   } else {

upb/util/required_fields.c

@@ -265,9 +265,8 @@ static void upb_util_FindUnsetRequiredInternal(upb_FindContext* ctx,
       if (!val_m) continue;
       const upb_Map* map = val.map_val;
       size_t iter = kUpb_Map_Begin;
-      while (upb_MapIterator_Next(map, &iter)) {
-        upb_MessageValue key = upb_MapIterator_Key(map, iter);
-        upb_MessageValue map_val = upb_MapIterator_Value(map, iter);
+      upb_MessageValue key, map_val;
+      while (upb_Map_Next(map, &key, &map_val, &iter)) {
         upb_FindContext_Push(ctx, (upb_FieldPathEntry){.map_key = key});
         upb_util_FindUnsetRequiredInternal(ctx, map_val.msg_val, val_m);
         upb_FindContext_Pop(ctx);