Add upb_Map_GetMutable API to upb

PiperOrigin-RevId: 742338581

Author: dbenson24

rust/upb/lib.rs

@@ -29,8 +29,8 @@ pub use extension_registry::{upb_ExtensionRegistry, RawExtensionRegistry};
 
 mod map;
 pub use map::{
-    upb_Map, upb_Map_Clear, upb_Map_Delete, upb_Map_Get, upb_Map_Insert, upb_Map_New, upb_Map_Next,
-    upb_Map_Size, MapInsertStatus, RawMap, UPB_MAP_BEGIN,
+    upb_Map, upb_Map_Clear, upb_Map_Delete, upb_Map_Get, upb_Map_GetMutable, upb_Map_Insert,
+    upb_Map_New, upb_Map_Next, upb_Map_Size, MapInsertStatus, RawMap, UPB_MAP_BEGIN,
 };
 
 mod message;

rust/upb/map.rs

@@ -6,7 +6,7 @@
 // https://developers.google.com/open-source/licenses/bsd
 
 use super::opaque_pointee::opaque_pointee;
-use super::{upb_MessageValue, CType, RawArena};
+use super::{upb_Message, upb_MessageValue, CType, RawArena};
 use core::ptr::NonNull;
 
 opaque_pointee!(upb_Map);
@@ -33,6 +33,7 @@ extern "C" {
         arena: RawArena,
     ) -> MapInsertStatus;
     pub fn upb_Map_Get(map: RawMap, key: upb_MessageValue, value: *mut upb_MessageValue) -> bool;
+    pub fn upb_Map_GetMutable(map: RawMap, key: upb_MessageValue) -> *mut upb_Message;
     pub fn upb_Map_Delete(
         map: RawMap,
         key: upb_MessageValue,
@@ -60,6 +61,7 @@ mod tests {
         assert_linked!(upb_Map_Size);
         assert_linked!(upb_Map_Insert);
         assert_linked!(upb_Map_Get);
+        assert_linked!(upb_Map_GetMutable);
         assert_linked!(upb_Map_Delete);
         assert_linked!(upb_Map_Clear);
         assert_linked!(upb_Map_Next);

upb/message/map.c

@@ -17,6 +17,7 @@
 #include "upb/hash/str_table.h"
 #include "upb/mem/arena.h"
 #include "upb/message/internal/map.h"
+#include "upb/message/internal/types.h"
 #include "upb/message/map.h"
 #include "upb/message/message.h"
 #include "upb/message/value.h"
@@ -52,6 +53,16 @@ bool upb_Map_Get(const upb_Map* map, upb_MessageValue key,
   return _upb_Map_Get(map, &key, map->key_size, val, map->val_size);
 }
 
+struct upb_Message* upb_Map_GetMutable(upb_Map* map, upb_MessageValue key) {
+  UPB_ASSERT(map->val_size == sizeof(upb_Message*));
+  upb_Message* val = NULL;
+  if (_upb_Map_Get(map, &key, map->key_size, &val, sizeof(upb_Message*))) {
+    return val;
+  } else {
+    return NULL;
+  }
+}
+
 void upb_Map_Clear(upb_Map* map) { _upb_Map_Clear(map); }
 
 upb_MapInsertStatus upb_Map_Insert(upb_Map* map, upb_MessageValue key,

upb/message/map.h

@@ -13,6 +13,7 @@
 #include "upb/base/descriptor_constants.h"
 #include "upb/mem/arena.h"
 #include "upb/message/internal/map.h"
+#include "upb/message/internal/types.h"
 #include "upb/message/value.h"
 #include "upb/mini_table/field.h"
 #include "upb/mini_table/message.h"
@@ -39,6 +40,12 @@ UPB_API size_t upb_Map_Size(const upb_Map* map);
 UPB_API bool upb_Map_Get(const upb_Map* map, upb_MessageValue key,
                          upb_MessageValue* val);
 
+// Returns a mutable pointer to the value for the given key. Returns NULL if the
+// key is not present.
+// This function is only legal to call for maps that contain messages.
+UPB_API struct upb_Message* upb_Map_GetMutable(upb_Map* map,
+                                               upb_MessageValue key);
+
 // Removes all entries in the map.
 UPB_API void upb_Map_Clear(upb_Map* map);
 
@@ -71,7 +78,7 @@ UPB_API bool upb_Map_Delete(upb_Map* map, upb_MessageValue key,
 //   ...
 // }
 
-#define kUpb_Map_Begin ((size_t) - 1)
+#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.