diff --git a/upb/hash/common.c b/upb/hash/common.c index 17b68235e116d..c20c61badd299 100644 --- a/upb/hash/common.c +++ b/upb/hash/common.c @@ -52,7 +52,7 @@ UPB_INLINE int _upb_popcnt32(uint32_t i) { #undef UPB_FAST_POPCOUNT32 -UPB_INLINE uint8_t _upb_log2_table_size(upb_table* t) { +UPB_INLINE uint8_t _upb_log2_table_size(const upb_table* t) { return _upb_popcnt32(t->mask); } @@ -589,6 +589,42 @@ bool upb_strtable_resize(upb_strtable* t, size_t size_lg2, upb_Arena* a) { return true; } +bool upb_strtable_copy(upb_strtable* dest, const upb_strtable* src, + upb_Arena* a) { + if (src->t.count == 0) { + return upb_strtable_init(dest, 0, a); + } + dest->t.count = src->t.count; + dest->t.mask = src->t.mask; + dest->t.entries = + upb_Arena_Malloc(a, upb_table_size(&src->t) * sizeof(upb_tabent)); + if (!dest->t.entries) return false; + upb_tabent* restrict dest_entries = dest->t.entries; + const upb_tabent* restrict src_entries = src->t.entries; + size_t table_size = upb_table_size(&src->t); + for (size_t i = 0; i < table_size; i++) { + upb_tabent* dest_ent = &dest_entries[i]; + const upb_tabent* src_ent = &src_entries[i]; + if (!upb_tabent_isempty(src_ent)) { + upb_StringView sv = upb_key_strview(src_ent->key); + upb_SizePrefixString* size_prefix_string = + upb_SizePrefixString_Copy(sv, a); + if (!size_prefix_string) return false; + dest_ent->key.str = size_prefix_string; + dest_ent->val = src_ent->val; + if (UPB_UNPREDICTABLE(upb_tabent_hasnext(src_ent))) { + size_t offset = upb_tabent_next(src_ent) - src_entries; + upb_tabent_setnext(dest_ent, dest_entries + offset); + } else { + upb_tabent_clearnext(dest_ent); + } + } else { + *dest_ent = (upb_tabent){}; + } + } + return true; +} + bool upb_strtable_insert(upb_strtable* t, const char* k, size_t len, upb_value v, upb_Arena* a) { if (isfull(&t->t)) { @@ -812,6 +848,40 @@ bool upb_inttable_init(upb_inttable* t, upb_Arena* a) { return upb_inttable_sizedinit(t, 3, a); } +bool upb_inttable_copy(upb_inttable* dest, const upb_inttable* src, + upb_Arena* a) { + if (src->t.count == 0) { + return upb_inttable_sizedinit(dest, 0, a); + } + + if (!upb_inttable_sizedinit( + dest, src->t.mask ? _upb_log2_table_size(&src->t) : 0, a)) { + return false; + } + dest->t.count = src->t.count; + + upb_tabent* restrict dest_entries = dest->t.entries; + const upb_tabent* restrict src_entries = src->t.entries; + size_t table_size = upb_table_size(&src->t); + for (size_t i = 0; i < table_size; i++) { + upb_tabent* dest_ent = &dest_entries[i]; + const upb_tabent* src_ent = &src_entries[i]; + if (!upb_tabent_isempty(src_ent)) { + dest_ent->key = src_ent->key; + dest_ent->val = src_ent->val; + if (UPB_UNPREDICTABLE(upb_tabent_hasnext(src_ent))) { + size_t offset = upb_tabent_next(src_ent) - src_entries; + upb_tabent_setnext(dest_ent, dest_entries + offset); + } else { + upb_tabent_clearnext(dest_ent); + } + } else { + *dest_ent = (upb_tabent){}; + } + } + return true; +} + bool upb_inttable_insert(upb_inttable* t, uintptr_t key, upb_value val, upb_Arena* a) { if (isfull(&t->t)) { diff --git a/upb/hash/int_table.h b/upb/hash/int_table.h index 521f048f1866f..eaaa78c4fe365 100644 --- a/upb/hash/int_table.h +++ b/upb/hash/int_table.h @@ -40,6 +40,10 @@ size_t upb_inttable_count(const upb_inttable* t); UPB_NODISCARD bool upb_inttable_insert(upb_inttable* t, uintptr_t key, upb_value val, upb_Arena* a); +// Copies the table without rehashing. +bool upb_inttable_copy(upb_inttable* dest, const upb_inttable* src, + upb_Arena* a); + // Looks up key in this table, returning "true" if the key was found. // If v is non-NULL, copies the value for this key into *v. bool upb_inttable_lookup(const upb_inttable* t, uintptr_t key, upb_value* v); diff --git a/upb/hash/str_table.h b/upb/hash/str_table.h index e3d428482c6e5..85cb4815a8ed8 100644 --- a/upb/hash/str_table.h +++ b/upb/hash/str_table.h @@ -48,6 +48,10 @@ void upb_strtable_clear(upb_strtable* t); UPB_NODISCARD bool upb_strtable_insert(upb_strtable* t, const char* key, size_t len, upb_value val, upb_Arena* a); +// Copies the table and its keys without rehashing. +bool upb_strtable_copy(upb_strtable* dest, const upb_strtable* src, + upb_Arena* a); + // Looks up key in this table, returning "true" if the key was found. // If v is non-NULL, copies the value for this key into *v. bool upb_strtable_lookup2(const upb_strtable* t, const char* key, size_t len, diff --git a/upb/message/BUILD b/upb/message/BUILD index e36b5c237cb46..535017765cf4e 100644 --- a/upb/message/BUILD +++ b/upb/message/BUILD @@ -166,6 +166,7 @@ cc_library( ":message", "//upb/base", "//upb/base:internal", + "//upb/hash", "//upb/mem", "//upb/mini_table", "//upb/mini_table:internal", diff --git a/upb/message/copy.c b/upb/message/copy.c index 8c7f2397e9f95..f18666a2dc43a 100644 --- a/upb/message/copy.c +++ b/upb/message/copy.c @@ -13,6 +13,9 @@ #include "upb/base/descriptor_constants.h" #include "upb/base/string_view.h" +#include "upb/hash/common.h" +#include "upb/hash/int_table.h" +#include "upb/hash/str_table.h" #include "upb/mem/arena.h" #include "upb/message/accessors.h" #include "upb/message/array.h" @@ -86,26 +89,72 @@ upb_Map* upb_Map_DeepClone(const upb_Map* map, upb_CType key_type, upb_CType value_type, const upb_MiniTable* map_entry_table, upb_Arena* arena) { - upb_Map* cloned_map = _upb_Map_New(arena, map->key_size, map->val_size); + upb_Map* cloned_map = upb_Arena_Malloc(arena, sizeof(upb_Map)); if (cloned_map == NULL) { return NULL; } - upb_MessageValue key, val; - size_t iter = kUpb_Map_Begin; - while (upb_Map_Next(map, &key, &val, &iter)) { - const upb_MiniTableField* value_field = - upb_MiniTable_MapValue(map_entry_table); - const upb_MiniTable* value_sub = - upb_MiniTableField_CType(value_field) == kUpb_CType_Message - ? upb_MiniTable_GetSubMessageTable(value_field) - : NULL; - upb_CType value_field_type = upb_MiniTableField_CType(value_field); - if (!upb_Clone_MessageValue(&val, value_field_type, value_sub, arena)) { + cloned_map->key_size = map->key_size; + cloned_map->val_size = map->val_size; + cloned_map->UPB_PRIVATE(is_frozen) = false; + cloned_map->UPB_PRIVATE(is_strtable) = map->UPB_PRIVATE(is_strtable); + + const upb_MiniTableField* value_field = + upb_MiniTable_MapValue(map_entry_table); + const upb_MiniTable* value_sub = + upb_MiniTableField_CType(value_field) == kUpb_CType_Message + ? upb_MiniTable_GetSubMessageTable(value_field) + : NULL; + upb_CType value_field_type = upb_MiniTableField_CType(value_field); + + bool is_primitive = value_field_type != kUpb_CType_Message && + value_field_type != kUpb_CType_String && + value_field_type != kUpb_CType_Bytes; + + if (map->UPB_PRIVATE(is_strtable)) { + if (!upb_strtable_copy(&cloned_map->t.strtable, &map->t.strtable, arena)) { return NULL; } - if (!upb_Map_Set(cloned_map, key, val, arena)) { + if (!is_primitive) { + intptr_t iter = UPB_STRTABLE_BEGIN; + upb_StringView key; + upb_value tabval; + while ( + upb_strtable_next2(&cloned_map->t.strtable, &key, &tabval, &iter)) { + upb_MessageValue val; + _upb_map_fromvalue(tabval, &val, map->val_size); + if (!upb_Clone_MessageValue(&val, value_field_type, value_sub, arena)) { + return NULL; + } + upb_value cloned_tabval = {0}; + if (!_upb_map_tovalue(&val, map->val_size, &cloned_tabval, arena)) { + return NULL; + } + upb_strtable_setentryvalue(&cloned_map->t.strtable, iter, + cloned_tabval); + } + } + } else { + if (!upb_inttable_copy(&cloned_map->t.inttable, &map->t.inttable, arena)) { return NULL; } + if (!is_primitive) { + intptr_t iter = UPB_INTTABLE_BEGIN; + uintptr_t key; + upb_value tabval; + while (upb_inttable_next(&cloned_map->t.inttable, &key, &tabval, &iter)) { + upb_MessageValue val; + _upb_map_fromvalue(tabval, &val, map->val_size); + if (!upb_Clone_MessageValue(&val, value_field_type, value_sub, arena)) { + return NULL; + } + upb_value cloned_tabval = {0}; + if (!_upb_map_tovalue(&val, map->val_size, &cloned_tabval, arena)) { + return NULL; + } + upb_inttable_setentryvalue(&cloned_map->t.inttable, iter, + cloned_tabval); + } + } } return cloned_map; }