diff --git a/BUILD b/BUILD index 045ded0abc..89d9e3c22d 100644 --- a/BUILD +++ b/BUILD @@ -285,6 +285,30 @@ cc_library( ], ) +cc_library( + name = "message_promote", + srcs = [ + "upb/message/promote.c", + ], + hdrs = [ + "upb/message/promote.h", + ], + copts = UPB_DEFAULT_COPTS, + visibility = ["//visibility:public"], + deps = [ + ":collections_internal", + ":eps_copy_input_stream", + ":hash", + ":message_accessors", + ":message_internal", + ":mini_table_internal", + ":port", + ":upb", + ":wire", + ":wire_reader", + ], +) + cc_library( name = "message_copy", srcs = [ @@ -341,6 +365,25 @@ cc_test( ], ) +cc_test( + name = "message_promote_test", + srcs = ["upb/message/promote_test.cc"], + deps = [ + ":collections", + ":message_accessors", + ":message_promote", + ":mini_table_internal", + ":port", + ":upb", + "//upb/test:test_messages_proto2_upb_proto", + "//upb/test:test_messages_proto3_upb_proto", + "//upb/test:test_upb_proto", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_googletest//:gtest_main", + "@com_google_protobuf//:protobuf", + ], +) + cc_test( name = "message_copy_test", srcs = ["upb/message/copy_test.cc"], diff --git a/upb/message/accessors.c b/upb/message/accessors.c index 2f1dd7cb6e..d2c74df454 100644 --- a/upb/message/accessors.c +++ b/upb/message/accessors.c @@ -41,246 +41,6 @@ // Must be last. #include "upb/port/def.inc" -// Parses unknown data by merging into existing base_message or creating a -// new message usingg mini_table. -static upb_UnknownToMessageRet upb_MiniTable_ParseUnknownMessage( - const char* unknown_data, size_t unknown_size, - const upb_MiniTable* mini_table, upb_Message* base_message, - int decode_options, upb_Arena* arena) { - upb_UnknownToMessageRet ret; - ret.message = - base_message ? base_message : _upb_Message_New(mini_table, arena); - if (!ret.message) { - ret.status = kUpb_UnknownToMessage_OutOfMemory; - return ret; - } - // Decode sub message using unknown field contents. - const char* data = unknown_data; - uint32_t tag; - uint64_t message_len = 0; - data = upb_WireReader_ReadTag(data, &tag); - data = upb_WireReader_ReadVarint(data, &message_len); - upb_DecodeStatus status = upb_Decode(data, message_len, ret.message, - mini_table, NULL, decode_options, arena); - if (status == kUpb_DecodeStatus_OutOfMemory) { - ret.status = kUpb_UnknownToMessage_OutOfMemory; - } else if (status == kUpb_DecodeStatus_Ok) { - ret.status = kUpb_UnknownToMessage_Ok; - } else { - ret.status = kUpb_UnknownToMessage_ParseError; - } - return ret; -} - -upb_GetExtension_Status upb_MiniTable_GetOrPromoteExtension( - upb_Message* msg, const upb_MiniTableExtension* ext_table, - int decode_options, upb_Arena* arena, - const upb_Message_Extension** extension) { - UPB_ASSERT(upb_MiniTableField_CType(&ext_table->field) == kUpb_CType_Message); - *extension = _upb_Message_Getext(msg, ext_table); - if (*extension) { - return kUpb_GetExtension_Ok; - } - - // Check unknown fields, if available promote. - int field_number = ext_table->field.number; - upb_FindUnknownRet result = upb_MiniTable_FindUnknown( - msg, field_number, kUpb_WireFormat_DefaultDepthLimit); - if (result.status != kUpb_FindUnknown_Ok) { - return kUpb_GetExtension_NotPresent; - } - size_t len; - size_t ofs = result.ptr - upb_Message_GetUnknown(msg, &len); - // Decode and promote from unknown. - const upb_MiniTable* extension_table = ext_table->sub.submsg; - upb_UnknownToMessageRet parse_result = upb_MiniTable_ParseUnknownMessage( - result.ptr, result.len, extension_table, - /* base_message= */ NULL, decode_options, arena); - switch (parse_result.status) { - case kUpb_UnknownToMessage_OutOfMemory: - return kUpb_GetExtension_OutOfMemory; - case kUpb_UnknownToMessage_ParseError: - return kUpb_GetExtension_ParseError; - case kUpb_UnknownToMessage_NotFound: - return kUpb_GetExtension_NotPresent; - case kUpb_UnknownToMessage_Ok: - break; - } - upb_Message* extension_msg = parse_result.message; - // Add to extensions. - upb_Message_Extension* ext = - _upb_Message_GetOrCreateExtension(msg, ext_table, arena); - if (!ext) { - return kUpb_GetExtension_OutOfMemory; - } - memcpy(&ext->data, &extension_msg, sizeof(extension_msg)); - *extension = ext; - const char* delete_ptr = upb_Message_GetUnknown(msg, &len) + ofs; - upb_Message_DeleteUnknown(msg, delete_ptr, result.len); - return kUpb_GetExtension_Ok; -} - -upb_GetExtensionAsBytes_Status upb_MiniTable_GetExtensionAsBytes( - const upb_Message* msg, const upb_MiniTableExtension* ext_table, - int encode_options, upb_Arena* arena, const char** extension_data, - size_t* len) { - const upb_Message_Extension* msg_ext = _upb_Message_Getext(msg, ext_table); - UPB_ASSERT(upb_MiniTableField_CType(&ext_table->field) == kUpb_CType_Message); - if (msg_ext) { - upb_EncodeStatus status = - upb_Encode(msg_ext->data.ptr, msg_ext->ext->sub.submsg, encode_options, - arena, (char**)extension_data, len); - if (status != kUpb_EncodeStatus_Ok) { - return kUpb_GetExtensionAsBytes_EncodeError; - } - return kUpb_GetExtensionAsBytes_Ok; - } - int field_number = ext_table->field.number; - upb_FindUnknownRet result = upb_MiniTable_FindUnknown( - msg, field_number, upb_DecodeOptions_GetMaxDepth(encode_options)); - if (result.status != kUpb_FindUnknown_Ok) { - return kUpb_GetExtensionAsBytes_NotPresent; - } - const char* data = result.ptr; - uint32_t tag; - uint64_t message_len = 0; - data = upb_WireReader_ReadTag(data, &tag); - data = upb_WireReader_ReadVarint(data, &message_len); - *extension_data = data; - *len = message_len; - return kUpb_GetExtensionAsBytes_Ok; -} - -static upb_FindUnknownRet upb_FindUnknownRet_ParseError(void) { - return (upb_FindUnknownRet){.status = kUpb_FindUnknown_ParseError}; -} - -upb_FindUnknownRet upb_MiniTable_FindUnknown(const upb_Message* msg, - uint32_t field_number, - int depth_limit) { - size_t size; - upb_FindUnknownRet ret; - - const char* ptr = upb_Message_GetUnknown(msg, &size); - upb_EpsCopyInputStream stream; - upb_EpsCopyInputStream_Init(&stream, &ptr, size, true); - - while (!upb_EpsCopyInputStream_IsDone(&stream, &ptr)) { - uint32_t tag; - const char* unknown_begin = ptr; - ptr = upb_WireReader_ReadTag(ptr, &tag); - if (!ptr) return upb_FindUnknownRet_ParseError(); - if (field_number == upb_WireReader_GetFieldNumber(tag)) { - ret.status = kUpb_FindUnknown_Ok; - ret.ptr = upb_EpsCopyInputStream_GetAliasedPtr(&stream, unknown_begin); - ptr = _upb_WireReader_SkipValue(ptr, tag, depth_limit, &stream); - // Because we know that the input is a flat buffer, it is safe to perform - // pointer arithmetic on aliased pointers. - ret.len = upb_EpsCopyInputStream_GetAliasedPtr(&stream, ptr) - ret.ptr; - return ret; - } - - ptr = _upb_WireReader_SkipValue(ptr, tag, depth_limit, &stream); - if (!ptr) return upb_FindUnknownRet_ParseError(); - } - ret.status = kUpb_FindUnknown_NotPresent; - ret.ptr = NULL; - ret.len = 0; - return ret; -} - -// Warning: See TODO(b/267655898) -upb_UnknownToMessageRet upb_MiniTable_PromoteUnknownToMessage( - upb_Message* msg, const upb_MiniTable* mini_table, - const upb_MiniTableField* field, const upb_MiniTable* sub_mini_table, - int decode_options, upb_Arena* arena) { - upb_FindUnknownRet unknown; - // We need to loop and merge unknowns that have matching tag field->number. - upb_Message* message = NULL; - // Callers should check that message is not set first before calling - // PromotoUnknownToMessage. - UPB_ASSERT(upb_MiniTable_GetSubMessageTable(mini_table, field) == - sub_mini_table); - bool is_oneof = _upb_MiniTableField_InOneOf(field); - if (!is_oneof || _upb_getoneofcase_field(msg, field) == field->number) { - UPB_ASSERT(upb_Message_GetMessage(msg, field, NULL) == NULL); - } - upb_UnknownToMessageRet ret; - ret.status = kUpb_UnknownToMessage_Ok; - do { - unknown = upb_MiniTable_FindUnknown( - msg, field->number, upb_DecodeOptions_GetMaxDepth(decode_options)); - switch (unknown.status) { - case kUpb_FindUnknown_Ok: { - const char* unknown_data = unknown.ptr; - size_t unknown_size = unknown.len; - ret = upb_MiniTable_ParseUnknownMessage(unknown_data, unknown_size, - sub_mini_table, message, - decode_options, arena); - if (ret.status == kUpb_UnknownToMessage_Ok) { - message = ret.message; - upb_Message_DeleteUnknown(msg, unknown_data, unknown_size); - } - } break; - case kUpb_FindUnknown_ParseError: - ret.status = kUpb_UnknownToMessage_ParseError; - break; - case kUpb_FindUnknown_NotPresent: - // If we parsed at least one unknown, we are done. - ret.status = - message ? kUpb_UnknownToMessage_Ok : kUpb_UnknownToMessage_NotFound; - break; - } - } while (unknown.status == kUpb_FindUnknown_Ok); - if (message) { - if (is_oneof) { - *_upb_oneofcase_field(msg, field) = field->number; - } - upb_Message_SetMessage(msg, mini_table, field, message); - ret.message = message; - } - return ret; -} - -// Moves repeated messages in unknowns to a upb_Array. -// -// Since the repeated field is not a scalar type we don't check for -// kUpb_LabelFlags_IsPacked. -// TODO(b/251007554): Optimize. Instead of converting messages one at a time, -// scan all unknown data once and compact. -upb_UnknownToMessage_Status upb_MiniTable_PromoteUnknownToMessageArray( - upb_Message* msg, const upb_MiniTableField* field, - const upb_MiniTable* mini_table, int decode_options, upb_Arena* arena) { - upb_Array* repeated_messages = upb_Message_GetMutableArray(msg, field); - // Find all unknowns with given field number and parse. - upb_FindUnknownRet unknown; - do { - unknown = upb_MiniTable_FindUnknown( - msg, field->number, upb_DecodeOptions_GetMaxDepth(decode_options)); - if (unknown.status == kUpb_FindUnknown_Ok) { - upb_UnknownToMessageRet ret = upb_MiniTable_ParseUnknownMessage( - unknown.ptr, unknown.len, mini_table, - /* base_message= */ NULL, decode_options, arena); - if (ret.status == kUpb_UnknownToMessage_Ok) { - upb_MessageValue value; - value.msg_val = ret.message; - // Allocate array on demand before append. - if (!repeated_messages) { - upb_Message_ResizeArray(msg, field, 0, arena); - repeated_messages = upb_Message_GetMutableArray(msg, field); - } - if (!upb_Array_Append(repeated_messages, value, arena)) { - return kUpb_UnknownToMessage_OutOfMemory; - } - upb_Message_DeleteUnknown(msg, unknown.ptr, unknown.len); - } else { - return ret.status; - } - } - } while (unknown.status == kUpb_FindUnknown_Ok); - return kUpb_UnknownToMessage_Ok; -} - upb_MapInsertStatus upb_Message_InsertMapEntry(upb_Map* map, const upb_MiniTable* mini_table, const upb_MiniTableField* field, @@ -306,39 +66,3 @@ upb_MapInsertStatus upb_Message_InsertMapEntry(upb_Map* map, &map_entry_value); return upb_Map_Insert(map, map_entry_key, map_entry_value, arena); } - -// Moves repeated messages in unknowns to a upb_Map. -upb_UnknownToMessage_Status upb_MiniTable_PromoteUnknownToMap( - upb_Message* msg, const upb_MiniTable* mini_table, - const upb_MiniTableField* field, int decode_options, upb_Arena* arena) { - const upb_MiniTable* map_entry_mini_table = - mini_table->subs[field->UPB_PRIVATE(submsg_index)].submsg; - UPB_ASSERT(map_entry_mini_table); - UPB_ASSERT(map_entry_mini_table); - UPB_ASSERT(map_entry_mini_table->field_count == 2); - UPB_ASSERT(upb_FieldMode_Get(field) == kUpb_FieldMode_Map); - // Find all unknowns with given field number and parse. - upb_FindUnknownRet unknown; - while (1) { - unknown = upb_MiniTable_FindUnknown( - msg, field->number, upb_DecodeOptions_GetMaxDepth(decode_options)); - if (unknown.status != kUpb_FindUnknown_Ok) break; - upb_UnknownToMessageRet ret = upb_MiniTable_ParseUnknownMessage( - unknown.ptr, unknown.len, map_entry_mini_table, - /* base_message= */ NULL, decode_options, arena); - if (ret.status != kUpb_UnknownToMessage_Ok) return ret.status; - // Allocate map on demand before append. - upb_Map* map = upb_Message_GetOrCreateMutableMap(msg, map_entry_mini_table, - field, arena); - upb_Message* map_entry_message = ret.message; - upb_MapInsertStatus insert_status = upb_Message_InsertMapEntry( - map, mini_table, field, map_entry_message, arena); - if (insert_status == kUpb_MapInsertStatus_OutOfMemory) { - return kUpb_UnknownToMessage_OutOfMemory; - } - UPB_ASSUME(insert_status == kUpb_MapInsertStatus_Inserted || - insert_status == kUpb_MapInsertStatus_Replaced); - upb_Message_DeleteUnknown(msg, unknown.ptr, unknown.len); - } - return kUpb_UnknownToMessage_Ok; -} diff --git a/upb/message/accessors.h b/upb/message/accessors.h index 6d42d88cfe..0432d9e08e 100644 --- a/upb/message/accessors.h +++ b/upb/message/accessors.h @@ -369,100 +369,6 @@ upb_MapInsertStatus upb_Message_InsertMapEntry(upb_Map* map, upb_Message* map_entry_message, upb_Arena* arena); -typedef enum { - kUpb_GetExtension_Ok, - kUpb_GetExtension_NotPresent, - kUpb_GetExtension_ParseError, - kUpb_GetExtension_OutOfMemory, -} upb_GetExtension_Status; - -typedef enum { - kUpb_GetExtensionAsBytes_Ok, - kUpb_GetExtensionAsBytes_NotPresent, - kUpb_GetExtensionAsBytes_EncodeError, -} upb_GetExtensionAsBytes_Status; - -// Returns a message extension or promotes an unknown field to -// an extension. -// -// TODO(ferhat): Only supports extension fields that are messages, -// expand support to include non-message types. -upb_GetExtension_Status upb_MiniTable_GetOrPromoteExtension( - upb_Message* msg, const upb_MiniTableExtension* ext_table, - int decode_options, upb_Arena* arena, - const upb_Message_Extension** extension); - -// Returns a message extension or unknown field matching the extension -// data as bytes. -// -// If an extension has already been decoded it will be re-encoded -// to bytes. -upb_GetExtensionAsBytes_Status upb_MiniTable_GetExtensionAsBytes( - const upb_Message* msg, const upb_MiniTableExtension* ext_table, - int encode_options, upb_Arena* arena, const char** extension_data, - size_t* len); - -typedef enum { - kUpb_FindUnknown_Ok, - kUpb_FindUnknown_NotPresent, - kUpb_FindUnknown_ParseError, -} upb_FindUnknown_Status; - -typedef struct { - upb_FindUnknown_Status status; - // Start of unknown field data in message arena. - const char* ptr; - // Size of unknown field data. - size_t len; -} upb_FindUnknownRet; - -// Finds first occurrence of unknown data by tag id in message. -upb_FindUnknownRet upb_MiniTable_FindUnknown(const upb_Message* msg, - uint32_t field_number, - int depth_limit); - -typedef enum { - kUpb_UnknownToMessage_Ok, - kUpb_UnknownToMessage_ParseError, - kUpb_UnknownToMessage_OutOfMemory, - kUpb_UnknownToMessage_NotFound, -} upb_UnknownToMessage_Status; - -typedef struct { - upb_UnknownToMessage_Status status; - upb_Message* message; -} upb_UnknownToMessageRet; - -// Promotes unknown data inside message to a upb_Message parsing the unknown. -// -// The unknown data is removed from message after field value is set -// using upb_Message_SetMessage. -// -// WARNING!: See b/267655898 -upb_UnknownToMessageRet upb_MiniTable_PromoteUnknownToMessage( - upb_Message* msg, const upb_MiniTable* mini_table, - const upb_MiniTableField* field, const upb_MiniTable* sub_mini_table, - int decode_options, upb_Arena* arena); - -// Promotes all unknown data that matches field tag id to repeated messages -// in upb_Array. -// -// The unknown data is removed from message after upb_Array is populated. -// Since repeated messages can't be packed we remove each unknown that -// contains the target tag id. -upb_UnknownToMessage_Status upb_MiniTable_PromoteUnknownToMessageArray( - upb_Message* msg, const upb_MiniTableField* field, - const upb_MiniTable* mini_table, int decode_options, upb_Arena* arena); - -// Promotes all unknown data that matches field tag id to upb_Map. -// -// The unknown data is removed from message after upb_Map is populated. -// Since repeated messages can't be packed we remove each unknown that -// contains the target tag id. -upb_UnknownToMessage_Status upb_MiniTable_PromoteUnknownToMap( - upb_Message* msg, const upb_MiniTable* mini_table, - const upb_MiniTableField* field, int decode_options, upb_Arena* arena); - #ifdef __cplusplus } /* extern "C" */ #endif diff --git a/upb/message/accessors_test.cc b/upb/message/accessors_test.cc index 910a700994..a46f60254e 100644 --- a/upb/message/accessors_test.cc +++ b/upb/message/accessors_test.cc @@ -392,395 +392,6 @@ TEST(GeneratedCode, GetMutableMessage) { upb_Arena_Free(arena); } -TEST(GeneratedCode, FindUnknown) { - upb_Arena* arena = upb_Arena_New(); - upb_test_ModelWithExtensions* msg = upb_test_ModelWithExtensions_new(arena); - upb_test_ModelWithExtensions_set_random_int32(msg, 10); - upb_test_ModelWithExtensions_set_random_name( - msg, upb_StringView_FromString("Hello")); - - upb_test_ModelExtension1* extension1 = upb_test_ModelExtension1_new(arena); - upb_test_ModelExtension1_set_str(extension1, - upb_StringView_FromString("World")); - - upb_test_ModelExtension1_set_model_ext(msg, extension1, arena); - - size_t serialized_size; - char* serialized = - upb_test_ModelWithExtensions_serialize(msg, arena, &serialized_size); - - upb_test_EmptyMessageWithExtensions* base_msg = - upb_test_EmptyMessageWithExtensions_parse(serialized, serialized_size, - arena); - - upb_FindUnknownRet result = upb_MiniTable_FindUnknown( - base_msg, upb_test_ModelExtension1_model_ext_ext.field.number, - kUpb_WireFormat_DefaultDepthLimit); - EXPECT_EQ(kUpb_FindUnknown_Ok, result.status); - - result = upb_MiniTable_FindUnknown( - base_msg, upb_test_ModelExtension2_model_ext_ext.field.number, - kUpb_WireFormat_DefaultDepthLimit); - EXPECT_EQ(kUpb_FindUnknown_NotPresent, result.status); - - upb_Arena_Free(arena); -} - -TEST(GeneratedCode, Extensions) { - upb_Arena* arena = upb_Arena_New(); - upb_test_ModelWithExtensions* msg = upb_test_ModelWithExtensions_new(arena); - upb_test_ModelWithExtensions_set_random_int32(msg, 10); - upb_test_ModelWithExtensions_set_random_name( - msg, upb_StringView_FromString("Hello")); - - upb_test_ModelExtension1* extension1 = upb_test_ModelExtension1_new(arena); - upb_test_ModelExtension1_set_str(extension1, - upb_StringView_FromString("World")); - - upb_test_ModelExtension2* extension2 = upb_test_ModelExtension2_new(arena); - upb_test_ModelExtension2_set_i(extension2, 5); - - upb_test_ModelExtension2* extension3 = upb_test_ModelExtension2_new(arena); - upb_test_ModelExtension2_set_i(extension3, 6); - - upb_test_ModelExtension2* extension4 = upb_test_ModelExtension2_new(arena); - upb_test_ModelExtension2_set_i(extension4, 7); - - upb_test_ModelExtension2* extension5 = upb_test_ModelExtension2_new(arena); - upb_test_ModelExtension2_set_i(extension5, 8); - - upb_test_ModelExtension2* extension6 = upb_test_ModelExtension2_new(arena); - upb_test_ModelExtension2_set_i(extension6, 9); - - // Set many extensions, to exercise code paths that involve reallocating the - // extensions and unknown fields array. - upb_test_ModelExtension1_set_model_ext(msg, extension1, arena); - upb_test_ModelExtension2_set_model_ext(msg, extension2, arena); - upb_test_ModelExtension2_set_model_ext_2(msg, extension3, arena); - upb_test_ModelExtension2_set_model_ext_3(msg, extension4, arena); - upb_test_ModelExtension2_set_model_ext_4(msg, extension5, arena); - upb_test_ModelExtension2_set_model_ext_5(msg, extension6, arena); - - size_t serialized_size; - char* serialized = - upb_test_ModelWithExtensions_serialize(msg, arena, &serialized_size); - - const upb_Message_Extension* upb_ext2; - upb_test_ModelExtension1* ext1; - upb_test_ModelExtension2* ext2; - upb_GetExtension_Status promote_status; - - // Test known GetExtension 1 - promote_status = upb_MiniTable_GetOrPromoteExtension( - msg, &upb_test_ModelExtension1_model_ext_ext, 0, arena, &upb_ext2); - ext1 = (upb_test_ModelExtension1*)upb_ext2->data.ptr; - EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); - EXPECT_TRUE(upb_StringView_IsEqual(upb_StringView_FromString("World"), - upb_test_ModelExtension1_str(ext1))); - - // Test known GetExtension 2 - promote_status = upb_MiniTable_GetOrPromoteExtension( - msg, &upb_test_ModelExtension2_model_ext_ext, 0, arena, &upb_ext2); - ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; - EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); - EXPECT_EQ(5, upb_test_ModelExtension2_i(ext2)); - - // Test known GetExtension 3 - promote_status = upb_MiniTable_GetOrPromoteExtension( - msg, &upb_test_ModelExtension2_model_ext_2_ext, 0, arena, &upb_ext2); - ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; - EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); - EXPECT_EQ(6, upb_test_ModelExtension2_i(ext2)); - - // Test known GetExtension 4 - promote_status = upb_MiniTable_GetOrPromoteExtension( - msg, &upb_test_ModelExtension2_model_ext_3_ext, 0, arena, &upb_ext2); - ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; - EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); - EXPECT_EQ(7, upb_test_ModelExtension2_i(ext2)); - - // Test known GetExtension 5 - promote_status = upb_MiniTable_GetOrPromoteExtension( - msg, &upb_test_ModelExtension2_model_ext_4_ext, 0, arena, &upb_ext2); - ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; - EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); - EXPECT_EQ(8, upb_test_ModelExtension2_i(ext2)); - - // Test known GetExtension 6 - promote_status = upb_MiniTable_GetOrPromoteExtension( - msg, &upb_test_ModelExtension2_model_ext_5_ext, 0, arena, &upb_ext2); - ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; - EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); - EXPECT_EQ(9, upb_test_ModelExtension2_i(ext2)); - - upb_test_EmptyMessageWithExtensions* base_msg = - upb_test_EmptyMessageWithExtensions_parse(serialized, serialized_size, - arena); - - // Get unknown extension bytes before promotion. - const char* extension_data; - size_t len; - upb_GetExtensionAsBytes_Status status = upb_MiniTable_GetExtensionAsBytes( - base_msg, &upb_test_ModelExtension2_model_ext_ext, 0, arena, - &extension_data, &len); - EXPECT_EQ(kUpb_GetExtensionAsBytes_Ok, status); - EXPECT_EQ(0x48, extension_data[0]); - EXPECT_EQ(5, extension_data[1]); - - // Test unknown GetExtension. - promote_status = upb_MiniTable_GetOrPromoteExtension( - base_msg, &upb_test_ModelExtension1_model_ext_ext, 0, arena, &upb_ext2); - ext1 = (upb_test_ModelExtension1*)upb_ext2->data.ptr; - EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); - EXPECT_TRUE(upb_StringView_IsEqual(upb_StringView_FromString("World"), - upb_test_ModelExtension1_str(ext1))); - - // Test unknown GetExtension. - promote_status = upb_MiniTable_GetOrPromoteExtension( - base_msg, &upb_test_ModelExtension2_model_ext_ext, 0, arena, &upb_ext2); - ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; - EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); - EXPECT_EQ(5, upb_test_ModelExtension2_i(ext2)); - - // Test unknown GetExtension. - promote_status = upb_MiniTable_GetOrPromoteExtension( - base_msg, &upb_test_ModelExtension2_model_ext_2_ext, 0, arena, &upb_ext2); - ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; - EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); - EXPECT_EQ(6, upb_test_ModelExtension2_i(ext2)); - - // Test unknown GetExtension. - promote_status = upb_MiniTable_GetOrPromoteExtension( - base_msg, &upb_test_ModelExtension2_model_ext_3_ext, 0, arena, &upb_ext2); - ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; - EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); - EXPECT_EQ(7, upb_test_ModelExtension2_i(ext2)); - - // Test unknown GetExtension. - promote_status = upb_MiniTable_GetOrPromoteExtension( - base_msg, &upb_test_ModelExtension2_model_ext_4_ext, 0, arena, &upb_ext2); - ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; - EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); - EXPECT_EQ(8, upb_test_ModelExtension2_i(ext2)); - - upb_Arena_Free(arena); -} - -// Create a minitable to mimic ModelWithSubMessages with unlinked subs -// to lazily promote unknowns after parsing. -upb_MiniTable* CreateMiniTableWithEmptySubTables(upb_Arena* arena) { - upb::MtDataEncoder e; - e.StartMessage(0); - e.PutField(kUpb_FieldType_Int32, 4, 0); - e.PutField(kUpb_FieldType_Message, 5, 0); - e.PutField(kUpb_FieldType_Message, 6, kUpb_FieldModifier_IsRepeated); - - upb_Status status; - upb_Status_Clear(&status); - upb_MiniTable* table = - upb_MiniTable_Build(e.data().data(), e.data().size(), arena, &status); - EXPECT_EQ(status.ok, true); - // Initialize sub table to null. Not using upb_MiniTable_SetSubMessage - // since it checks ->ext on parameter. - upb_MiniTableSub* sub = const_cast( - &table->subs[table->fields[1].UPB_PRIVATE(submsg_index)]); - sub->submsg = nullptr; - sub = const_cast( - &table->subs[table->fields[2].UPB_PRIVATE(submsg_index)]); - sub->submsg = nullptr; - return table; -} - -// Create a minitable to mimic ModelWithMaps with unlinked subs -// to lazily promote unknowns after parsing. -upb_MiniTable* CreateMiniTableWithEmptySubTablesForMaps(upb_Arena* arena) { - upb::MtDataEncoder e; - e.StartMessage(0); - e.PutField(kUpb_FieldType_Int32, 1, 0); - e.PutField(kUpb_FieldType_Message, 3, kUpb_FieldModifier_IsRepeated); - e.PutField(kUpb_FieldType_Message, 4, kUpb_FieldModifier_IsRepeated); - - upb_Status status; - upb_Status_Clear(&status); - upb_MiniTable* table = - upb_MiniTable_Build(e.data().data(), e.data().size(), arena, &status); - EXPECT_EQ(status.ok, true); - // Initialize sub table to null. Not using upb_MiniTable_SetSubMessage - // since it checks ->ext on parameter. - upb_MiniTableSub* sub = const_cast( - &table->subs[table->fields[1].UPB_PRIVATE(submsg_index)]); - sub->submsg = nullptr; - sub = const_cast( - &table->subs[table->fields[2].UPB_PRIVATE(submsg_index)]); - sub->submsg = nullptr; - return table; -} - -upb_MiniTable* CreateMapEntryMiniTable(upb_Arena* arena) { - upb::MtDataEncoder e; - e.EncodeMap(kUpb_FieldType_String, kUpb_FieldType_String, 0, 0); - upb_Status status; - upb_Status_Clear(&status); - upb_MiniTable* table = - upb_MiniTable_Build(e.data().data(), e.data().size(), arena, &status); - EXPECT_EQ(status.ok, true); - return table; -} - -TEST(GeneratedCode, PromoteUnknownMessage) { - upb_Arena* arena = upb_Arena_New(); - upb_test_ModelWithSubMessages* input_msg = - upb_test_ModelWithSubMessages_new(arena); - upb_test_ModelWithExtensions* sub_message = - upb_test_ModelWithExtensions_new(arena); - upb_test_ModelWithSubMessages_set_id(input_msg, 11); - upb_test_ModelWithExtensions_set_random_int32(sub_message, 12); - upb_test_ModelWithSubMessages_set_optional_child(input_msg, sub_message); - size_t serialized_size; - char* serialized = upb_test_ModelWithSubMessages_serialize(input_msg, arena, - &serialized_size); - - upb_MiniTable* mini_table = CreateMiniTableWithEmptySubTables(arena); - upb_Message* msg = _upb_Message_New(mini_table, arena); - upb_DecodeStatus decode_status = upb_Decode(serialized, serialized_size, msg, - mini_table, nullptr, 0, arena); - EXPECT_EQ(decode_status, kUpb_DecodeStatus_Ok); - int32_t val = upb_Message_GetInt32( - msg, upb_MiniTable_FindFieldByNumber(mini_table, 4), 0); - EXPECT_EQ(val, 11); - upb_FindUnknownRet unknown = - upb_MiniTable_FindUnknown(msg, 5, kUpb_WireFormat_DefaultDepthLimit); - EXPECT_EQ(unknown.status, kUpb_FindUnknown_Ok); - // Update mini table and promote unknown to a message. - EXPECT_TRUE(upb_MiniTable_SetSubMessage( - mini_table, (upb_MiniTableField*)&mini_table->fields[1], - &upb_test_ModelWithExtensions_msg_init)); - const int decode_options = upb_DecodeOptions_MaxDepth( - kUpb_WireFormat_DefaultDepthLimit); // UPB_DECODE_ALIAS disabled. - upb_UnknownToMessageRet promote_result = - upb_MiniTable_PromoteUnknownToMessage( - msg, mini_table, &mini_table->fields[1], - &upb_test_ModelWithExtensions_msg_init, decode_options, arena); - EXPECT_EQ(promote_result.status, kUpb_UnknownToMessage_Ok); - const upb_Message* promoted_message = - upb_Message_GetMessage(msg, &mini_table->fields[1], nullptr); - EXPECT_EQ(upb_test_ModelWithExtensions_random_int32( - (upb_test_ModelWithExtensions*)promoted_message), - 12); - upb_Arena_Free(arena); -} - -TEST(GeneratedCode, PromoteUnknownRepeatedMessage) { - upb_Arena* arena = upb_Arena_New(); - upb_test_ModelWithSubMessages* input_msg = - upb_test_ModelWithSubMessages_new(arena); - upb_test_ModelWithSubMessages_set_id(input_msg, 123); - - // Add 2 repeated messages to input_msg. - upb_test_ModelWithExtensions* item = - upb_test_ModelWithSubMessages_add_items(input_msg, arena); - upb_test_ModelWithExtensions_set_random_int32(item, 5); - item = upb_test_ModelWithSubMessages_add_items(input_msg, arena); - upb_test_ModelWithExtensions_set_random_int32(item, 6); - - size_t serialized_size; - char* serialized = upb_test_ModelWithSubMessages_serialize(input_msg, arena, - &serialized_size); - - upb_MiniTable* mini_table = CreateMiniTableWithEmptySubTables(arena); - upb_Message* msg = _upb_Message_New(mini_table, arena); - upb_DecodeStatus decode_status = upb_Decode(serialized, serialized_size, msg, - mini_table, nullptr, 0, arena); - EXPECT_EQ(decode_status, kUpb_DecodeStatus_Ok); - int32_t val = upb_Message_GetInt32( - msg, upb_MiniTable_FindFieldByNumber(mini_table, 4), 0); - EXPECT_EQ(val, 123); - - // Check that we have repeated field data in an unknown. - upb_FindUnknownRet unknown = - upb_MiniTable_FindUnknown(msg, 6, kUpb_WireFormat_DefaultDepthLimit); - EXPECT_EQ(unknown.status, kUpb_FindUnknown_Ok); - - // Update mini table and promote unknown to a message. - EXPECT_TRUE(upb_MiniTable_SetSubMessage( - mini_table, (upb_MiniTableField*)&mini_table->fields[2], - &upb_test_ModelWithExtensions_msg_init)); - const int decode_options = upb_DecodeOptions_MaxDepth( - kUpb_WireFormat_DefaultDepthLimit); // UPB_DECODE_ALIAS disabled. - upb_UnknownToMessage_Status promote_result = - upb_MiniTable_PromoteUnknownToMessageArray( - msg, &mini_table->fields[2], &upb_test_ModelWithExtensions_msg_init, - decode_options, arena); - EXPECT_EQ(promote_result, kUpb_UnknownToMessage_Ok); - - upb_Array* array = upb_Message_GetMutableArray(msg, &mini_table->fields[2]); - const upb_Message* promoted_message = upb_Array_Get(array, 0).msg_val; - EXPECT_EQ(upb_test_ModelWithExtensions_random_int32( - (upb_test_ModelWithExtensions*)promoted_message), - 5); - promoted_message = upb_Array_Get(array, 1).msg_val; - EXPECT_EQ(upb_test_ModelWithExtensions_random_int32( - (upb_test_ModelWithExtensions*)promoted_message), - 6); - upb_Arena_Free(arena); -} - -TEST(GeneratedCode, PromoteUnknownToMap) { - upb_Arena* arena = upb_Arena_New(); - upb_test_ModelWithMaps* input_msg = upb_test_ModelWithMaps_new(arena); - upb_test_ModelWithMaps_set_id(input_msg, 123); - - // Add 2 map entries. - upb_test_ModelWithMaps_map_ss_set(input_msg, - upb_StringView_FromString("key1"), - upb_StringView_FromString("value1"), arena); - upb_test_ModelWithMaps_map_ss_set(input_msg, - upb_StringView_FromString("key2"), - upb_StringView_FromString("value2"), arena); - - size_t serialized_size; - char* serialized = - upb_test_ModelWithMaps_serialize(input_msg, arena, &serialized_size); - - upb_MiniTable* mini_table = CreateMiniTableWithEmptySubTablesForMaps(arena); - upb_MiniTable* map_entry_mini_table = CreateMapEntryMiniTable(arena); - upb_Message* msg = _upb_Message_New(mini_table, arena); - const int decode_options = - upb_DecodeOptions_MaxDepth(kUpb_WireFormat_DefaultDepthLimit); - upb_DecodeStatus decode_status = - upb_Decode(serialized, serialized_size, msg, mini_table, nullptr, - decode_options, arena); - EXPECT_EQ(decode_status, kUpb_DecodeStatus_Ok); - int32_t val = upb_Message_GetInt32( - msg, upb_MiniTable_FindFieldByNumber(mini_table, 1), 0); - EXPECT_EQ(val, 123); - - // Check that we have map data in an unknown. - upb_FindUnknownRet unknown = - upb_MiniTable_FindUnknown(msg, 3, kUpb_WireFormat_DefaultDepthLimit); - EXPECT_EQ(unknown.status, kUpb_FindUnknown_Ok); - - // Update mini table and promote unknown to a message. - EXPECT_TRUE(upb_MiniTable_SetSubMessage( - mini_table, (upb_MiniTableField*)&mini_table->fields[1], - map_entry_mini_table)); - upb_UnknownToMessage_Status promote_result = - upb_MiniTable_PromoteUnknownToMap(msg, mini_table, &mini_table->fields[1], - decode_options, arena); - EXPECT_EQ(promote_result, kUpb_UnknownToMessage_Ok); - - upb_Map* map = upb_Message_GetOrCreateMutableMap( - msg, map_entry_mini_table, &mini_table->fields[1], arena); - EXPECT_NE(map, nullptr); - // Lookup in map. - upb_MessageValue key; - key.str_val = upb_StringView_FromString("key2"); - upb_MessageValue value; - EXPECT_TRUE(upb_Map_Get(map, key, &value)); - EXPECT_EQ(0, strncmp(value.str_val.data, "value2", 5)); - upb_Arena_Free(arena); -} - TEST(GeneratedCode, EnumClosedCheck) { upb_Arena* arena = upb_Arena_New(); diff --git a/upb/message/message.h b/upb/message/message.h index 1cf7249c2d..dbf0e13979 100644 --- a/upb/message/message.h +++ b/upb/message/message.h @@ -25,11 +25,9 @@ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. */ -// Public APIs for message operations that do not require descriptors. -// These functions can be used even in build that does not want to depend on -// reflection or descriptors. +// Public APIs for message operations that do not depend on the schema. // -// Descriptor-based reflection functionality lives in reflection.h. +// MiniTable-based accessors live in accessors.h. #ifndef UPB_MESSAGE_MESSAGE_H_ #define UPB_MESSAGE_MESSAGE_H_ diff --git a/upb/message/promote.c b/upb/message/promote.c new file mode 100644 index 0000000000..3f2a2db930 --- /dev/null +++ b/upb/message/promote.c @@ -0,0 +1,319 @@ +/* + * Copyright (c) 2009-2021, Google LLC + * All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * * Neither the name of Google LLC nor the + * names of its contributors may be used to endorse or promote products + * derived from this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED. IN NO EVENT SHALL Google LLC BE LIABLE FOR ANY DIRECT, + * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES + * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; + * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND + * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS + * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ + +#include "upb/message/promote.h" + +#include "upb/collections/array.h" +#include "upb/collections/array_internal.h" +#include "upb/collections/map.h" +#include "upb/message/accessors.h" +#include "upb/message/message.h" +#include "upb/mini_table/field_internal.h" +#include "upb/wire/common.h" +#include "upb/wire/decode.h" +#include "upb/wire/encode.h" +#include "upb/wire/eps_copy_input_stream.h" +#include "upb/wire/reader.h" + +// Must be last. +#include "upb/port/def.inc" + +// Parses unknown data by merging into existing base_message or creating a +// new message usingg mini_table. +static upb_UnknownToMessageRet upb_MiniTable_ParseUnknownMessage( + const char* unknown_data, size_t unknown_size, + const upb_MiniTable* mini_table, upb_Message* base_message, + int decode_options, upb_Arena* arena) { + upb_UnknownToMessageRet ret; + ret.message = + base_message ? base_message : _upb_Message_New(mini_table, arena); + if (!ret.message) { + ret.status = kUpb_UnknownToMessage_OutOfMemory; + return ret; + } + // Decode sub message using unknown field contents. + const char* data = unknown_data; + uint32_t tag; + uint64_t message_len = 0; + data = upb_WireReader_ReadTag(data, &tag); + data = upb_WireReader_ReadVarint(data, &message_len); + upb_DecodeStatus status = upb_Decode(data, message_len, ret.message, + mini_table, NULL, decode_options, arena); + if (status == kUpb_DecodeStatus_OutOfMemory) { + ret.status = kUpb_UnknownToMessage_OutOfMemory; + } else if (status == kUpb_DecodeStatus_Ok) { + ret.status = kUpb_UnknownToMessage_Ok; + } else { + ret.status = kUpb_UnknownToMessage_ParseError; + } + return ret; +} + +upb_GetExtension_Status upb_MiniTable_GetOrPromoteExtension( + upb_Message* msg, const upb_MiniTableExtension* ext_table, + int decode_options, upb_Arena* arena, + const upb_Message_Extension** extension) { + UPB_ASSERT(upb_MiniTableField_CType(&ext_table->field) == kUpb_CType_Message); + *extension = _upb_Message_Getext(msg, ext_table); + if (*extension) { + return kUpb_GetExtension_Ok; + } + + // Check unknown fields, if available promote. + int field_number = ext_table->field.number; + upb_FindUnknownRet result = upb_MiniTable_FindUnknown( + msg, field_number, kUpb_WireFormat_DefaultDepthLimit); + if (result.status != kUpb_FindUnknown_Ok) { + return kUpb_GetExtension_NotPresent; + } + size_t len; + size_t ofs = result.ptr - upb_Message_GetUnknown(msg, &len); + // Decode and promote from unknown. + const upb_MiniTable* extension_table = ext_table->sub.submsg; + upb_UnknownToMessageRet parse_result = upb_MiniTable_ParseUnknownMessage( + result.ptr, result.len, extension_table, + /* base_message= */ NULL, decode_options, arena); + switch (parse_result.status) { + case kUpb_UnknownToMessage_OutOfMemory: + return kUpb_GetExtension_OutOfMemory; + case kUpb_UnknownToMessage_ParseError: + return kUpb_GetExtension_ParseError; + case kUpb_UnknownToMessage_NotFound: + return kUpb_GetExtension_NotPresent; + case kUpb_UnknownToMessage_Ok: + break; + } + upb_Message* extension_msg = parse_result.message; + // Add to extensions. + upb_Message_Extension* ext = + _upb_Message_GetOrCreateExtension(msg, ext_table, arena); + if (!ext) { + return kUpb_GetExtension_OutOfMemory; + } + memcpy(&ext->data, &extension_msg, sizeof(extension_msg)); + *extension = ext; + const char* delete_ptr = upb_Message_GetUnknown(msg, &len) + ofs; + upb_Message_DeleteUnknown(msg, delete_ptr, result.len); + return kUpb_GetExtension_Ok; +} + +upb_GetExtensionAsBytes_Status upb_MiniTable_GetExtensionAsBytes( + const upb_Message* msg, const upb_MiniTableExtension* ext_table, + int encode_options, upb_Arena* arena, const char** extension_data, + size_t* len) { + const upb_Message_Extension* msg_ext = _upb_Message_Getext(msg, ext_table); + UPB_ASSERT(upb_MiniTableField_CType(&ext_table->field) == kUpb_CType_Message); + if (msg_ext) { + upb_EncodeStatus status = + upb_Encode(msg_ext->data.ptr, msg_ext->ext->sub.submsg, encode_options, + arena, (char**)extension_data, len); + if (status != kUpb_EncodeStatus_Ok) { + return kUpb_GetExtensionAsBytes_EncodeError; + } + return kUpb_GetExtensionAsBytes_Ok; + } + int field_number = ext_table->field.number; + upb_FindUnknownRet result = upb_MiniTable_FindUnknown( + msg, field_number, upb_DecodeOptions_GetMaxDepth(encode_options)); + if (result.status != kUpb_FindUnknown_Ok) { + return kUpb_GetExtensionAsBytes_NotPresent; + } + const char* data = result.ptr; + uint32_t tag; + uint64_t message_len = 0; + data = upb_WireReader_ReadTag(data, &tag); + data = upb_WireReader_ReadVarint(data, &message_len); + *extension_data = data; + *len = message_len; + return kUpb_GetExtensionAsBytes_Ok; +} + +static upb_FindUnknownRet upb_FindUnknownRet_ParseError(void) { + return (upb_FindUnknownRet){.status = kUpb_FindUnknown_ParseError}; +} + +upb_FindUnknownRet upb_MiniTable_FindUnknown(const upb_Message* msg, + uint32_t field_number, + int depth_limit) { + size_t size; + upb_FindUnknownRet ret; + + const char* ptr = upb_Message_GetUnknown(msg, &size); + upb_EpsCopyInputStream stream; + upb_EpsCopyInputStream_Init(&stream, &ptr, size, true); + + while (!upb_EpsCopyInputStream_IsDone(&stream, &ptr)) { + uint32_t tag; + const char* unknown_begin = ptr; + ptr = upb_WireReader_ReadTag(ptr, &tag); + if (!ptr) return upb_FindUnknownRet_ParseError(); + if (field_number == upb_WireReader_GetFieldNumber(tag)) { + ret.status = kUpb_FindUnknown_Ok; + ret.ptr = upb_EpsCopyInputStream_GetAliasedPtr(&stream, unknown_begin); + ptr = _upb_WireReader_SkipValue(ptr, tag, depth_limit, &stream); + // Because we know that the input is a flat buffer, it is safe to perform + // pointer arithmetic on aliased pointers. + ret.len = upb_EpsCopyInputStream_GetAliasedPtr(&stream, ptr) - ret.ptr; + return ret; + } + + ptr = _upb_WireReader_SkipValue(ptr, tag, depth_limit, &stream); + if (!ptr) return upb_FindUnknownRet_ParseError(); + } + ret.status = kUpb_FindUnknown_NotPresent; + ret.ptr = NULL; + ret.len = 0; + return ret; +} + +// Warning: See TODO(b/267655898) +upb_UnknownToMessageRet upb_MiniTable_PromoteUnknownToMessage( + upb_Message* msg, const upb_MiniTable* mini_table, + const upb_MiniTableField* field, const upb_MiniTable* sub_mini_table, + int decode_options, upb_Arena* arena) { + upb_FindUnknownRet unknown; + // We need to loop and merge unknowns that have matching tag field->number. + upb_Message* message = NULL; + // Callers should check that message is not set first before calling + // PromotoUnknownToMessage. + UPB_ASSERT(upb_MiniTable_GetSubMessageTable(mini_table, field) == + sub_mini_table); + bool is_oneof = _upb_MiniTableField_InOneOf(field); + if (!is_oneof || _upb_getoneofcase_field(msg, field) == field->number) { + UPB_ASSERT(upb_Message_GetMessage(msg, field, NULL) == NULL); + } + upb_UnknownToMessageRet ret; + ret.status = kUpb_UnknownToMessage_Ok; + do { + unknown = upb_MiniTable_FindUnknown( + msg, field->number, upb_DecodeOptions_GetMaxDepth(decode_options)); + switch (unknown.status) { + case kUpb_FindUnknown_Ok: { + const char* unknown_data = unknown.ptr; + size_t unknown_size = unknown.len; + ret = upb_MiniTable_ParseUnknownMessage(unknown_data, unknown_size, + sub_mini_table, message, + decode_options, arena); + if (ret.status == kUpb_UnknownToMessage_Ok) { + message = ret.message; + upb_Message_DeleteUnknown(msg, unknown_data, unknown_size); + } + } break; + case kUpb_FindUnknown_ParseError: + ret.status = kUpb_UnknownToMessage_ParseError; + break; + case kUpb_FindUnknown_NotPresent: + // If we parsed at least one unknown, we are done. + ret.status = + message ? kUpb_UnknownToMessage_Ok : kUpb_UnknownToMessage_NotFound; + break; + } + } while (unknown.status == kUpb_FindUnknown_Ok); + if (message) { + if (is_oneof) { + *_upb_oneofcase_field(msg, field) = field->number; + } + upb_Message_SetMessage(msg, mini_table, field, message); + ret.message = message; + } + return ret; +} + +// Moves repeated messages in unknowns to a upb_Array. +// +// Since the repeated field is not a scalar type we don't check for +// kUpb_LabelFlags_IsPacked. +// TODO(b/251007554): Optimize. Instead of converting messages one at a time, +// scan all unknown data once and compact. +upb_UnknownToMessage_Status upb_MiniTable_PromoteUnknownToMessageArray( + upb_Message* msg, const upb_MiniTableField* field, + const upb_MiniTable* mini_table, int decode_options, upb_Arena* arena) { + upb_Array* repeated_messages = upb_Message_GetMutableArray(msg, field); + // Find all unknowns with given field number and parse. + upb_FindUnknownRet unknown; + do { + unknown = upb_MiniTable_FindUnknown( + msg, field->number, upb_DecodeOptions_GetMaxDepth(decode_options)); + if (unknown.status == kUpb_FindUnknown_Ok) { + upb_UnknownToMessageRet ret = upb_MiniTable_ParseUnknownMessage( + unknown.ptr, unknown.len, mini_table, + /* base_message= */ NULL, decode_options, arena); + if (ret.status == kUpb_UnknownToMessage_Ok) { + upb_MessageValue value; + value.msg_val = ret.message; + // Allocate array on demand before append. + if (!repeated_messages) { + upb_Message_ResizeArray(msg, field, 0, arena); + repeated_messages = upb_Message_GetMutableArray(msg, field); + } + if (!upb_Array_Append(repeated_messages, value, arena)) { + return kUpb_UnknownToMessage_OutOfMemory; + } + upb_Message_DeleteUnknown(msg, unknown.ptr, unknown.len); + } else { + return ret.status; + } + } + } while (unknown.status == kUpb_FindUnknown_Ok); + return kUpb_UnknownToMessage_Ok; +} + +// Moves repeated messages in unknowns to a upb_Map. +upb_UnknownToMessage_Status upb_MiniTable_PromoteUnknownToMap( + upb_Message* msg, const upb_MiniTable* mini_table, + const upb_MiniTableField* field, int decode_options, upb_Arena* arena) { + const upb_MiniTable* map_entry_mini_table = + mini_table->subs[field->UPB_PRIVATE(submsg_index)].submsg; + UPB_ASSERT(map_entry_mini_table); + UPB_ASSERT(map_entry_mini_table); + UPB_ASSERT(map_entry_mini_table->field_count == 2); + UPB_ASSERT(upb_FieldMode_Get(field) == kUpb_FieldMode_Map); + // Find all unknowns with given field number and parse. + upb_FindUnknownRet unknown; + while (1) { + unknown = upb_MiniTable_FindUnknown( + msg, field->number, upb_DecodeOptions_GetMaxDepth(decode_options)); + if (unknown.status != kUpb_FindUnknown_Ok) break; + upb_UnknownToMessageRet ret = upb_MiniTable_ParseUnknownMessage( + unknown.ptr, unknown.len, map_entry_mini_table, + /* base_message= */ NULL, decode_options, arena); + if (ret.status != kUpb_UnknownToMessage_Ok) return ret.status; + // Allocate map on demand before append. + upb_Map* map = upb_Message_GetOrCreateMutableMap(msg, map_entry_mini_table, + field, arena); + upb_Message* map_entry_message = ret.message; + upb_MapInsertStatus insert_status = upb_Message_InsertMapEntry( + map, mini_table, field, map_entry_message, arena); + if (insert_status == kUpb_MapInsertStatus_OutOfMemory) { + return kUpb_UnknownToMessage_OutOfMemory; + } + UPB_ASSUME(insert_status == kUpb_MapInsertStatus_Inserted || + insert_status == kUpb_MapInsertStatus_Replaced); + upb_Message_DeleteUnknown(msg, unknown.ptr, unknown.len); + } + return kUpb_UnknownToMessage_Ok; +} diff --git a/upb/message/promote.h b/upb/message/promote.h new file mode 100644 index 0000000000..36b45aab45 --- /dev/null +++ b/upb/message/promote.h @@ -0,0 +1,130 @@ +/* + * Copyright (c) 2009-2022, Google LLC + * All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * * Neither the name of Google LLC nor the + * names of its contributors may be used to endorse or promote products + * derived from this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED. IN NO EVENT SHALL Google LLC BE LIABLE FOR ANY DIRECT, + * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES + * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; + * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND + * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS + * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ + +#ifndef UPB_MESSAGE_PROMOTE_H_ +#define UPB_MESSAGE_PROMOTE_H_ + +#include "upb/message/extension_internal.h" + +// Must be last. +#include "upb/port/def.inc" + +#ifdef __cplusplus +extern "C" { +#endif + +typedef enum { + kUpb_GetExtension_Ok, + kUpb_GetExtension_NotPresent, + kUpb_GetExtension_ParseError, + kUpb_GetExtension_OutOfMemory, +} upb_GetExtension_Status; + +typedef enum { + kUpb_GetExtensionAsBytes_Ok, + kUpb_GetExtensionAsBytes_NotPresent, + kUpb_GetExtensionAsBytes_EncodeError, +} upb_GetExtensionAsBytes_Status; + +// Returns a message extension or promotes an unknown field to +// an extension. +// +// TODO(ferhat): Only supports extension fields that are messages, +// expand support to include non-message types. +upb_GetExtension_Status upb_MiniTable_GetOrPromoteExtension( + upb_Message* msg, const upb_MiniTableExtension* ext_table, + int decode_options, upb_Arena* arena, + const upb_Message_Extension** extension); + +typedef enum { + kUpb_FindUnknown_Ok, + kUpb_FindUnknown_NotPresent, + kUpb_FindUnknown_ParseError, +} upb_FindUnknown_Status; + +typedef struct { + upb_FindUnknown_Status status; + // Start of unknown field data in message arena. + const char* ptr; + // Size of unknown field data. + size_t len; +} upb_FindUnknownRet; + +// Finds first occurrence of unknown data by tag id in message. +upb_FindUnknownRet upb_MiniTable_FindUnknown(const upb_Message* msg, + uint32_t field_number, + int depth_limit); + +typedef enum { + kUpb_UnknownToMessage_Ok, + kUpb_UnknownToMessage_ParseError, + kUpb_UnknownToMessage_OutOfMemory, + kUpb_UnknownToMessage_NotFound, +} upb_UnknownToMessage_Status; + +typedef struct { + upb_UnknownToMessage_Status status; + upb_Message* message; +} upb_UnknownToMessageRet; + +// Promotes unknown data inside message to a upb_Message parsing the unknown. +// +// The unknown data is removed from message after field value is set +// using upb_Message_SetMessage. +// +// WARNING!: See b/267655898 +upb_UnknownToMessageRet upb_MiniTable_PromoteUnknownToMessage( + upb_Message* msg, const upb_MiniTable* mini_table, + const upb_MiniTableField* field, const upb_MiniTable* sub_mini_table, + int decode_options, upb_Arena* arena); + +// Promotes all unknown data that matches field tag id to repeated messages +// in upb_Array. +// +// The unknown data is removed from message after upb_Array is populated. +// Since repeated messages can't be packed we remove each unknown that +// contains the target tag id. +upb_UnknownToMessage_Status upb_MiniTable_PromoteUnknownToMessageArray( + upb_Message* msg, const upb_MiniTableField* field, + const upb_MiniTable* mini_table, int decode_options, upb_Arena* arena); + +// Promotes all unknown data that matches field tag id to upb_Map. +// +// The unknown data is removed from message after upb_Map is populated. +// Since repeated messages can't be packed we remove each unknown that +// contains the target tag id. +upb_UnknownToMessage_Status upb_MiniTable_PromoteUnknownToMap( + upb_Message* msg, const upb_MiniTable* mini_table, + const upb_MiniTableField* field, int decode_options, upb_Arena* arena); + +#ifdef __cplusplus +} /* extern "C" */ +#endif + +#include "upb/port/undef.inc" + +#endif // UPB_MESSAGE_PROMOTE_H_ diff --git a/upb/message/promote_test.cc b/upb/message/promote_test.cc new file mode 100644 index 0000000000..98d98c56f4 --- /dev/null +++ b/upb/message/promote_test.cc @@ -0,0 +1,475 @@ +/* + * Copyright (c) 2009-2021, Google LLC + * All rights reserved. + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright + * notice, this list of conditions and the following disclaimer in the + * documentation and/or other materials provided with the distribution. + * * Neither the name of Google LLC nor the + * names of its contributors may be used to endorse or promote products + * derived from this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + * ARE DISCLAIMED. IN NO EVENT SHALL Google LLC BE LIABLE FOR ANY DIRECT, + * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES + * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; + * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND + * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS + * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ + +/* Test of mini table accessors. + * + * Messages are created and mutated using generated code, and then + * accessed through reflective APIs exposed through mini table accessors. + */ + +#include "upb/message/promote.h" + +#include + +#include "gtest/gtest.h" +#include "google/protobuf/test_messages_proto2.upb.h" +#include "google/protobuf/test_messages_proto3.upb.h" +#include "upb/base/string_view.h" +#include "upb/collections/array.h" +#include "upb/message/accessors.h" +#include "upb/mini_table/common.h" +#include "upb/mini_table/decode.h" +#include "upb/mini_table/encode_internal.hpp" +#include "upb/mini_table/field_internal.h" +#include "upb/test/test.upb.h" +#include "upb/upb.h" +#include "upb/wire/common.h" +#include "upb/wire/decode.h" + +// Must be last +#include "upb/port/def.inc" + +namespace { + +// Proto2 test messages field numbers used for reflective access. +const uint32_t kFieldOptionalInt32 = 1; +const uint32_t kFieldOptionalUInt32 = 3; +const uint32_t kFieldOptionalBool = 13; +const uint32_t kFieldOptionalString = 14; +const uint32_t kFieldOptionalNestedMessage = 18; +const uint32_t kFieldOptionalRepeatedInt32 = 31; +const uint32_t kFieldOptionalNestedMessageA = 1; +const uint32_t kFieldOptionalOneOfUInt32 = 111; +const uint32_t kFieldOptionalOneOfString = 113; + +const uint32_t kFieldProto3OptionalInt64 = 2; +const uint32_t kFieldProto3OptionalUInt64 = 4; + +const char kTestStr1[] = "Hello1"; +const char kTestStr2[] = "Hello2"; +const int32_t kTestInt32 = 567; +const int32_t kTestUInt32 = 0xF1234567; +const uint64_t kTestUInt64 = 0xFEDCBAFF87654321; + +TEST(GeneratedCode, FindUnknown) { + upb_Arena* arena = upb_Arena_New(); + upb_test_ModelWithExtensions* msg = upb_test_ModelWithExtensions_new(arena); + upb_test_ModelWithExtensions_set_random_int32(msg, 10); + upb_test_ModelWithExtensions_set_random_name( + msg, upb_StringView_FromString("Hello")); + + upb_test_ModelExtension1* extension1 = upb_test_ModelExtension1_new(arena); + upb_test_ModelExtension1_set_str(extension1, + upb_StringView_FromString("World")); + + upb_test_ModelExtension1_set_model_ext(msg, extension1, arena); + + size_t serialized_size; + char* serialized = + upb_test_ModelWithExtensions_serialize(msg, arena, &serialized_size); + + upb_test_EmptyMessageWithExtensions* base_msg = + upb_test_EmptyMessageWithExtensions_parse(serialized, serialized_size, + arena); + + upb_FindUnknownRet result = upb_MiniTable_FindUnknown( + base_msg, upb_test_ModelExtension1_model_ext_ext.field.number, + kUpb_WireFormat_DefaultDepthLimit); + EXPECT_EQ(kUpb_FindUnknown_Ok, result.status); + + result = upb_MiniTable_FindUnknown( + base_msg, upb_test_ModelExtension2_model_ext_ext.field.number, + kUpb_WireFormat_DefaultDepthLimit); + EXPECT_EQ(kUpb_FindUnknown_NotPresent, result.status); + + upb_Arena_Free(arena); +} + +TEST(GeneratedCode, Extensions) { + upb_Arena* arena = upb_Arena_New(); + upb_test_ModelWithExtensions* msg = upb_test_ModelWithExtensions_new(arena); + upb_test_ModelWithExtensions_set_random_int32(msg, 10); + upb_test_ModelWithExtensions_set_random_name( + msg, upb_StringView_FromString("Hello")); + + upb_test_ModelExtension1* extension1 = upb_test_ModelExtension1_new(arena); + upb_test_ModelExtension1_set_str(extension1, + upb_StringView_FromString("World")); + + upb_test_ModelExtension2* extension2 = upb_test_ModelExtension2_new(arena); + upb_test_ModelExtension2_set_i(extension2, 5); + + upb_test_ModelExtension2* extension3 = upb_test_ModelExtension2_new(arena); + upb_test_ModelExtension2_set_i(extension3, 6); + + upb_test_ModelExtension2* extension4 = upb_test_ModelExtension2_new(arena); + upb_test_ModelExtension2_set_i(extension4, 7); + + upb_test_ModelExtension2* extension5 = upb_test_ModelExtension2_new(arena); + upb_test_ModelExtension2_set_i(extension5, 8); + + upb_test_ModelExtension2* extension6 = upb_test_ModelExtension2_new(arena); + upb_test_ModelExtension2_set_i(extension6, 9); + + // Set many extensions, to exercise code paths that involve reallocating the + // extensions and unknown fields array. + upb_test_ModelExtension1_set_model_ext(msg, extension1, arena); + upb_test_ModelExtension2_set_model_ext(msg, extension2, arena); + upb_test_ModelExtension2_set_model_ext_2(msg, extension3, arena); + upb_test_ModelExtension2_set_model_ext_3(msg, extension4, arena); + upb_test_ModelExtension2_set_model_ext_4(msg, extension5, arena); + upb_test_ModelExtension2_set_model_ext_5(msg, extension6, arena); + + size_t serialized_size; + char* serialized = + upb_test_ModelWithExtensions_serialize(msg, arena, &serialized_size); + + const upb_Message_Extension* upb_ext2; + upb_test_ModelExtension1* ext1; + upb_test_ModelExtension2* ext2; + upb_GetExtension_Status promote_status; + + // Test known GetExtension 1 + promote_status = upb_MiniTable_GetOrPromoteExtension( + msg, &upb_test_ModelExtension1_model_ext_ext, 0, arena, &upb_ext2); + ext1 = (upb_test_ModelExtension1*)upb_ext2->data.ptr; + EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); + EXPECT_TRUE(upb_StringView_IsEqual(upb_StringView_FromString("World"), + upb_test_ModelExtension1_str(ext1))); + + // Test known GetExtension 2 + promote_status = upb_MiniTable_GetOrPromoteExtension( + msg, &upb_test_ModelExtension2_model_ext_ext, 0, arena, &upb_ext2); + ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; + EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); + EXPECT_EQ(5, upb_test_ModelExtension2_i(ext2)); + + // Test known GetExtension 3 + promote_status = upb_MiniTable_GetOrPromoteExtension( + msg, &upb_test_ModelExtension2_model_ext_2_ext, 0, arena, &upb_ext2); + ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; + EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); + EXPECT_EQ(6, upb_test_ModelExtension2_i(ext2)); + + // Test known GetExtension 4 + promote_status = upb_MiniTable_GetOrPromoteExtension( + msg, &upb_test_ModelExtension2_model_ext_3_ext, 0, arena, &upb_ext2); + ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; + EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); + EXPECT_EQ(7, upb_test_ModelExtension2_i(ext2)); + + // Test known GetExtension 5 + promote_status = upb_MiniTable_GetOrPromoteExtension( + msg, &upb_test_ModelExtension2_model_ext_4_ext, 0, arena, &upb_ext2); + ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; + EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); + EXPECT_EQ(8, upb_test_ModelExtension2_i(ext2)); + + // Test known GetExtension 6 + promote_status = upb_MiniTable_GetOrPromoteExtension( + msg, &upb_test_ModelExtension2_model_ext_5_ext, 0, arena, &upb_ext2); + ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; + EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); + EXPECT_EQ(9, upb_test_ModelExtension2_i(ext2)); + + upb_test_EmptyMessageWithExtensions* base_msg = + upb_test_EmptyMessageWithExtensions_parse(serialized, serialized_size, + arena); + + // Get unknown extension bytes before promotion. + size_t start_len; + upb_Message_GetUnknown(base_msg, &start_len); + EXPECT_GT(start_len, 0); + EXPECT_EQ(0, upb_Message_ExtensionCount(base_msg)); + + // Test unknown GetExtension. + promote_status = upb_MiniTable_GetOrPromoteExtension( + base_msg, &upb_test_ModelExtension1_model_ext_ext, 0, arena, &upb_ext2); + ext1 = (upb_test_ModelExtension1*)upb_ext2->data.ptr; + EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); + EXPECT_TRUE(upb_StringView_IsEqual(upb_StringView_FromString("World"), + upb_test_ModelExtension1_str(ext1))); + + // Test unknown GetExtension. + promote_status = upb_MiniTable_GetOrPromoteExtension( + base_msg, &upb_test_ModelExtension2_model_ext_ext, 0, arena, &upb_ext2); + ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; + EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); + EXPECT_EQ(5, upb_test_ModelExtension2_i(ext2)); + + // Test unknown GetExtension. + promote_status = upb_MiniTable_GetOrPromoteExtension( + base_msg, &upb_test_ModelExtension2_model_ext_2_ext, 0, arena, &upb_ext2); + ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; + EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); + EXPECT_EQ(6, upb_test_ModelExtension2_i(ext2)); + + // Test unknown GetExtension. + promote_status = upb_MiniTable_GetOrPromoteExtension( + base_msg, &upb_test_ModelExtension2_model_ext_3_ext, 0, arena, &upb_ext2); + ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; + EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); + EXPECT_EQ(7, upb_test_ModelExtension2_i(ext2)); + + // Test unknown GetExtension. + promote_status = upb_MiniTable_GetOrPromoteExtension( + base_msg, &upb_test_ModelExtension2_model_ext_4_ext, 0, arena, &upb_ext2); + ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; + EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); + EXPECT_EQ(8, upb_test_ModelExtension2_i(ext2)); + + // Test unknown GetExtension. + promote_status = upb_MiniTable_GetOrPromoteExtension( + base_msg, &upb_test_ModelExtension2_model_ext_5_ext, 0, arena, &upb_ext2); + ext2 = (upb_test_ModelExtension2*)upb_ext2->data.ptr; + EXPECT_EQ(kUpb_GetExtension_Ok, promote_status); + EXPECT_EQ(9, upb_test_ModelExtension2_i(ext2)); + + size_t end_len; + upb_Message_GetUnknown(base_msg, &end_len); + EXPECT_LT(end_len, start_len); + EXPECT_EQ(6, upb_Message_ExtensionCount(base_msg)); + + upb_Arena_Free(arena); +} + +// Create a minitable to mimic ModelWithSubMessages with unlinked subs +// to lazily promote unknowns after parsing. +upb_MiniTable* CreateMiniTableWithEmptySubTables(upb_Arena* arena) { + upb::MtDataEncoder e; + e.StartMessage(0); + e.PutField(kUpb_FieldType_Int32, 4, 0); + e.PutField(kUpb_FieldType_Message, 5, 0); + e.PutField(kUpb_FieldType_Message, 6, kUpb_FieldModifier_IsRepeated); + + upb_Status status; + upb_Status_Clear(&status); + upb_MiniTable* table = + upb_MiniTable_Build(e.data().data(), e.data().size(), arena, &status); + EXPECT_EQ(status.ok, true); + // Initialize sub table to null. Not using upb_MiniTable_SetSubMessage + // since it checks ->ext on parameter. + upb_MiniTableSub* sub = const_cast( + &table->subs[table->fields[1].UPB_PRIVATE(submsg_index)]); + sub->submsg = nullptr; + sub = const_cast( + &table->subs[table->fields[2].UPB_PRIVATE(submsg_index)]); + sub->submsg = nullptr; + return table; +} + +// Create a minitable to mimic ModelWithMaps with unlinked subs +// to lazily promote unknowns after parsing. +upb_MiniTable* CreateMiniTableWithEmptySubTablesForMaps(upb_Arena* arena) { + upb::MtDataEncoder e; + e.StartMessage(0); + e.PutField(kUpb_FieldType_Int32, 1, 0); + e.PutField(kUpb_FieldType_Message, 3, kUpb_FieldModifier_IsRepeated); + e.PutField(kUpb_FieldType_Message, 4, kUpb_FieldModifier_IsRepeated); + + upb_Status status; + upb_Status_Clear(&status); + upb_MiniTable* table = + upb_MiniTable_Build(e.data().data(), e.data().size(), arena, &status); + EXPECT_EQ(status.ok, true); + // Initialize sub table to null. Not using upb_MiniTable_SetSubMessage + // since it checks ->ext on parameter. + upb_MiniTableSub* sub = const_cast( + &table->subs[table->fields[1].UPB_PRIVATE(submsg_index)]); + sub->submsg = nullptr; + sub = const_cast( + &table->subs[table->fields[2].UPB_PRIVATE(submsg_index)]); + sub->submsg = nullptr; + return table; +} + +upb_MiniTable* CreateMapEntryMiniTable(upb_Arena* arena) { + upb::MtDataEncoder e; + e.EncodeMap(kUpb_FieldType_String, kUpb_FieldType_String, 0, 0); + upb_Status status; + upb_Status_Clear(&status); + upb_MiniTable* table = + upb_MiniTable_Build(e.data().data(), e.data().size(), arena, &status); + EXPECT_EQ(status.ok, true); + return table; +} + +TEST(GeneratedCode, PromoteUnknownMessage) { + upb_Arena* arena = upb_Arena_New(); + upb_test_ModelWithSubMessages* input_msg = + upb_test_ModelWithSubMessages_new(arena); + upb_test_ModelWithExtensions* sub_message = + upb_test_ModelWithExtensions_new(arena); + upb_test_ModelWithSubMessages_set_id(input_msg, 11); + upb_test_ModelWithExtensions_set_random_int32(sub_message, 12); + upb_test_ModelWithSubMessages_set_optional_child(input_msg, sub_message); + size_t serialized_size; + char* serialized = upb_test_ModelWithSubMessages_serialize(input_msg, arena, + &serialized_size); + + upb_MiniTable* mini_table = CreateMiniTableWithEmptySubTables(arena); + upb_Message* msg = _upb_Message_New(mini_table, arena); + upb_DecodeStatus decode_status = upb_Decode(serialized, serialized_size, msg, + mini_table, nullptr, 0, arena); + EXPECT_EQ(decode_status, kUpb_DecodeStatus_Ok); + int32_t val = upb_Message_GetInt32( + msg, upb_MiniTable_FindFieldByNumber(mini_table, 4), 0); + EXPECT_EQ(val, 11); + upb_FindUnknownRet unknown = + upb_MiniTable_FindUnknown(msg, 5, kUpb_WireFormat_DefaultDepthLimit); + EXPECT_EQ(unknown.status, kUpb_FindUnknown_Ok); + // Update mini table and promote unknown to a message. + EXPECT_TRUE(upb_MiniTable_SetSubMessage( + mini_table, (upb_MiniTableField*)&mini_table->fields[1], + &upb_test_ModelWithExtensions_msg_init)); + const int decode_options = upb_DecodeOptions_MaxDepth( + kUpb_WireFormat_DefaultDepthLimit); // UPB_DECODE_ALIAS disabled. + upb_UnknownToMessageRet promote_result = + upb_MiniTable_PromoteUnknownToMessage( + msg, mini_table, &mini_table->fields[1], + &upb_test_ModelWithExtensions_msg_init, decode_options, arena); + EXPECT_EQ(promote_result.status, kUpb_UnknownToMessage_Ok); + const upb_Message* promoted_message = + upb_Message_GetMessage(msg, &mini_table->fields[1], nullptr); + EXPECT_EQ(upb_test_ModelWithExtensions_random_int32( + (upb_test_ModelWithExtensions*)promoted_message), + 12); + upb_Arena_Free(arena); +} + +TEST(GeneratedCode, PromoteUnknownRepeatedMessage) { + upb_Arena* arena = upb_Arena_New(); + upb_test_ModelWithSubMessages* input_msg = + upb_test_ModelWithSubMessages_new(arena); + upb_test_ModelWithSubMessages_set_id(input_msg, 123); + + // Add 2 repeated messages to input_msg. + upb_test_ModelWithExtensions* item = + upb_test_ModelWithSubMessages_add_items(input_msg, arena); + upb_test_ModelWithExtensions_set_random_int32(item, 5); + item = upb_test_ModelWithSubMessages_add_items(input_msg, arena); + upb_test_ModelWithExtensions_set_random_int32(item, 6); + + size_t serialized_size; + char* serialized = upb_test_ModelWithSubMessages_serialize(input_msg, arena, + &serialized_size); + + upb_MiniTable* mini_table = CreateMiniTableWithEmptySubTables(arena); + upb_Message* msg = _upb_Message_New(mini_table, arena); + upb_DecodeStatus decode_status = upb_Decode(serialized, serialized_size, msg, + mini_table, nullptr, 0, arena); + EXPECT_EQ(decode_status, kUpb_DecodeStatus_Ok); + int32_t val = upb_Message_GetInt32( + msg, upb_MiniTable_FindFieldByNumber(mini_table, 4), 0); + EXPECT_EQ(val, 123); + + // Check that we have repeated field data in an unknown. + upb_FindUnknownRet unknown = + upb_MiniTable_FindUnknown(msg, 6, kUpb_WireFormat_DefaultDepthLimit); + EXPECT_EQ(unknown.status, kUpb_FindUnknown_Ok); + + // Update mini table and promote unknown to a message. + EXPECT_TRUE(upb_MiniTable_SetSubMessage( + mini_table, (upb_MiniTableField*)&mini_table->fields[2], + &upb_test_ModelWithExtensions_msg_init)); + const int decode_options = upb_DecodeOptions_MaxDepth( + kUpb_WireFormat_DefaultDepthLimit); // UPB_DECODE_ALIAS disabled. + upb_UnknownToMessage_Status promote_result = + upb_MiniTable_PromoteUnknownToMessageArray( + msg, &mini_table->fields[2], &upb_test_ModelWithExtensions_msg_init, + decode_options, arena); + EXPECT_EQ(promote_result, kUpb_UnknownToMessage_Ok); + + upb_Array* array = upb_Message_GetMutableArray(msg, &mini_table->fields[2]); + const upb_Message* promoted_message = upb_Array_Get(array, 0).msg_val; + EXPECT_EQ(upb_test_ModelWithExtensions_random_int32( + (upb_test_ModelWithExtensions*)promoted_message), + 5); + promoted_message = upb_Array_Get(array, 1).msg_val; + EXPECT_EQ(upb_test_ModelWithExtensions_random_int32( + (upb_test_ModelWithExtensions*)promoted_message), + 6); + upb_Arena_Free(arena); +} + +TEST(GeneratedCode, PromoteUnknownToMap) { + upb_Arena* arena = upb_Arena_New(); + upb_test_ModelWithMaps* input_msg = upb_test_ModelWithMaps_new(arena); + upb_test_ModelWithMaps_set_id(input_msg, 123); + + // Add 2 map entries. + upb_test_ModelWithMaps_map_ss_set(input_msg, + upb_StringView_FromString("key1"), + upb_StringView_FromString("value1"), arena); + upb_test_ModelWithMaps_map_ss_set(input_msg, + upb_StringView_FromString("key2"), + upb_StringView_FromString("value2"), arena); + + size_t serialized_size; + char* serialized = + upb_test_ModelWithMaps_serialize(input_msg, arena, &serialized_size); + + upb_MiniTable* mini_table = CreateMiniTableWithEmptySubTablesForMaps(arena); + upb_MiniTable* map_entry_mini_table = CreateMapEntryMiniTable(arena); + upb_Message* msg = _upb_Message_New(mini_table, arena); + const int decode_options = + upb_DecodeOptions_MaxDepth(kUpb_WireFormat_DefaultDepthLimit); + upb_DecodeStatus decode_status = + upb_Decode(serialized, serialized_size, msg, mini_table, nullptr, + decode_options, arena); + EXPECT_EQ(decode_status, kUpb_DecodeStatus_Ok); + int32_t val = upb_Message_GetInt32( + msg, upb_MiniTable_FindFieldByNumber(mini_table, 1), 0); + EXPECT_EQ(val, 123); + + // Check that we have map data in an unknown. + upb_FindUnknownRet unknown = + upb_MiniTable_FindUnknown(msg, 3, kUpb_WireFormat_DefaultDepthLimit); + EXPECT_EQ(unknown.status, kUpb_FindUnknown_Ok); + + // Update mini table and promote unknown to a message. + EXPECT_TRUE(upb_MiniTable_SetSubMessage( + mini_table, (upb_MiniTableField*)&mini_table->fields[1], + map_entry_mini_table)); + upb_UnknownToMessage_Status promote_result = + upb_MiniTable_PromoteUnknownToMap(msg, mini_table, &mini_table->fields[1], + decode_options, arena); + EXPECT_EQ(promote_result, kUpb_UnknownToMessage_Ok); + + upb_Map* map = upb_Message_GetOrCreateMutableMap( + msg, map_entry_mini_table, &mini_table->fields[1], arena); + EXPECT_NE(map, nullptr); + // Lookup in map. + upb_MessageValue key; + key.str_val = upb_StringView_FromString("key2"); + upb_MessageValue value; + EXPECT_TRUE(upb_Map_Get(map, key, &value)); + EXPECT_EQ(0, strncmp(value.str_val.data, "value2", 5)); + upb_Arena_Free(arena); +} + +} // namespace