upb: fix gencode bug with Map<*, bytes>

PiperOrigin-RevId: 597678769
protocolbuffers · Jan 12, 2024 · e6d8669 · e6d8669
1 parent 7bff169
commit e6d8669
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 5 deletions.
diff --git a/upb/test/test.proto b/upb/test/test.proto
@@ -92,6 +92,7 @@ message ModelWithSubMessages {
 
 message ModelWithMaps {
   optional int32 id = 1;
+  map<string, bytes> map_sb = 2;
   map<string, string> map_ss = 3;
   map<int32, int32> map_ii = 4;
   map<int32, ModelWithExtensions> map_im = 5;

diff --git a/upb/test/test_generated_code.cc b/upb/test/test_generated_code.cc
@@ -20,6 +20,7 @@
 #include "upb/base/string_view.h"
 #include "upb/mem/arena.hpp"
 #include "upb/message/array.h"
+#include "upb/message/map.h"
 #include "upb/test/test.upb.h"
 
 // Must be last.
@@ -923,3 +924,22 @@ TEST(GeneratedCode, Extensions) {
   ASSERT_EQ(size1, size2);
   ASSERT_EQ(0, memcmp(pb1, pb2, size1));
 }
+
+TEST(GeneratedCode, Maps) {
+  upb::Arena arena;
+  upb_test_ModelWithMaps* msg = upb_test_ModelWithMaps_new(arena.ptr());
+
+  auto sb = _upb_test_ModelWithMaps_map_sb_mutable_upb_map(msg, arena.ptr());
+  auto ss = _upb_test_ModelWithMaps_map_ss_mutable_upb_map(msg, arena.ptr());
+  auto ii = _upb_test_ModelWithMaps_map_ii_mutable_upb_map(msg, arena.ptr());
+
+  ASSERT_NE(sb, nullptr);
+  ASSERT_NE(ss, nullptr);
+  ASSERT_NE(ii, nullptr);
+
+  upb_MessageValue key, val;
+  key.str_val = test_str_view;
+  val.str_val = test_str_view2;
+
+  upb_Map_Set(sb, key, val, arena.ptr());
+}
diff --git a/upb_generator/protoc-gen-upb.cc b/upb_generator/protoc-gen-upb.cc
@@ -213,16 +213,20 @@ std::string MapValueCType(upb::FieldDefPtr map_field) {
   return CType(map_field.message_type().map_value());
 }
 
-std::string MapKeySize(upb::FieldDefPtr map_field, absl::string_view expr) {
-  return map_field.message_type().map_key().ctype() == kUpb_CType_String
+std::string MapKeyValueSize(upb_CType ctype, absl::string_view expr) {
+  return ctype == kUpb_CType_String || ctype == kUpb_CType_Bytes
              ? "0"
              : absl::StrCat("sizeof(", expr, ")");
 }
 
+std::string MapKeySize(upb::FieldDefPtr map_field, absl::string_view expr) {
+  const upb_CType ctype = map_field.message_type().map_key().ctype();
+  return MapKeyValueSize(ctype, expr);
+}
+
 std::string MapValueSize(upb::FieldDefPtr map_field, absl::string_view expr) {
-  return map_field.message_type().map_value().ctype() == kUpb_CType_String
-             ? "0"
-             : absl::StrCat("sizeof(", expr, ")");
+  const upb_CType ctype = map_field.message_type().map_value().ctype();
+  return MapKeyValueSize(ctype, expr);
 }
 
 std::string FieldInitializer(const DefPoolPair& pools, upb::FieldDefPtr field,