diff --git a/rust/upb/lib.rs b/rust/upb/lib.rs index 49cea15e5b..e116727705 100644 --- a/rust/upb/lib.rs +++ b/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; diff --git a/rust/upb/map.rs b/rust/upb/map.rs index afddb42c17..d03c61ee3a 100644 --- a/rust/upb/map.rs +++ b/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); diff --git a/upb/message/map.c b/upb/message/map.c index dcc1caff91..56a5fcc62d 100644 --- a/upb/message/map.c +++ b/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, diff --git a/upb/message/map.h b/upb/message/map.h index c758a695ce..77710d93dc 100644 --- a/upb/message/map.h +++ b/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.