hpb: implement SetExtension for scalars i32 and i64

PiperOrigin-RevId: 693799584

csharp/src/Google.Protobuf/Reflection/FeatureSetDescriptor.g.cs

@@ -1,17 +0,0 @@
-#region Copyright notice and license
-// Protocol Buffers - Google's data interchange format
-// Copyright 2008 Google Inc.  All rights reserved.
-//
-// Use of this source code is governed by a BSD-style
-// license that can be found in the LICENSE file or at
-// https://developers.google.com/open-source/licenses/bsd
-#endregion
-
-namespace Google.Protobuf.Reflection;
-
-internal sealed partial class FeatureSetDescriptor
-{
-    // Canonical serialized form of the edition defaults, generated by embed_edition_defaults.
-    private const string DefaultsBase64 =
-        "ChMYhAciACoMCAEQAhgCIAMoATACChMY5wciACoMCAIQARgBIAIoATABChMY6AciDAgBEAEYASACKAEwASoAIOYHKOgH";
-}

hpb/extension.h

@@ -13,6 +13,8 @@
 #include <vector>
 
 #include "absl/base/attributes.h"
+#include "absl/log/absl_log.h"
+#include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "google/protobuf/hpb/backend/upb/interop.h"
 #include "google/protobuf/hpb/internal/message_lock.h"
@@ -55,6 +57,7 @@ struct UpbExtensionTrait<int32_t> {
   using DefaultType = int32_t;
   using ReturnType = int32_t;
   static constexpr auto kGetter = upb_Message_GetExtensionInt32;
+  static constexpr auto kSetter = upb_Message_SetExtensionInt32;
 };
 
 template <>
@@ -62,6 +65,7 @@ struct UpbExtensionTrait<int64_t> {
   using DefaultType = int64_t;
   using ReturnType = int64_t;
   static constexpr auto kGetter = upb_Message_GetExtensionInt64;
+  static constexpr auto kSetter = upb_Message_SetExtensionInt64;
 };
 
 // TODO: b/375460289 - flesh out non-promotional msg support that does
@@ -184,11 +188,18 @@ absl::Status SetExtension(
     Ptr<T> message,
     const ::hpb::internal::ExtensionIdentifier<T, Extension>& id,
     const Extension& value) {
-  static_assert(!std::is_const_v<T>);
-  auto* message_arena = hpb::interop::upb::GetArena(message);
-  return ::hpb::internal::SetExtension(hpb::interop::upb::GetMessage(message),
-                                       message_arena, id.mini_table_ext(),
-                                       hpb::interop::upb::GetMessage(&value));
+  if constexpr (std::is_integral_v<Extension>) {
+    bool res = hpb::internal::UpbExtensionTrait<Extension>::kSetter(
+        hpb::interop::upb::GetMessage(message), id.mini_table_ext(), value,
+        hpb::interop::upb::GetArena(message));
+    return res ? absl::OkStatus() : MessageAllocationError();
+  } else {
+    static_assert(!std::is_const_v<T>);
+    auto* message_arena = hpb::interop::upb::GetArena(message);
+    return ::hpb::internal::SetExtension(hpb::interop::upb::GetMessage(message),
+                                         message_arena, id.mini_table_ext(),
+                                         hpb::interop::upb::GetMessage(&value));
+  }
 }
 
 template <typename T, typename Extension,
@@ -212,14 +223,21 @@ absl::Status SetExtension(
     Ptr<T> message,
     const ::hpb::internal::ExtensionIdentifier<T, Extension>& id,
     Extension&& value) {
-  Extension ext = std::move(value);
-  static_assert(!std::is_const_v<T>);
-  auto* message_arena = hpb::interop::upb::GetArena(message);
-  auto* extension_arena = hpb::interop::upb::GetArena(&ext);
-  return ::hpb::internal::MoveExtension(hpb::interop::upb::GetMessage(message),
-                                        message_arena, id.mini_table_ext(),
-                                        hpb::interop::upb::GetMessage(&ext),
-                                        extension_arena);
+  if constexpr (std::is_integral_v<Extension>) {
+    bool res = hpb::internal::UpbExtensionTrait<Extension>::kSetter(
+        hpb::interop::upb::GetMessage(message), id.mini_table_ext(), value,
+        hpb::interop::upb::GetArena(message));
+    return res ? absl::OkStatus() : MessageAllocationError();
+  } else {
+    Extension ext = std::forward<Extension>(value);
+    static_assert(!std::is_const_v<T>);
+    auto* message_arena = hpb::interop::upb::GetArena(message);
+    auto* extension_arena = hpb::interop::upb::GetArena(&ext);
+    return ::hpb::internal::MoveExtension(
+        hpb::interop::upb::GetMessage(message), message_arena,
+        id.mini_table_ext(), hpb::interop::upb::GetMessage(&ext),
+        extension_arena);
+  }
 }
 
 template <typename T, typename Extension,

hpb_generator/tests/test_generated.cc

@@ -5,13 +5,11 @@
 // license that can be found in the LICENSE file or at
 // https://developers.google.com/open-source/licenses/bsd
 
-#include <climits>
 #include <cstdint>
 #include <iterator>
 #include <limits>
 #include <memory>
 #include <string>
-#include <type_traits>
 #include <utility>
 #include <vector>
 
@@ -58,6 +56,7 @@ using ::hpb_unittest::protos::ThemeExtension;
 using ::hpb_unittest::someotherpackage::protos::int32_ext;
 using ::hpb_unittest::someotherpackage::protos::int64_ext;
 using ::testing::ElementsAre;
+using ::testing::status::IsOkAndHolds;
 
 TEST(CppGeneratedCode, Constructor) { TestModel test_model; }
 
@@ -734,6 +733,24 @@ TEST(CppGeneratedCode, ClearExtensionWithEmptyExtensionPtr) {
   EXPECT_EQ(false, ::hpb::HasExtension(recursive_child, theme));
 }
 
+TEST(CppGeneratedCode, SetExtensionInt32) {
+  TestModel model;
+  EXPECT_EQ(false, hpb::HasExtension(&model, int32_ext));
+  int32_t val = 55;
+  auto x = hpb::SetExtension(&model, int32_ext, val);
+  EXPECT_EQ(true, hpb::HasExtension(&model, int32_ext));
+  EXPECT_THAT(hpb::GetExtension(&model, int32_ext), IsOkAndHolds(val));
+}
+
+TEST(CppGeneratedCode, SetExtensionInt64) {
+  TestModel model;
+  EXPECT_EQ(false, hpb::HasExtension(&model, int64_ext));
+  int64_t val = std::numeric_limits<int32_t>::max() + int64_t{1};
+  auto x = hpb::SetExtension(&model, int64_ext, val);
+  EXPECT_EQ(true, hpb::HasExtension(&model, int64_ext));
+  EXPECT_THAT(hpb::GetExtension(&model, int64_ext), IsOkAndHolds(val));
+}
+
 TEST(CppGeneratedCode, SetExtension) {
   TestModel model;
   void* prior_message;