xref: /aosp_15_r20/external/grpc-grpc/third_party/upb/upb/message/promote_test.cc (revision cc02d7e222339f7a4f6ba5f422e6413f4bd931f2)
1 // Protocol Buffers - Google's data interchange format
2 // Copyright 2023 Google LLC.  All rights reserved.
3 //
4 // Use of this source code is governed by a BSD-style
5 // license that can be found in the LICENSE file or at
6 // https://developers.google.com/open-source/licenses/bsd
7 
8 /* Test of mini table accessors.
9  *
10  * Messages are created and mutated using generated code, and then
11  * accessed through reflective APIs exposed through mini table accessors.
12  */
13 
14 #include "upb/message/promote.h"
15 
16 #include <string.h>
17 
18 #include <cstddef>
19 #include <cstdint>
20 #include <string>
21 
22 #include <gtest/gtest.h>
23 #include "upb/base/descriptor_constants.h"
24 #include "upb/base/status.h"
25 #include "upb/base/string_view.h"
26 #include "upb/base/upcast.h"
27 #include "upb/mem/arena.h"
28 #include "upb/mem/arena.hpp"
29 #include "upb/message/accessors.h"
30 #include "upb/message/array.h"
31 #include "upb/message/copy.h"
32 #include "upb/message/internal/extension.h"
33 #include "upb/message/internal/message.h"
34 #include "upb/message/map.h"
35 #include "upb/message/message.h"
36 #include "upb/message/tagged_ptr.h"
37 #include "upb/mini_descriptor/decode.h"
38 #include "upb/mini_descriptor/internal/encode.hpp"
39 #include "upb/mini_descriptor/internal/modifiers.h"
40 #include "upb/mini_descriptor/link.h"
41 #include "upb/mini_table/extension.h"
42 #include "upb/mini_table/field.h"
43 #include "upb/mini_table/message.h"
44 #include "upb/test/test.upb.h"
45 #include "upb/test/test.upb_minitable.h"
46 #include "upb/wire/decode.h"
47 #include "upb/wire/encode.h"
48 
49 namespace {
50 
TEST(GeneratedCode,FindUnknown)51 TEST(GeneratedCode, FindUnknown) {
52   upb_Arena* arena = upb_Arena_New();
53   upb_test_ModelWithExtensions* msg = upb_test_ModelWithExtensions_new(arena);
54   upb_test_ModelWithExtensions_set_random_int32(msg, 10);
55   upb_test_ModelWithExtensions_set_random_name(
56       msg, upb_StringView_FromString("Hello"));
57 
58   upb_test_ModelExtension1* extension1 = upb_test_ModelExtension1_new(arena);
59   upb_test_ModelExtension1_set_str(extension1,
60                                    upb_StringView_FromString("World"));
61 
62   upb_test_ModelExtension1_set_model_ext(msg, extension1, arena);
63 
64   size_t serialized_size;
65   char* serialized =
66       upb_test_ModelWithExtensions_serialize(msg, arena, &serialized_size);
67 
68   upb_test_EmptyMessageWithExtensions* base_msg =
69       upb_test_EmptyMessageWithExtensions_parse(serialized, serialized_size,
70                                                 arena);
71 
72   upb_FindUnknownRet result = upb_Message_FindUnknown(
73       UPB_UPCAST(base_msg),
74       upb_MiniTableExtension_Number(&upb_test_ModelExtension1_model_ext_ext),
75       0);
76   EXPECT_EQ(kUpb_FindUnknown_Ok, result.status);
77 
78   result = upb_Message_FindUnknown(
79       UPB_UPCAST(base_msg),
80       upb_MiniTableExtension_Number(&upb_test_ModelExtension2_model_ext_ext),
81       0);
82   EXPECT_EQ(kUpb_FindUnknown_NotPresent, result.status);
83 
84   upb_Arena_Free(arena);
85 }
86 
TEST(GeneratedCode,Extensions)87 TEST(GeneratedCode, Extensions) {
88   upb_Arena* arena = upb_Arena_New();
89   upb_test_ModelWithExtensions* msg = upb_test_ModelWithExtensions_new(arena);
90   upb_test_ModelWithExtensions_set_random_int32(msg, 10);
91   upb_test_ModelWithExtensions_set_random_name(
92       msg, upb_StringView_FromString("Hello"));
93 
94   upb_test_ModelExtension1* extension1 = upb_test_ModelExtension1_new(arena);
95   upb_test_ModelExtension1_set_str(extension1,
96                                    upb_StringView_FromString("World"));
97 
98   upb_test_ModelExtension2* extension2 = upb_test_ModelExtension2_new(arena);
99   upb_test_ModelExtension2_set_i(extension2, 5);
100 
101   upb_test_ModelExtension2* extension3 = upb_test_ModelExtension2_new(arena);
102   upb_test_ModelExtension2_set_i(extension3, 6);
103 
104   upb_test_ModelExtension2* extension4 = upb_test_ModelExtension2_new(arena);
105   upb_test_ModelExtension2_set_i(extension4, 7);
106 
107   upb_test_ModelExtension2* extension5 = upb_test_ModelExtension2_new(arena);
108   upb_test_ModelExtension2_set_i(extension5, 8);
109 
110   upb_test_ModelExtension2* extension6 = upb_test_ModelExtension2_new(arena);
111   upb_test_ModelExtension2_set_i(extension6, 9);
112 
113   // Set many extensions, to exercise code paths that involve reallocating the
114   // extensions and unknown fields array.
115   upb_test_ModelExtension1_set_model_ext(msg, extension1, arena);
116   upb_test_ModelExtension2_set_model_ext(msg, extension2, arena);
117   upb_test_ModelExtension2_set_model_ext_2(msg, extension3, arena);
118   upb_test_ModelExtension2_set_model_ext_3(msg, extension4, arena);
119   upb_test_ModelExtension2_set_model_ext_4(msg, extension5, arena);
120   upb_test_ModelExtension2_set_model_ext_5(msg, extension6, arena);
121 
122   size_t serialized_size;
123   char* serialized =
124       upb_test_ModelWithExtensions_serialize(msg, arena, &serialized_size);
125 
126   upb_test_ModelExtension1* ext1;
127   upb_test_ModelExtension2* ext2;
128   upb_GetExtension_Status promote_status;
129   upb_MessageValue value;
130 
131   // Test known GetExtension 1
132   promote_status = upb_Message_GetOrPromoteExtension(
133       UPB_UPCAST(msg), &upb_test_ModelExtension1_model_ext_ext, 0, arena,
134       &value);
135   ext1 = (upb_test_ModelExtension1*)value.msg_val;
136   EXPECT_EQ(kUpb_GetExtension_Ok, promote_status);
137   EXPECT_TRUE(upb_StringView_IsEqual(upb_StringView_FromString("World"),
138                                      upb_test_ModelExtension1_str(ext1)));
139 
140   // Test known GetExtension 2
141   promote_status = upb_Message_GetOrPromoteExtension(
142       UPB_UPCAST(msg), &upb_test_ModelExtension2_model_ext_ext, 0, arena,
143       &value);
144   ext2 = (upb_test_ModelExtension2*)value.msg_val;
145   EXPECT_EQ(kUpb_GetExtension_Ok, promote_status);
146   EXPECT_EQ(5, upb_test_ModelExtension2_i(ext2));
147 
148   // Test known GetExtension 3
149   promote_status = upb_Message_GetOrPromoteExtension(
150       UPB_UPCAST(msg), &upb_test_ModelExtension2_model_ext_2_ext, 0, arena,
151       &value);
152   ext2 = (upb_test_ModelExtension2*)value.msg_val;
153   EXPECT_EQ(kUpb_GetExtension_Ok, promote_status);
154   EXPECT_EQ(6, upb_test_ModelExtension2_i(ext2));
155 
156   // Test known GetExtension 4
157   promote_status = upb_Message_GetOrPromoteExtension(
158       UPB_UPCAST(msg), &upb_test_ModelExtension2_model_ext_3_ext, 0, arena,
159       &value);
160   ext2 = (upb_test_ModelExtension2*)value.msg_val;
161   EXPECT_EQ(kUpb_GetExtension_Ok, promote_status);
162   EXPECT_EQ(7, upb_test_ModelExtension2_i(ext2));
163 
164   // Test known GetExtension 5
165   promote_status = upb_Message_GetOrPromoteExtension(
166       UPB_UPCAST(msg), &upb_test_ModelExtension2_model_ext_4_ext, 0, arena,
167       &value);
168   ext2 = (upb_test_ModelExtension2*)value.msg_val;
169   EXPECT_EQ(kUpb_GetExtension_Ok, promote_status);
170   EXPECT_EQ(8, upb_test_ModelExtension2_i(ext2));
171 
172   // Test known GetExtension 6
173   promote_status = upb_Message_GetOrPromoteExtension(
174       UPB_UPCAST(msg), &upb_test_ModelExtension2_model_ext_5_ext, 0, arena,
175       &value);
176   ext2 = (upb_test_ModelExtension2*)value.msg_val;
177   EXPECT_EQ(kUpb_GetExtension_Ok, promote_status);
178   EXPECT_EQ(9, upb_test_ModelExtension2_i(ext2));
179 
180   upb_test_EmptyMessageWithExtensions* base_msg =
181       upb_test_EmptyMessageWithExtensions_parse(serialized, serialized_size,
182                                                 arena);
183 
184   // Get unknown extension bytes before promotion.
185   size_t start_len;
186   upb_Message_GetUnknown(UPB_UPCAST(base_msg), &start_len);
187   EXPECT_GT(start_len, 0);
188   EXPECT_EQ(0, upb_Message_ExtensionCount(UPB_UPCAST(base_msg)));
189 
190   // Test unknown GetExtension.
191   promote_status = upb_Message_GetOrPromoteExtension(
192       UPB_UPCAST(base_msg), &upb_test_ModelExtension1_model_ext_ext, 0, arena,
193       &value);
194   ext1 = (upb_test_ModelExtension1*)value.msg_val;
195   EXPECT_EQ(kUpb_GetExtension_Ok, promote_status);
196   EXPECT_TRUE(upb_StringView_IsEqual(upb_StringView_FromString("World"),
197                                      upb_test_ModelExtension1_str(ext1)));
198 
199   // Test unknown GetExtension.
200   promote_status = upb_Message_GetOrPromoteExtension(
201       UPB_UPCAST(base_msg), &upb_test_ModelExtension2_model_ext_ext, 0, arena,
202       &value);
203   ext2 = (upb_test_ModelExtension2*)value.msg_val;
204   EXPECT_EQ(kUpb_GetExtension_Ok, promote_status);
205   EXPECT_EQ(5, upb_test_ModelExtension2_i(ext2));
206 
207   // Test unknown GetExtension.
208   promote_status = upb_Message_GetOrPromoteExtension(
209       UPB_UPCAST(base_msg), &upb_test_ModelExtension2_model_ext_2_ext, 0, arena,
210       &value);
211   ext2 = (upb_test_ModelExtension2*)value.msg_val;
212   EXPECT_EQ(kUpb_GetExtension_Ok, promote_status);
213   EXPECT_EQ(6, upb_test_ModelExtension2_i(ext2));
214 
215   // Test unknown GetExtension.
216   promote_status = upb_Message_GetOrPromoteExtension(
217       UPB_UPCAST(base_msg), &upb_test_ModelExtension2_model_ext_3_ext, 0, arena,
218       &value);
219   ext2 = (upb_test_ModelExtension2*)value.msg_val;
220   EXPECT_EQ(kUpb_GetExtension_Ok, promote_status);
221   EXPECT_EQ(7, upb_test_ModelExtension2_i(ext2));
222 
223   // Test unknown GetExtension.
224   promote_status = upb_Message_GetOrPromoteExtension(
225       UPB_UPCAST(base_msg), &upb_test_ModelExtension2_model_ext_4_ext, 0, arena,
226       &value);
227   ext2 = (upb_test_ModelExtension2*)value.msg_val;
228   EXPECT_EQ(kUpb_GetExtension_Ok, promote_status);
229   EXPECT_EQ(8, upb_test_ModelExtension2_i(ext2));
230 
231   // Test unknown GetExtension.
232   promote_status = upb_Message_GetOrPromoteExtension(
233       UPB_UPCAST(base_msg), &upb_test_ModelExtension2_model_ext_5_ext, 0, arena,
234       &value);
235   ext2 = (upb_test_ModelExtension2*)value.msg_val;
236   EXPECT_EQ(kUpb_GetExtension_Ok, promote_status);
237   EXPECT_EQ(9, upb_test_ModelExtension2_i(ext2));
238 
239   size_t end_len;
240   upb_Message_GetUnknown(UPB_UPCAST(base_msg), &end_len);
241   EXPECT_LT(end_len, start_len);
242   EXPECT_EQ(6, upb_Message_ExtensionCount(UPB_UPCAST(base_msg)));
243 
244   upb_Arena_Free(arena);
245 }
246 
247 // Create a minitable to mimic ModelWithSubMessages with unlinked subs
248 // to lazily promote unknowns after parsing.
CreateMiniTableWithEmptySubTables(upb_Arena * arena)249 upb_MiniTable* CreateMiniTableWithEmptySubTables(upb_Arena* arena) {
250   upb::MtDataEncoder e;
251   e.StartMessage(0);
252   e.PutField(kUpb_FieldType_Int32, 4, 0);
253   e.PutField(kUpb_FieldType_Message, 5, 0);
254   e.PutField(kUpb_FieldType_Message, 6, kUpb_FieldModifier_IsRepeated);
255 
256   upb_Status status;
257   upb_Status_Clear(&status);
258   upb_MiniTable* table =
259       upb_MiniTable_Build(e.data().data(), e.data().size(), arena, &status);
260   EXPECT_EQ(status.ok, true);
261   return table;
262 }
263 
CreateMapEntryMiniTable(upb_Arena * arena)264 upb_MiniTable* CreateMapEntryMiniTable(upb_Arena* arena) {
265   upb::MtDataEncoder e;
266   e.EncodeMap(kUpb_FieldType_Int32, kUpb_FieldType_Message, 0, 0);
267   upb_Status status;
268   upb_Status_Clear(&status);
269   upb_MiniTable* table =
270       upb_MiniTable_Build(e.data().data(), e.data().size(), arena, &status);
271   EXPECT_EQ(status.ok, true);
272   return table;
273 }
274 
275 // Create a minitable to mimic ModelWithMaps with unlinked subs
276 // to lazily promote unknowns after parsing.
CreateMiniTableWithEmptySubTablesForMaps(upb_Arena * arena)277 upb_MiniTable* CreateMiniTableWithEmptySubTablesForMaps(upb_Arena* arena) {
278   upb::MtDataEncoder e;
279   e.StartMessage(0);
280   e.PutField(kUpb_FieldType_Int32, 1, 0);
281   e.PutField(kUpb_FieldType_Message, 3, kUpb_FieldModifier_IsRepeated);
282   e.PutField(kUpb_FieldType_Message, 5, kUpb_FieldModifier_IsRepeated);
283 
284   upb_Status status;
285   upb_Status_Clear(&status);
286   upb_MiniTable* table =
287       upb_MiniTable_Build(e.data().data(), e.data().size(), arena, &status);
288 
289   // Field 5 corresponds to ModelWithMaps.map_sm.
290   upb_MiniTableField* map_field = const_cast<upb_MiniTableField*>(
291       upb_MiniTable_FindFieldByNumber(table, 5));
292   EXPECT_NE(map_field, nullptr);
293   upb_MiniTable* sub_table = CreateMapEntryMiniTable(arena);
294   upb_MiniTable_SetSubMessage(table, map_field, sub_table);
295   EXPECT_EQ(status.ok, true);
296   return table;
297 }
298 
CheckReserialize(const upb_Message * msg,const upb_MiniTable * mini_table,upb_Arena * arena,char * serialized,size_t serialized_size)299 void CheckReserialize(const upb_Message* msg, const upb_MiniTable* mini_table,
300                       upb_Arena* arena, char* serialized,
301                       size_t serialized_size) {
302   // We can safely encode the "empty" message.  We expect to get the same bytes
303   // out as were parsed.
304   size_t reserialized_size;
305   char* reserialized;
306   upb_EncodeStatus encode_status =
307       upb_Encode(msg, mini_table, kUpb_EncodeOption_Deterministic, arena,
308                  &reserialized, &reserialized_size);
309   EXPECT_EQ(encode_status, kUpb_EncodeStatus_Ok);
310   EXPECT_EQ(reserialized_size, serialized_size);
311   EXPECT_EQ(0, memcmp(reserialized, serialized, serialized_size));
312 
313   // We should get the same result if we copy+reserialize.
314   upb_Message* clone = upb_Message_DeepClone(msg, mini_table, arena);
315   encode_status = upb_Encode(clone, mini_table, kUpb_EncodeOption_Deterministic,
316                              arena, &reserialized, &reserialized_size);
317   EXPECT_EQ(encode_status, kUpb_EncodeStatus_Ok);
318   EXPECT_EQ(reserialized_size, serialized_size);
319   EXPECT_EQ(0, memcmp(reserialized, serialized, serialized_size));
320 }
321 
TEST(GeneratedCode,PromoteUnknownMessage)322 TEST(GeneratedCode, PromoteUnknownMessage) {
323   upb::Arena arena;
324   upb_test_ModelWithSubMessages* input_msg =
325       upb_test_ModelWithSubMessages_new(arena.ptr());
326   upb_test_ModelWithExtensions* sub_message =
327       upb_test_ModelWithExtensions_new(arena.ptr());
328   upb_test_ModelWithSubMessages_set_id(input_msg, 11);
329   upb_test_ModelWithExtensions_set_random_int32(sub_message, 12);
330   upb_test_ModelWithSubMessages_set_optional_child(input_msg, sub_message);
331   size_t serialized_size;
332   char* serialized = upb_test_ModelWithSubMessages_serialize(
333       input_msg, arena.ptr(), &serialized_size);
334 
335   upb_MiniTable* mini_table = CreateMiniTableWithEmptySubTables(arena.ptr());
336   upb_DecodeStatus decode_status;
337 
338   // If we parse without allowing unlinked objects, the parse will fail.
339   // TODO: re-enable this test once the old method of tree shaking is
340   // removed
341   // upb_Message* fail_msg = _upb_Message_New(mini_table, arena.ptr());
342   // decode_status =
343   //     upb_Decode(serialized, serialized_size, fail_msg, mini_table, nullptr,
344   //     0,
345   //                arena.ptr());
346   // EXPECT_EQ(decode_status, kUpb_DecodeStatus_UnlinkedSubMessage);
347 
348   // if we parse while allowing unlinked objects, the parse will succeed.
349   upb_Message* msg = _upb_Message_New(mini_table, arena.ptr());
350   decode_status =
351       upb_Decode(serialized, serialized_size, msg, mini_table, nullptr,
352                  kUpb_DecodeOption_ExperimentalAllowUnlinked, arena.ptr());
353   EXPECT_EQ(decode_status, kUpb_DecodeStatus_Ok);
354 
355   CheckReserialize(msg, mini_table, arena.ptr(), serialized, serialized_size);
356 
357   // We can encode the "empty" message and get the same output bytes.
358   size_t reserialized_size;
359   char* reserialized;
360   upb_EncodeStatus encode_status = upb_Encode(
361       msg, mini_table, 0, arena.ptr(), &reserialized, &reserialized_size);
362   EXPECT_EQ(encode_status, kUpb_EncodeStatus_Ok);
363   EXPECT_EQ(reserialized_size, serialized_size);
364   EXPECT_EQ(0, memcmp(reserialized, serialized, serialized_size));
365 
366   // Int32 field is present, as normal.
367   int32_t val = upb_Message_GetInt32(
368       msg, upb_MiniTable_FindFieldByNumber(mini_table, 4), 0);
369   EXPECT_EQ(val, 11);
370 
371   // Unlinked sub-message is present, but getting the value returns NULL.
372   const upb_MiniTableField* submsg_field =
373       upb_MiniTable_FindFieldByNumber(mini_table, 5);
374   ASSERT_TRUE(submsg_field != nullptr);
375   EXPECT_TRUE(upb_Message_HasBaseField(msg, submsg_field));
376   upb_TaggedMessagePtr tagged =
377       upb_Message_GetTaggedMessagePtr(msg, submsg_field, nullptr);
378   EXPECT_TRUE(upb_TaggedMessagePtr_IsEmpty(tagged));
379 
380   // Update mini table and promote unknown to a message.
381   EXPECT_TRUE(
382       upb_MiniTable_SetSubMessage(mini_table, (upb_MiniTableField*)submsg_field,
383                                   &upb_0test__ModelWithExtensions_msg_init));
384 
385   const int decode_options =
386       upb_DecodeOptions_MaxDepth(0);  // UPB_DECODE_ALIAS disabled.
387   upb_test_ModelWithExtensions* promoted;
388   upb_DecodeStatus promote_result =
389       upb_Message_PromoteMessage(msg, mini_table, submsg_field, decode_options,
390                                  arena.ptr(), (upb_Message**)&promoted);
391   EXPECT_EQ(promote_result, kUpb_DecodeStatus_Ok);
392   EXPECT_NE(nullptr, promoted);
393   EXPECT_EQ(UPB_UPCAST(promoted),
394             upb_Message_GetMessage(msg, submsg_field, nullptr));
395   EXPECT_EQ(upb_test_ModelWithExtensions_random_int32(promoted), 12);
396 }
397 
398 // Tests a second parse that reuses an empty/unlinked message while the message
399 // is still unlinked.
TEST(GeneratedCode,ReparseUnlinked)400 TEST(GeneratedCode, ReparseUnlinked) {
401   upb::Arena arena;
402   upb_test_ModelWithSubMessages* input_msg =
403       upb_test_ModelWithSubMessages_new(arena.ptr());
404   upb_test_ModelWithExtensions* sub_message =
405       upb_test_ModelWithExtensions_new(arena.ptr());
406   upb_test_ModelWithSubMessages_set_id(input_msg, 11);
407   upb_test_ModelWithExtensions_add_repeated_int32(sub_message, 12, arena.ptr());
408   upb_test_ModelWithSubMessages_set_optional_child(input_msg, sub_message);
409   size_t serialized_size;
410   char* serialized = upb_test_ModelWithSubMessages_serialize(
411       input_msg, arena.ptr(), &serialized_size);
412 
413   upb_MiniTable* mini_table = CreateMiniTableWithEmptySubTables(arena.ptr());
414 
415   // Parse twice without linking the MiniTable.
416   upb_Message* msg = _upb_Message_New(mini_table, arena.ptr());
417   upb_DecodeStatus decode_status =
418       upb_Decode(serialized, serialized_size, msg, mini_table, nullptr,
419                  kUpb_DecodeOption_ExperimentalAllowUnlinked, arena.ptr());
420   EXPECT_EQ(decode_status, kUpb_DecodeStatus_Ok);
421 
422   decode_status =
423       upb_Decode(serialized, serialized_size, msg, mini_table, nullptr,
424                  kUpb_DecodeOption_ExperimentalAllowUnlinked, arena.ptr());
425   EXPECT_EQ(decode_status, kUpb_DecodeStatus_Ok);
426 
427   // Update mini table and promote unknown to a message.
428   const upb_MiniTableField* submsg_field =
429       upb_MiniTable_FindFieldByNumber(mini_table, 5);
430   EXPECT_TRUE(
431       upb_MiniTable_SetSubMessage(mini_table, (upb_MiniTableField*)submsg_field,
432                                   &upb_0test__ModelWithExtensions_msg_init));
433 
434   const int decode_options =
435       upb_DecodeOptions_MaxDepth(0);  // UPB_DECODE_ALIAS disabled.
436   upb_test_ModelWithExtensions* promoted;
437   upb_DecodeStatus promote_result =
438       upb_Message_PromoteMessage(msg, mini_table, submsg_field, decode_options,
439                                  arena.ptr(), (upb_Message**)&promoted);
440   EXPECT_EQ(promote_result, kUpb_DecodeStatus_Ok);
441   EXPECT_NE(nullptr, promoted);
442   EXPECT_EQ(UPB_UPCAST(promoted),
443             upb_Message_GetMessage(msg, submsg_field, nullptr));
444 
445   // The repeated field should have two entries for the two parses.
446   size_t repeated_size;
447   const int32_t* entries =
448       upb_test_ModelWithExtensions_repeated_int32(promoted, &repeated_size);
449   EXPECT_EQ(repeated_size, 2);
450   EXPECT_EQ(entries[0], 12);
451   EXPECT_EQ(entries[1], 12);
452 }
453 
454 // Tests a second parse that promotes a message within the parser because we are
455 // merging into an empty/unlinked message after the message has been linked.
TEST(GeneratedCode,PromoteInParser)456 TEST(GeneratedCode, PromoteInParser) {
457   upb::Arena arena;
458   upb_test_ModelWithSubMessages* input_msg =
459       upb_test_ModelWithSubMessages_new(arena.ptr());
460   upb_test_ModelWithExtensions* sub_message =
461       upb_test_ModelWithExtensions_new(arena.ptr());
462   upb_test_ModelWithSubMessages_set_id(input_msg, 11);
463   upb_test_ModelWithExtensions_add_repeated_int32(sub_message, 12, arena.ptr());
464   upb_test_ModelWithSubMessages_set_optional_child(input_msg, sub_message);
465   size_t serialized_size;
466   char* serialized = upb_test_ModelWithSubMessages_serialize(
467       input_msg, arena.ptr(), &serialized_size);
468 
469   upb_MiniTable* mini_table = CreateMiniTableWithEmptySubTables(arena.ptr());
470 
471   // Parse once without linking the MiniTable.
472   upb_Message* msg = _upb_Message_New(mini_table, arena.ptr());
473   upb_DecodeStatus decode_status =
474       upb_Decode(serialized, serialized_size, msg, mini_table, nullptr,
475                  kUpb_DecodeOption_ExperimentalAllowUnlinked, arena.ptr());
476   EXPECT_EQ(decode_status, kUpb_DecodeStatus_Ok);
477 
478   // Link the MiniTable.
479   const upb_MiniTableField* submsg_field =
480       upb_MiniTable_FindFieldByNumber(mini_table, 5);
481   EXPECT_TRUE(
482       upb_MiniTable_SetSubMessage(mini_table, (upb_MiniTableField*)submsg_field,
483                                   &upb_0test__ModelWithExtensions_msg_init));
484 
485   // Parse again.  This will promote the message.  An explicit promote will not
486   // be required.
487   decode_status =
488       upb_Decode(serialized, serialized_size, msg, mini_table, nullptr,
489                  kUpb_DecodeOption_ExperimentalAllowUnlinked, arena.ptr());
490   EXPECT_EQ(decode_status, kUpb_DecodeStatus_Ok);
491   upb_test_ModelWithExtensions* promoted =
492       (upb_test_ModelWithExtensions*)upb_Message_GetMessage(msg, submsg_field,
493                                                             nullptr);
494 
495   EXPECT_NE(nullptr, promoted);
496   EXPECT_EQ(UPB_UPCAST(promoted),
497             upb_Message_GetMessage(msg, submsg_field, nullptr));
498 
499   // The repeated field should have two entries for the two parses.
500   size_t repeated_size;
501   const int32_t* entries =
502       upb_test_ModelWithExtensions_repeated_int32(promoted, &repeated_size);
503   EXPECT_EQ(repeated_size, 2);
504   EXPECT_EQ(entries[0], 12);
505   EXPECT_EQ(entries[1], 12);
506 }
507 
TEST(GeneratedCode,PromoteUnknownRepeatedMessage)508 TEST(GeneratedCode, PromoteUnknownRepeatedMessage) {
509   upb::Arena arena;
510   upb_test_ModelWithSubMessages* input_msg =
511       upb_test_ModelWithSubMessages_new(arena.ptr());
512   upb_test_ModelWithSubMessages_set_id(input_msg, 123);
513 
514   // Add 2 repeated messages to input_msg.
515   upb_test_ModelWithExtensions* item =
516       upb_test_ModelWithSubMessages_add_items(input_msg, arena.ptr());
517   upb_test_ModelWithExtensions_set_random_int32(item, 5);
518   item = upb_test_ModelWithSubMessages_add_items(input_msg, arena.ptr());
519   upb_test_ModelWithExtensions_set_random_int32(item, 6);
520 
521   size_t serialized_size;
522   char* serialized = upb_test_ModelWithSubMessages_serialize(
523       input_msg, arena.ptr(), &serialized_size);
524 
525   upb_MiniTable* mini_table = CreateMiniTableWithEmptySubTables(arena.ptr());
526   upb_DecodeStatus decode_status;
527 
528   // If we parse without allowing unlinked objects, the parse will fail.
529   // TODO: re-enable this test once the old method of tree shaking is
530   // removed
531   // upb_Message* fail_msg = _upb_Message_New(mini_table, arena.ptr());
532   // decode_status =
533   //     upb_Decode(serialized, serialized_size, fail_msg, mini_table, nullptr,
534   //     0,
535   //                arena.ptr());
536   // EXPECT_EQ(decode_status, kUpb_DecodeStatus_UnlinkedSubMessage);
537 
538   // if we parse while allowing unlinked objects, the parse will succeed.
539   upb_Message* msg = _upb_Message_New(mini_table, arena.ptr());
540   decode_status =
541       upb_Decode(serialized, serialized_size, msg, mini_table, nullptr,
542                  kUpb_DecodeOption_ExperimentalAllowUnlinked, arena.ptr());
543 
544   CheckReserialize(msg, mini_table, arena.ptr(), serialized, serialized_size);
545 
546   // Int32 field is present, as normal.
547   EXPECT_EQ(decode_status, kUpb_DecodeStatus_Ok);
548   int32_t val = upb_Message_GetInt32(
549       msg, upb_MiniTable_FindFieldByNumber(mini_table, 4), 0);
550   EXPECT_EQ(val, 123);
551 
552   const upb_MiniTableField* repeated_field =
553       upb_MiniTable_FindFieldByNumber(mini_table, 6);
554 
555   upb_Array* array = upb_Message_GetMutableArray(msg, repeated_field);
556 
557   // Array length is 2 even though the messages are empty.
558   EXPECT_EQ(2, upb_Array_Size(array));
559 
560   // Update mini table and promote unknown to a message.
561   EXPECT_TRUE(upb_MiniTable_SetSubMessage(
562       mini_table, (upb_MiniTableField*)repeated_field,
563       &upb_0test__ModelWithExtensions_msg_init));
564   const int decode_options =
565       upb_DecodeOptions_MaxDepth(0);  // UPB_DECODE_ALIAS disabled.
566   upb_DecodeStatus promote_result =
567       upb_Array_PromoteMessages(array, &upb_0test__ModelWithExtensions_msg_init,
568                                 decode_options, arena.ptr());
569   EXPECT_EQ(promote_result, kUpb_DecodeStatus_Ok);
570   const upb_Message* promoted_message = upb_Array_Get(array, 0).msg_val;
571   EXPECT_EQ(upb_test_ModelWithExtensions_random_int32(
572                 (upb_test_ModelWithExtensions*)promoted_message),
573             5);
574   promoted_message = upb_Array_Get(array, 1).msg_val;
575   EXPECT_EQ(upb_test_ModelWithExtensions_random_int32(
576                 (upb_test_ModelWithExtensions*)promoted_message),
577             6);
578 }
579 
TEST(GeneratedCode,PromoteUnknownToMap)580 TEST(GeneratedCode, PromoteUnknownToMap) {
581   upb::Arena arena;
582   upb_test_ModelWithMaps* input_msg = upb_test_ModelWithMaps_new(arena.ptr());
583   upb_test_ModelWithMaps_set_id(input_msg, 123);
584 
585   upb_test_ModelWithExtensions* submsg1 =
586       upb_test_ModelWithExtensions_new(arena.ptr());
587   upb_test_ModelWithExtensions_set_random_int32(submsg1, 123);
588   upb_test_ModelWithExtensions* submsg2 =
589       upb_test_ModelWithExtensions_new(arena.ptr());
590   upb_test_ModelWithExtensions_set_random_int32(submsg2, 456);
591 
592   // Add 2 map entries.
593   upb_test_ModelWithMaps_map_im_set(input_msg, 111, submsg1, arena.ptr());
594   upb_test_ModelWithMaps_map_im_set(input_msg, 222, submsg2, arena.ptr());
595 
596   size_t serialized_size;
597   char* serialized = upb_test_ModelWithMaps_serialize_ex(
598       input_msg, kUpb_EncodeOption_Deterministic, arena.ptr(),
599       &serialized_size);
600 
601   upb_MiniTable* mini_table =
602       CreateMiniTableWithEmptySubTablesForMaps(arena.ptr());
603 
604   // If we parse without allowing unlinked objects, the parse will fail.
605   upb_Message* fail_msg1 = _upb_Message_New(mini_table, arena.ptr());
606   upb_DecodeStatus decode_status =
607       upb_Decode(serialized, serialized_size, fail_msg1, mini_table, nullptr, 0,
608                  arena.ptr());
609   EXPECT_EQ(decode_status, kUpb_DecodeStatus_UnlinkedSubMessage);
610 
611   // if we parse while allowing unlinked objects, the parse will succeed.
612   upb_Message* msg = _upb_Message_New(mini_table, arena.ptr());
613   decode_status =
614       upb_Decode(serialized, serialized_size, msg, mini_table, nullptr,
615                  kUpb_DecodeOption_ExperimentalAllowUnlinked, arena.ptr());
616   EXPECT_EQ(decode_status, kUpb_DecodeStatus_Ok);
617 
618   CheckReserialize(msg, mini_table, arena.ptr(), serialized, serialized_size);
619 
620   upb_MiniTableField* map_field = const_cast<upb_MiniTableField*>(
621       upb_MiniTable_FindFieldByNumber(mini_table, 5));
622 
623   upb_Map* map = upb_Message_GetMutableMap(msg, map_field);
624 
625   // Map size is 2 even though messages are unlinked.
626   EXPECT_EQ(2, upb_Map_Size(map));
627 
628   // Update mini table and promote unknown to a message.
629   upb_MiniTable* entry = const_cast<upb_MiniTable*>(
630       upb_MiniTable_GetSubMessageTable(mini_table, map_field));
631   upb_MiniTableField* entry_value = const_cast<upb_MiniTableField*>(
632       upb_MiniTable_FindFieldByNumber(entry, 2));
633   upb_MiniTable_SetSubMessage(entry, entry_value,
634                               &upb_0test__ModelWithExtensions_msg_init);
635   upb_DecodeStatus promote_result = upb_Map_PromoteMessages(
636       map, &upb_0test__ModelWithExtensions_msg_init, 0, arena.ptr());
637   EXPECT_EQ(promote_result, kUpb_DecodeStatus_Ok);
638 
639   upb_MessageValue key;
640   upb_MessageValue val;
641   key.int32_val = 111;
642   EXPECT_TRUE(upb_Map_Get(map, key, &val));
643   EXPECT_EQ(123, upb_test_ModelWithExtensions_random_int32(
644                      static_cast<const upb_test_ModelWithExtensions*>(
645                          (void*)(val.msg_val))));
646 
647   key.int32_val = 222;
648   EXPECT_TRUE(upb_Map_Get(map, key, &val));
649   EXPECT_EQ(456, upb_test_ModelWithExtensions_random_int32(
650                      static_cast<const upb_test_ModelWithExtensions*>(
651                          (void*)(val.msg_val))));
652 }
653 
654 }  // namespace
655 
656 // OLD tests, to be removed!
657 
658 namespace {
659 
660 // Create a minitable to mimic ModelWithSubMessages with unlinked subs
661 // to lazily promote unknowns after parsing.
CreateMiniTableWithEmptySubTablesOld(upb_Arena * arena)662 upb_MiniTable* CreateMiniTableWithEmptySubTablesOld(upb_Arena* arena) {
663   upb::MtDataEncoder e;
664   e.StartMessage(0);
665   e.PutField(kUpb_FieldType_Int32, 4, 0);
666   e.PutField(kUpb_FieldType_Message, 5, 0);
667   e.PutField(kUpb_FieldType_Message, 6, kUpb_FieldModifier_IsRepeated);
668 
669   upb_Status status;
670   upb_Status_Clear(&status);
671   upb_MiniTable* table =
672       upb_MiniTable_Build(e.data().data(), e.data().size(), arena, &status);
673   EXPECT_EQ(status.ok, true);
674   return table;
675 }
676 
677 // Create a minitable to mimic ModelWithMaps with unlinked subs
678 // to lazily promote unknowns after parsing.
CreateMiniTableWithEmptySubTablesForMapsOld(upb_Arena * arena)679 upb_MiniTable* CreateMiniTableWithEmptySubTablesForMapsOld(upb_Arena* arena) {
680   upb::MtDataEncoder e;
681   e.StartMessage(0);
682   e.PutField(kUpb_FieldType_Int32, 1, 0);
683   e.PutField(kUpb_FieldType_Message, 3, kUpb_FieldModifier_IsRepeated);
684   e.PutField(kUpb_FieldType_Message, 4, kUpb_FieldModifier_IsRepeated);
685 
686   upb_Status status;
687   upb_Status_Clear(&status);
688   upb_MiniTable* table =
689       upb_MiniTable_Build(e.data().data(), e.data().size(), arena, &status);
690   EXPECT_EQ(status.ok, true);
691   return table;
692 }
693 
CreateMapEntryMiniTableOld(upb_Arena * arena)694 upb_MiniTable* CreateMapEntryMiniTableOld(upb_Arena* arena) {
695   upb::MtDataEncoder e;
696   e.EncodeMap(kUpb_FieldType_String, kUpb_FieldType_String, 0, 0);
697   upb_Status status;
698   upb_Status_Clear(&status);
699   upb_MiniTable* table =
700       upb_MiniTable_Build(e.data().data(), e.data().size(), arena, &status);
701   EXPECT_EQ(status.ok, true);
702   return table;
703 }
704 
TEST(GeneratedCode,PromoteUnknownMessageOld)705 TEST(GeneratedCode, PromoteUnknownMessageOld) {
706   upb_Arena* arena = upb_Arena_New();
707   upb_test_ModelWithSubMessages* input_msg =
708       upb_test_ModelWithSubMessages_new(arena);
709   upb_test_ModelWithExtensions* sub_message =
710       upb_test_ModelWithExtensions_new(arena);
711   upb_test_ModelWithSubMessages_set_id(input_msg, 11);
712   upb_test_ModelWithExtensions_set_random_int32(sub_message, 12);
713   upb_test_ModelWithSubMessages_set_optional_child(input_msg, sub_message);
714   size_t serialized_size;
715   char* serialized = upb_test_ModelWithSubMessages_serialize(input_msg, arena,
716                                                              &serialized_size);
717 
718   upb_MiniTable* mini_table = CreateMiniTableWithEmptySubTablesOld(arena);
719   upb_Message* msg = _upb_Message_New(mini_table, arena);
720   upb_DecodeStatus decode_status = upb_Decode(serialized, serialized_size, msg,
721                                               mini_table, nullptr, 0, arena);
722   EXPECT_EQ(decode_status, kUpb_DecodeStatus_Ok);
723   int32_t val = upb_Message_GetInt32(
724       msg, upb_MiniTable_FindFieldByNumber(mini_table, 4), 0);
725   EXPECT_EQ(val, 11);
726   upb_FindUnknownRet unknown = upb_Message_FindUnknown(msg, 5, 0);
727   EXPECT_EQ(unknown.status, kUpb_FindUnknown_Ok);
728   // Update mini table and promote unknown to a message.
729   EXPECT_TRUE(upb_MiniTable_SetSubMessage(
730       mini_table,
731       (upb_MiniTableField*)upb_MiniTable_GetFieldByIndex(mini_table, 1),
732       &upb_0test__ModelWithExtensions_msg_init));
733   const int decode_options =
734       upb_DecodeOptions_MaxDepth(0);  // UPB_DECODE_ALIAS disabled.
735   upb_UnknownToMessageRet promote_result =
736       upb_MiniTable_PromoteUnknownToMessage(
737           msg, mini_table, upb_MiniTable_GetFieldByIndex(mini_table, 1),
738           &upb_0test__ModelWithExtensions_msg_init, decode_options, arena);
739   EXPECT_EQ(promote_result.status, kUpb_UnknownToMessage_Ok);
740   const upb_Message* promoted_message = upb_Message_GetMessage(
741       msg, upb_MiniTable_GetFieldByIndex(mini_table, 1), nullptr);
742   EXPECT_EQ(upb_test_ModelWithExtensions_random_int32(
743                 (upb_test_ModelWithExtensions*)promoted_message),
744             12);
745   upb_Arena_Free(arena);
746 }
747 
TEST(GeneratedCode,PromoteUnknownRepeatedMessageOld)748 TEST(GeneratedCode, PromoteUnknownRepeatedMessageOld) {
749   upb_Arena* arena = upb_Arena_New();
750   upb_test_ModelWithSubMessages* input_msg =
751       upb_test_ModelWithSubMessages_new(arena);
752   upb_test_ModelWithSubMessages_set_id(input_msg, 123);
753 
754   // Add 2 repeated messages to input_msg.
755   upb_test_ModelWithExtensions* item =
756       upb_test_ModelWithSubMessages_add_items(input_msg, arena);
757   upb_test_ModelWithExtensions_set_random_int32(item, 5);
758   item = upb_test_ModelWithSubMessages_add_items(input_msg, arena);
759   upb_test_ModelWithExtensions_set_random_int32(item, 6);
760 
761   size_t serialized_size;
762   char* serialized = upb_test_ModelWithSubMessages_serialize(input_msg, arena,
763                                                              &serialized_size);
764 
765   upb_MiniTable* mini_table = CreateMiniTableWithEmptySubTablesOld(arena);
766   upb_Message* msg = _upb_Message_New(mini_table, arena);
767   upb_DecodeStatus decode_status = upb_Decode(serialized, serialized_size, msg,
768                                               mini_table, nullptr, 0, arena);
769   EXPECT_EQ(decode_status, kUpb_DecodeStatus_Ok);
770   int32_t val = upb_Message_GetInt32(
771       msg, upb_MiniTable_FindFieldByNumber(mini_table, 4), 0);
772   EXPECT_EQ(val, 123);
773 
774   // Check that we have repeated field data in an unknown.
775   upb_FindUnknownRet unknown = upb_Message_FindUnknown(msg, 6, 0);
776   EXPECT_EQ(unknown.status, kUpb_FindUnknown_Ok);
777 
778   // Update mini table and promote unknown to a message.
779   EXPECT_TRUE(upb_MiniTable_SetSubMessage(
780       mini_table,
781       (upb_MiniTableField*)upb_MiniTable_GetFieldByIndex(mini_table, 2),
782       &upb_0test__ModelWithExtensions_msg_init));
783   const int decode_options =
784       upb_DecodeOptions_MaxDepth(0);  // UPB_DECODE_ALIAS disabled.
785   upb_UnknownToMessage_Status promote_result =
786       upb_MiniTable_PromoteUnknownToMessageArray(
787           msg, upb_MiniTable_GetFieldByIndex(mini_table, 2),
788           &upb_0test__ModelWithExtensions_msg_init, decode_options, arena);
789   EXPECT_EQ(promote_result, kUpb_UnknownToMessage_Ok);
790 
791   upb_Array* array = upb_Message_GetMutableArray(
792       msg, upb_MiniTable_GetFieldByIndex(mini_table, 2));
793   const upb_Message* promoted_message = upb_Array_Get(array, 0).msg_val;
794   EXPECT_EQ(upb_test_ModelWithExtensions_random_int32(
795                 (upb_test_ModelWithExtensions*)promoted_message),
796             5);
797   promoted_message = upb_Array_Get(array, 1).msg_val;
798   EXPECT_EQ(upb_test_ModelWithExtensions_random_int32(
799                 (upb_test_ModelWithExtensions*)promoted_message),
800             6);
801   upb_Arena_Free(arena);
802 }
803 
TEST(GeneratedCode,PromoteUnknownToMapOld)804 TEST(GeneratedCode, PromoteUnknownToMapOld) {
805   upb_Arena* arena = upb_Arena_New();
806   upb_test_ModelWithMaps* input_msg = upb_test_ModelWithMaps_new(arena);
807   upb_test_ModelWithMaps_set_id(input_msg, 123);
808 
809   // Add 2 map entries.
810   upb_test_ModelWithMaps_map_ss_set(input_msg,
811                                     upb_StringView_FromString("key1"),
812                                     upb_StringView_FromString("value1"), arena);
813   upb_test_ModelWithMaps_map_ss_set(input_msg,
814                                     upb_StringView_FromString("key2"),
815                                     upb_StringView_FromString("value2"), arena);
816 
817   size_t serialized_size;
818   char* serialized =
819       upb_test_ModelWithMaps_serialize(input_msg, arena, &serialized_size);
820 
821   upb_MiniTable* mini_table =
822       CreateMiniTableWithEmptySubTablesForMapsOld(arena);
823   upb_MiniTable* map_entry_mini_table = CreateMapEntryMiniTableOld(arena);
824   upb_Message* msg = _upb_Message_New(mini_table, arena);
825   const int decode_options = upb_DecodeOptions_MaxDepth(0);
826   upb_DecodeStatus decode_status =
827       upb_Decode(serialized, serialized_size, msg, mini_table, nullptr,
828                  decode_options, arena);
829   EXPECT_EQ(decode_status, kUpb_DecodeStatus_Ok);
830   int32_t val = upb_Message_GetInt32(
831       msg, upb_MiniTable_FindFieldByNumber(mini_table, 1), 0);
832   EXPECT_EQ(val, 123);
833 
834   // Check that we have map data in an unknown.
835   upb_FindUnknownRet unknown = upb_Message_FindUnknown(msg, 3, 0);
836   EXPECT_EQ(unknown.status, kUpb_FindUnknown_Ok);
837 
838   // Update mini table and promote unknown to a message.
839   EXPECT_TRUE(upb_MiniTable_SetSubMessage(
840       mini_table,
841       (upb_MiniTableField*)upb_MiniTable_GetFieldByIndex(mini_table, 1),
842       map_entry_mini_table));
843   upb_UnknownToMessage_Status promote_result =
844       upb_MiniTable_PromoteUnknownToMap(
845           msg, mini_table, upb_MiniTable_GetFieldByIndex(mini_table, 1),
846           decode_options, arena);
847   EXPECT_EQ(promote_result, kUpb_UnknownToMessage_Ok);
848 
849   upb_Map* map = upb_Message_GetOrCreateMutableMap(
850       msg, map_entry_mini_table, upb_MiniTable_GetFieldByIndex(mini_table, 1),
851       arena);
852   EXPECT_NE(map, nullptr);
853   // Lookup in map.
854   upb_MessageValue key;
855   key.str_val = upb_StringView_FromString("key2");
856   upb_MessageValue value;
857   EXPECT_TRUE(upb_Map_Get(map, key, &value));
858   EXPECT_EQ(0, strncmp(value.str_val.data, "value2", 5));
859   upb_Arena_Free(arena);
860 }
861 
862 }  // namespace
863