From b72722a8d6068e529b5d55d18d7fc32bd0acc6c3 Mon Sep 17 00:00:00 2001
From: Hong Shin <hongshin@google.com>
Date: Thu, 31 Oct 2024 09:10:07 -0700
Subject: [PATCH] hpb: utilize layout_index as opposed to raw index

.index() is dependent on the order specified in the .proto file.
Minitables create their own ordering, which represent the true index that we're interested in set_alias.

This can be fetched via .layout_index when we have a upb::FieldDefPtr.

PiperOrigin-RevId: 691825782
---
 hpb_generator/context.h         | 22 ++++++++++++++++++----
 hpb_generator/gen_accessors.cc  |  4 ++--
 hpb_generator/protoc-gen-hpb.cc |  8 +++++---
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/hpb_generator/context.h b/hpb_generator/context.h
index d7a7ae9bb0..f1e989fb62 100644
--- a/hpb_generator/context.h
+++ b/hpb_generator/context.h
@@ -9,7 +9,6 @@
 #define GOOGLE_PROTOBUF_COMPILER_HPB_CONTEXT_H__
 
 #include <string>
-#include <utility>
 
 #include "absl/strings/ascii.h"
 #include "absl/strings/str_cat.h"
@@ -21,7 +20,9 @@
 #include "google/protobuf/descriptor.h"
 #include "google/protobuf/io/printer.h"
 #include "google/protobuf/io/zero_copy_stream.h"
-#include "upb_generator/c/names.h"
+#include "upb/reflection/def.hpp"
+#include "upb_generator/common/cpp_to_upb_def.h"
+
 namespace google::protobuf::hpb_generator {
 
 enum class Backend { UPB, CPP };
@@ -44,8 +45,11 @@ struct Options {
  */
 class Context final {
  public:
-  Context(io::ZeroCopyOutputStream* stream, const Options& options)
-      : stream_(stream), printer_(stream_), options_(options) {}
+  Context(const FileDescriptor* file, io::ZeroCopyOutputStream* stream,
+          const Options& options)
+      : stream_(stream), printer_(stream_), options_(options) {
+    BuildDefPool(file);
+  }
 
   void Emit(absl::Span<const io::Printer::Sub> vars, absl::string_view format,
             absl::SourceLocation loc = absl::SourceLocation::current()) {
@@ -68,15 +72,25 @@ class Context final {
   const Options& options() { return options_; }
   io::Printer& printer() { return printer_; }
 
+  inline std::string GetLayoutIndex(const FieldDescriptor* field) {
+    return absl::StrCat(
+        upb::generator::FindBaseFieldDef(pool_, field).layout_index());
+  }
+
   Context(const Context&) = delete;
   Context& operator=(const Context&) = delete;
   Context(Context&&) = delete;
   Context& operator=(Context&&) = delete;
 
  private:
+  inline void BuildDefPool(const FileDescriptor* file) {
+    upb::generator::AddFile(file, &pool_);
+  }
+
   io::ZeroCopyOutputStream* stream_;
   io::Printer printer_;
   const Options& options_;
+  upb::DefPool pool_;
 };
 
 // TODO: b/373438292 - re-house these 4 legacy funcs post io::Printer move
diff --git a/hpb_generator/gen_accessors.cc b/hpb_generator/gen_accessors.cc
index d038374963..d480c85265 100644
--- a/hpb_generator/gen_accessors.cc
+++ b/hpb_generator/gen_accessors.cc
@@ -286,7 +286,7 @@ void WriteAccessorsInSource(const protobuf::Descriptor* desc, Context& ctx) {
                 ABSL_CHECK(upb_Arena_IsFused(arena_, hpb::interop::upb::GetArena(target)));
                 upb_Message_SetBaseFieldMessage(
                     UPB_UPCAST(msg_),
-                    upb_MiniTable_FindFieldByNumber($7::minitable(), $8),
+                    upb_MiniTable_GetFieldByIndex($7::minitable(), $8),
                     hpb::interop::upb::GetMessage(target));
               }
             )cc",
@@ -294,7 +294,7 @@ void WriteAccessorsInSource(const protobuf::Descriptor* desc, Context& ctx) {
             resolved_field_name,
             upb::generator::CApiMessageType(desc->full_name()),
             MessageBaseType(field, /* maybe_const */ false), resolved_upbc_name,
-            arena_expression, ClassName(desc), field->number());
+            arena_expression, ClassName(desc), ctx.GetLayoutIndex(field));
       }
     }
   }
diff --git a/hpb_generator/protoc-gen-hpb.cc b/hpb_generator/protoc-gen-hpb.cc
index 30d6b4e31e..e7cb0f0a6e 100644
--- a/hpb_generator/protoc-gen-hpb.cc
+++ b/hpb_generator/protoc-gen-hpb.cc
@@ -81,20 +81,22 @@ bool Generator::Generate(const protobuf::FileDescriptor* file,
   // Write model.upb.fwd.h
   std::unique_ptr<google::protobuf::io::ZeroCopyOutputStream> output_stream(
       context->Open(ForwardingHeaderFilename(file)));
-  auto fwd_ctx = Context(output_stream.get(), Options{.backend = Backend::UPB});
+  Context fwd_ctx =
+      Context(file, output_stream.get(), Options{.backend = Backend::UPB});
   WriteForwardingHeader(file, fwd_ctx);
 
   // Write model.upb.proto.h
   std::unique_ptr<google::protobuf::io::ZeroCopyOutputStream> header_output_stream(
       context->Open(CppHeaderFilename(file)));
-  Context hdr_ctx(header_output_stream.get(), Options{.backend = Backend::UPB});
+  Context hdr_ctx(file, header_output_stream.get(),
+                  Options{.backend = Backend::UPB});
   WriteHeader(file, hdr_ctx, strip_nonfunctional_codegen);
 
   // Write model.upb.proto.cc
   std::unique_ptr<google::protobuf::io::ZeroCopyOutputStream> cc_output_stream(
       context->Open(CppSourceFilename(file)));
   auto cc_ctx =
-      Context(cc_output_stream.get(), Options{.backend = Backend::UPB});
+      Context(file, cc_output_stream.get(), Options{.backend = Backend::UPB});
   WriteSource(file, cc_ctx, fasttable_enabled, strip_nonfunctional_codegen);
   return true;
 }