hpb: Add support for string extensions

We introduce a UpbExtensionTrait specialization for absl::string_view.

PiperOrigin-RevId: 734232249

hpb/extension.h

@@ -14,11 +14,13 @@
 #include "absl/base/attributes.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
 #include "google/protobuf/hpb/backend/upb/interop.h"
 #include "google/protobuf/hpb/internal/message_lock.h"
 #include "google/protobuf/hpb/internal/template_help.h"
 #include "google/protobuf/hpb/ptr.h"
 #include "google/protobuf/hpb/status.h"
+#include "upb/base/string_view.h"
 #include "upb/mem/arena.h"
 #include "upb/mem/arena.hpp"
 #include "upb/message/accessors.h"
@@ -72,7 +74,6 @@ struct UpbExtensionTrait<hpb::RepeatedField<T>> {
     return ReturnType(upb_arr, hpb::interop::upb::GetArena(message));
   }
 };
-
 #define UPB_EXT_PRIMITIVE(CppType, UpbFunc)                              \
   template <>                                                            \
   struct UpbExtensionTrait<CppType> {                                    \
@@ -93,7 +94,7 @@ struct UpbExtensionTrait<hpb::RepeatedField<T>> {
           interop::upb::GetArena(message));                              \
       return res ? absl::OkStatus() : MessageAllocationError();          \
     }                                                                    \
-  };
+  }
 
 UPB_EXT_PRIMITIVE(bool, Bool);
 UPB_EXT_PRIMITIVE(int32_t, Int32);
@@ -105,6 +106,30 @@ UPB_EXT_PRIMITIVE(double, Double);
 
 #undef UPB_EXT_PRIMITIVE
 
+template <>
+struct UpbExtensionTrait<absl::string_view> {
+  using DefaultType = absl::string_view;
+  using ReturnType = absl::string_view;
+
+  template <typename Msg, typename Id>
+  static constexpr ReturnType Get(Msg message, const Id& id) {
+    auto default_val = hpb::internal::PrivateAccess::GetDefaultValue(id);
+    upb_StringView result = upb_Message_GetExtensionString(
+        hpb::interop::upb::GetMessage(message), id.mini_table_ext(),
+        upb_StringView_FromDataAndSize(default_val.data(), default_val.size()));
+    return absl::string_view(result.data, result.size);
+  }
+
+  template <typename Msg, typename Id>
+  static absl::Status Set(Msg message, const Id& id, absl::string_view value) {
+    auto upb_value = upb_StringView_FromDataAndSize(value.data(), value.size());
+    bool res = upb_Message_SetExtensionString(interop::upb::GetMessage(message),
+                                              id.mini_table_ext(), upb_value,
+                                              interop::upb::GetArena(message));
+    return res ? absl::OkStatus() : MessageAllocationError();
+  }
+};
+
 // TODO: b/375460289 - flesh out non-promotional msg support that does
 // not return an error if missing but the default msg
 template <typename T>

hpb_generator/gen_utils.cc

@@ -13,10 +13,18 @@
 
 #include "absl/log/absl_log.h"
 #include "absl/strings/ascii.h"
+#include "absl/strings/escaping.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_replace.h"
 #include "absl/strings/string_view.h"
 #include "google/protobuf/descriptor.h"
 
