optimize map copy PiperOrigin-RevId: 967902019
diff --git a/upb/hash/common.c b/upb/hash/common.c index 844efa3..1ea2272 100644 --- a/upb/hash/common.c +++ b/upb/hash/common.c
@@ -53,7 +53,7 @@ #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); } @@ -590,6 +590,42 @@ 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)) { @@ -813,6 +849,40 @@ 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 521f048..eaaa78c 100644 --- a/upb/hash/int_table.h +++ b/upb/hash/int_table.h
@@ -40,6 +40,10 @@ 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 e3d4284..85cb481 100644 --- a/upb/hash/str_table.h +++ b/upb/hash/str_table.h
@@ -48,6 +48,10 @@ 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 c3a7004..6a4d350 100644 --- a/upb/message/BUILD +++ b/upb/message/BUILD
@@ -227,6 +227,7 @@ ":types", "//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 da7e7f8..f628407 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" @@ -83,30 +86,81 @@ UPB_UNREACHABLE(); } +static bool upb_Clone_MapValue(upb_value* tabval, size_t map_val_size, + upb_CType value_field_type, + const upb_MiniTable* value_sub, upb_Arena* arena, + upb_value* cloned_tabval) { + upb_MessageValue val; + _upb_map_fromvalue(*tabval, &val, map_val_size); + if (!upb_Clone_MessageValue(&val, value_field_type, value_sub, arena)) { + return false; + } + *cloned_tabval = (upb_value){0}; // Initialize + if (!_upb_map_tovalue(&val, map_val_size, cloned_tabval, arena)) { + return false; + } + return true; +} + 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_MiniTable_SubMessage(value_field); + 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_value cloned_tabval; + if (!upb_Clone_MapValue(&tabval, map->val_size, value_field_type, + value_sub, arena, &cloned_tabval)) { + 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_value cloned_tabval; + if (!upb_Clone_MapValue(&tabval, map->val_size, value_field_type, + value_sub, arena, &cloned_tabval)) { + return NULL; + } + upb_inttable_setentryvalue(&cloned_map->t.inttable, iter, + cloned_tabval); + } + } } return cloned_map; }