From 550f9b021bd3a4aa13e31b3b59305cfb0e067735 Mon Sep 17 00:00:00 2001
From: Adam Cozzette <acozzette@google.com>
Date: Tue, 13 Jun 2023 16:07:58 -0700
Subject: [PATCH] Fix handling of public imports in upb `protos` generator

I recently ran into an issue with public imports and the `protos` generator.
The includes of fwd.h files are currently generated based on where the message
or extension is defined, but when public imports are involved, this fwd.h
header might not be in the direct dependencies. When the `#include` points to a
transitive dependency then this can result in a layering violation.

This CL fixes the problem by generating fwd.h includes for all direct
dependencies, and making sure the fwd.h files themselves contain `#include`
lines to export their public dependencies.

PiperOrigin-RevId: 540102455
---
 protos_generator/protoc-gen-upb-protos.cc | 48 ++++++-----------------
 1 file changed, 11 insertions(+), 37 deletions(-)

diff --git a/protos_generator/protoc-gen-upb-protos.cc b/protos_generator/protoc-gen-upb-protos.cc
index 0355b786c4..48e69ab3c1 100644
--- a/protos_generator/protoc-gen-upb-protos.cc
+++ b/protos_generator/protoc-gen-upb-protos.cc
@@ -57,8 +57,6 @@ void WriteTypedefForwardingHeader(
     Output& output);
 void WriteHeaderMessageForwardDecls(
     const protobuf::FileDescriptor* file,
-    const std::vector<const protobuf::Descriptor*>& file_messages,
-    const std::vector<const protobuf::FieldDescriptor*>& file_exts,
     Output& output);
 
 class Generator : public protoc::CodeGenerator {
@@ -115,6 +113,13 @@ void WriteForwardingHeader(const protobuf::FileDescriptor* file,
       )cc",
       ToPreproc(file->name()));
   output("\n");
+  for (int i = 0; i < file->public_dependency_count(); ++i) {
+    output("#include \"$0\"\n",
+           ForwardingHeaderFilename(file->public_dependency(i)));
+  }
+  if (file->public_dependency_count() > 0) {
+    output("\n");
+  }
   const std::vector<const protobuf::Descriptor*> this_file_messages =
       SortedMessages(file);
   WriteTypedefForwardingHeader(file, this_file_messages, output);
@@ -162,8 +167,7 @@ void WriteHeader(const protobuf::FileDescriptor* file, Output& output) {
     output("\n");
   }
 
-  WriteHeaderMessageForwardDecls(file, this_file_messages, this_file_exts,
-                                 output);
+  WriteHeaderMessageForwardDecls(file, output);
   WriteStartNamespace(file, output);
 
   std::vector<const protobuf::EnumDescriptor*> this_file_enums =
@@ -258,43 +262,13 @@ void WriteTypedefForwardingHeader(
 /// Writes includes for upb C minitables and fwd.h for transitive typedefs.
 void WriteHeaderMessageForwardDecls(
     const protobuf::FileDescriptor* file,
-    const std::vector<const protobuf::Descriptor*>& file_messages,
-    const std::vector<const protobuf::FieldDescriptor*>& file_exts,
     Output& output) {
   // Import forward-declaration of types defined in this file.
   output("#include \"$0\"\n", UpbCFilename(file));
   output("#include \"$0\"\n", ForwardingHeaderFilename(file));
-  // Forward-declare types not in this file, but used as submessages.
-  // Order by full name for consistent ordering.
-  std::map<std::string, const protobuf::Descriptor*> forward_messages;
-
-  for (auto* message : file_messages) {
-    for (int i = 0; i < message->field_count(); i++) {
-      const protobuf::FieldDescriptor* field = message->field(i);
-      if (field->cpp_type() == protobuf::FieldDescriptor::CPPTYPE_MESSAGE &&
-          field->file() != field->message_type()->file()) {
-        forward_messages[field->message_type()->full_name()] =
-            field->message_type();
-      }
-    }
-  }
-  for (auto* ext : file_exts) {
-    if (ext->file() != ext->containing_type()->file()) {
-      forward_messages[ext->containing_type()->full_name()] =
-          ext->containing_type();
-      if (ext->cpp_type() == protobuf::FieldDescriptor::CPPTYPE_MESSAGE) {
-        forward_messages[ext->message_type()->full_name()] =
-            ext->message_type();
-      }
-    }
-  }
-  std::map<std::string, const protobuf::FileDescriptor*> files_to_import;
-  for (const auto& pair : forward_messages) {
-    files_to_import[ForwardingHeaderFilename(pair.second->file())] = file;
-  }
-  for (const auto& pair : files_to_import) {
-    output("#include \"$0\"\n", UpbCFilename(pair.second));
-    output("#include \"$0\"\n", pair.first);
+  // Import forward-declaration of types in dependencies.
+  for (int i = 0; i < file->dependency_count(); ++i) {
+    output("#include \"$0\"\n", ForwardingHeaderFilename(file->dependency(i)));
   }
   output("\n");
 }