Add ::hpb::ExtensionNumber API for extension field constants.

PiperOrigin-RevId: 653385818

hpb/hpb.h

@@ -8,11 +8,13 @@
 #ifndef PROTOBUF_HPB_HPB_H_
 #define PROTOBUF_HPB_HPB_H_
 
+#include <cstdint>
 #include <type_traits>
 #include <vector>
 
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
 #include "upb/base/status.hpp"
 #include "upb/mem/arena.hpp"
 #include "upb/message/copy.h"
@@ -148,6 +150,10 @@ struct PrivateAccess {
   static auto CreateMessage(upb_Arena* arena) {
     return typename T::Proxy(upb_Message_New(T::minitable(), arena), arena);
   }
+  template <typename ExtensionId>
+  static constexpr uint32_t GetExtensionNumber(const ExtensionId& id) {
+    return id.number();
+  }
 };
 
 template <typename T>
@@ -200,6 +206,12 @@ class ExtensionIdentifier : public ExtensionMiniTableProvider {
   constexpr explicit ExtensionIdentifier(
       const upb_MiniTableExtension* mini_table_ext)
       : ExtensionMiniTableProvider(mini_table_ext) {}
+
+ private:
+  constexpr uint32_t number() const {
+    return upb_MiniTableExtension_Number(mini_table_ext());
+  }
+  friend class PrivateAccess;
 };
 
 template <typename T>
@@ -574,6 +586,12 @@ absl::StatusOr<absl::string_view> Serialize(Ptr<T> message, upb::Arena& arena,
       ::protos::internal::GetMiniTable(message), arena.ptr(), options);
 }
 
+template <typename T, typename Extension>
+constexpr uint32_t ExtensionNumber(
+    internal::ExtensionIdentifier<T, Extension> id) {
+  return internal::PrivateAccess::GetExtensionNumber(id);
+}
+
 }  // namespace protos
 
 #endif  // PROTOBUF_HPB_HPB_H_

hpb_generator/tests/BUILD

@@ -156,8 +156,9 @@
 #         "@com_google_absl//absl/strings",
 #         "//hpb:requires",
 #         "//hpb",
-#         "//upb:mem",
 #         "//hpb:repeated_field",
+#         "//protos",
+#         "//upb:mem",
 #     ],
 # )
 # end:google_only

hpb_generator/tests/test_generated.cc

@@ -20,8 +20,10 @@
 #include "absl/strings/string_view.h"
 #include "google/protobuf/compiler/hpb/tests/child_model.upb.proto.h"
 #include "google/protobuf/compiler/hpb/tests/no_package.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"
 #include "google/protobuf/hpb/hpb.h"
+#include "google/protobuf/hpb/repeated_field.h"
 #include "google/protobuf/hpb/requires.h"
 #include "upb/mem/arena.h"
 #include "upb/mem/arena.hpp"
@@ -154,11 +156,11 @@ TEST(CppGeneratedCode, ScalarInt64) {
   EXPECT_EQ(testModel.optional_int64(), 0);
   EXPECT_FALSE(testModel.has_optional_int64());
   // Set value.
-  testModel.set_optional_int64(0xFF00CCDDA0001000);
+  testModel.set_optional_int64(static_cast<int64_t>(0xFF00CCDDA0001000));
   EXPECT_TRUE(testModel.has_optional_int64());
   EXPECT_EQ(testModel.optional_int64(), 0xFF00CCDDA0001000);
   // Change value.
-  testModel.set_optional_int64(0xFF00CCDD70002000);
+  testModel.set_optional_int64(static_cast<int64_t>(0xFF00CCDD70002000));
   EXPECT_TRUE(testModel.has_optional_int64());
   EXPECT_EQ(testModel.optional_int64(), 0xFF00CCDD70002000);
   // Clear value.
@@ -166,7 +168,7 @@ TEST(CppGeneratedCode, ScalarInt64) {
   EXPECT_FALSE(testModel.has_optional_int64());
   EXPECT_EQ(testModel.optional_int64(), 0);
   // Set after clear.
-  testModel.set_optional_int64(0xFF00CCDDA0001000);
+  testModel.set_optional_int64(static_cast<int64_t>(0xFF00CCDDA0001000));
   EXPECT_TRUE(testModel.has_optional_int64());
   EXPECT_EQ(testModel.optional_int64(), 0xFF00CCDDA0001000);
 }
@@ -1239,6 +1241,10 @@ TEST(CppGeneratedCode, FieldNumberConstants) {
   EXPECT_EQ(225, TestModel::kChildMapFieldNumber);
 }
 
+TEST(CppGeneratedCode, ExtensionFieldNumberConstant) {
+  EXPECT_EQ(12003, ::protos::ExtensionNumber(ThemeExtension::theme_extension));
+}
+
 TEST(CppGeneratedCode, ClearConstMessageShouldFailForConstChild) {
   TestModel model;
   EXPECT_FALSE(CanCallClearMessage<decltype(model.child_model_1())>());