+namespace {
+std::string EscapeTrigraphs(absl::string_view to_escape) {
+  return absl::StrReplaceAll(to_escape, {{"?", "\\?"}});
+}
+}  // namespace
+
 namespace google::protobuf::hpb_generator {
 
 namespace protobuf = ::proto2;
@@ -151,6 +159,10 @@ std::string DefaultValue(const FieldDescriptor* field) {
       return field->default_value_bool() ? "true" : "false";
     case FieldDescriptor::CPPTYPE_MESSAGE:
       return "::std::false_type()";
+    case FieldDescriptor::CPPTYPE_STRING:
+      return absl::StrCat(
+          "\"", EscapeTrigraphs(absl::CEscape(field->default_value_string())),
+          "\"");
     default:
       // TODO: b/375460289 - implement rest of scalars
       ABSL_LOG(WARNING) << "Unsupported default value type (in-progress): <"

hpb_generator/tests/extension_test.cc

@@ -12,6 +12,7 @@
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/strings/string_view.h"
 #include "google/protobuf/compiler/hpb/tests/child_model.upb.proto.h"
 #include "google/protobuf/compiler/hpb/tests/test_extension.upb.proto.h"
 #include "google/protobuf/compiler/hpb/tests/test_model.upb.proto.h"
@@ -35,6 +36,8 @@ using ::hpb_unittest::someotherpackage::protos::int64_ext;
 using ::hpb_unittest::someotherpackage::protos::repeated_int32_ext;
 using ::hpb_unittest::someotherpackage::protos::repeated_int64_ext;
 using ::hpb_unittest::someotherpackage::protos::repeated_string_ext;
+using ::hpb_unittest::someotherpackage::protos::string_ext;
+using ::hpb_unittest::someotherpackage::protos::string_trigraph_ext;
 using ::hpb_unittest::someotherpackage::protos::uint32_ext;
 using ::hpb_unittest::someotherpackage::protos::uint64_ext;
 
@@ -126,6 +129,15 @@ TEST(CppGeneratedCode, GetSetExtensionBool) {
   EXPECT_THAT(hpb::GetExtension(&model, bool_ext), IsOkAndHolds(true));
 }
 
+TEST(CppGeneratedCode, GetSetExtensionString) {
+  TestModel model;
+  EXPECT_EQ(false, hpb::HasExtension(&model, string_ext));
+  absl::string_view val = "Hello World";
+  auto x = hpb::SetExtension(&model, string_ext, val);
+  EXPECT_EQ(true, hpb::HasExtension(&model, string_ext));
+  EXPECT_THAT(hpb::GetExtension(&model, string_ext), IsOkAndHolds(val));
+}
+
 TEST(CppGeneratedCode, SetExtension) {
   TestModel model;
   void* prior_message;
@@ -366,6 +378,20 @@ TEST(CppGeneratedCode, GetExtensionBoolWithDefault) {
   EXPECT_THAT(res, IsOkAndHolds(true));
 }
 
+TEST(CppGeneratedCode, GetExtensionStringWithDefault) {
+  TestModel model;
+  auto res = hpb::GetExtension(&model, string_ext);
+  EXPECT_TRUE(res.ok());
+  EXPECT_THAT(res, IsOkAndHolds("mishpacha"));
+}
+
+TEST(CppGeneratedCode, GetExtensionStringWithDefaultAndTrigraph) {
+  TestModel model;
+  auto res = hpb::GetExtension(&model, string_trigraph_ext);
+  EXPECT_TRUE(res.ok());
+  EXPECT_THAT(res, IsOkAndHolds("bseder??!bseder"));
+}
+
 TEST(CppGeneratedCode, GetExtensionOnMutableChild) {
   TestModel model;
   ThemeExtension extension1;
@@ -567,6 +593,18 @@ TEST(CppGeneratedCode, GetExtensionRepeatedi64) {
   EXPECT_EQ((*res)[0], 322);
 }
 
+TEST(CppGeneratedCode, GetExtensionSingularString) {
+  TestModel model;
+  hpb::Arena arena;
+  hpb::ExtensionRegistry extensions(arena);
+  extensions.AddExtension(string_ext);
+  // These bytes represent a singular string field: "todaraba" @index 13012.
+  auto bytes = "\242\255\006\010todaraba";
+  auto parsed_model = hpb::Parse<TestModel>(bytes, extensions).value();
+  auto res = hpb::GetExtension(&parsed_model, string_ext);
+  EXPECT_THAT(res, IsOkAndHolds("todaraba"));
+}
+
 TEST(CppGeneratedCode, GetExtensionRepeatedString) {
   TestModel model;
   upb::Arena arena;

hpb_generator/tests/test_extension.proto

@@ -32,4 +32,7 @@ extend TestModel {
   double double_ext = 13010
       [default = 340282000000000000000000000000000000001.23];
   bool bool_ext = 13011 [default = true];
+
+  string string_ext = 13012 [default = "mishpacha"];
+  string string_trigraph_ext = 13013 [default = "bseder??!bseder"];
 }