Add GenerateCodeInfo::Annotation::Semantic support to io::Printer.

PiperOrigin-RevId: 507606033

src/google/protobuf/compiler/cpp/helpers.h

@@ -757,6 +757,8 @@ inline bool HasImplData(const Descriptor* desc, const Options& options) {
   return !HasSimpleBaseClass(desc, options);
 }
 
+// DO NOT USE IN NEW CODE! Use io::Printer directly instead. See b/242326974.
+//
 // Formatter is a functor class which acts as a closure around printer and
 // the variable map. It's much like printer->Print except it supports both named
 // variables that are substituted using a key value map and direct arguments. In

src/google/protobuf/io/printer.cc

@@ -631,7 +631,7 @@ void Printer::PrintImpl(absl::string_view format,
           if (options_.annotation_collector != nullptr) {
             options_.annotation_collector->AddAnnotation(
                 record_var.second, sink_.bytes_written(), record->file_path,
-                record->path);
+                record->path, record->semantic);
           }
         }
 
@@ -743,7 +743,7 @@ void Printer::PrintImpl(absl::string_view format,
           options_.annotation_collector != nullptr) {
         options_.annotation_collector->AddAnnotation(
             range_start, range_end, same_name_record->file_path,
-            same_name_record->path);
+            same_name_record->path, same_name_record->semantic);
       }
 
       if (opts.use_substitution_map) {

src/google/protobuf/io/printer.h

@@ -49,6 +49,7 @@
 #include "absl/container/flat_hash_map.h"
 #include "absl/functional/function_ref.h"
 #include "absl/log/absl_check.h"
+#include "absl/meta/type_traits.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/optional.h"
@@ -70,6 +71,15 @@ class PROTOBUF_EXPORT AnnotationCollector {
   // specific to derived types of AnnotationCollector.
   using Annotation = std::pair<std::pair<size_t, size_t>, std::string>;
 
+  // The semantic meaning of an annotation. This enum mirrors
+  // google.protobuf.GeneratedCodeInfo.Annotation.Semantic, and the enumerator values
+  // should match it.
+  enum Semantic {
+    kNone = 0,
+    kSet = 1,
+    kAlias = 2,
+  };
+
   virtual ~AnnotationCollector() = default;
 
   // Records that the bytes in file_path beginning with begin_offset and ending
@@ -78,6 +88,13 @@ class PROTOBUF_EXPORT AnnotationCollector {
                              const std::string& file_path,
                              const std::vector<int>& path) = 0;
 
+  virtual void AddAnnotation(size_t begin_offset, size_t end_offset,
+                             const std::string& file_path,
+                             const std::vector<int>& path,
+                             absl::optional<Semantic> semantic) {
+    AddAnnotation(begin_offset, end_offset, file_path, path);
+  }
+
   // TODO(gerbens) I don't see why we need virtuals here. Just a vector of
   // range, payload pairs stored in a context should suffice.
   virtual void AddAnnotationNew(Annotation&) {}
@@ -92,9 +109,28 @@ class PROTOBUF_EXPORT AnnotationCollector {
 //   optional string source_file = 2;
 //   optional int32 begin = 3;
 //   optional int32 end = 4;
+//   optional int32 semantic = 5;
 // }
 template <typename AnnotationProto>
 class AnnotationProtoCollector : public AnnotationCollector {
+ private:
+  // Some users of this type use it with a proto that does not have a
+  // "semantic" field. Therefore, we need to detect it with SFINAE.
+
+  // go/ranked-overloads
+  struct Rank0 {};
+  struct Rank1 : Rank0 {};
+
+  template <typename Proto>
+  static auto SetSemantic(Proto* p, int semantic, Rank1)
+      -> decltype(p->set_semantic(
+          static_cast<typename Proto::Semantic>(semantic))) {
+    return p->set_semantic(static_cast<typename Proto::Semantic>(semantic));
+  }
+
+  template <typename Proto>
+  static void SetSemantic(Proto*, int, Rank0) {}
+
  public:
   explicit AnnotationProtoCollector(AnnotationProto* annotation_proto)
       : annotation_proto_(annotation_proto) {}
@@ -102,6 +138,12 @@ class AnnotationProtoCollector : public AnnotationCollector {
   void AddAnnotation(size_t begin_offset, size_t end_offset,
                      const std::string& file_path,
                      const std::vector<int>& path) override {
+    AddAnnotation(begin_offset, end_offset, file_path, path, absl::nullopt);
+  }
+
+  void AddAnnotation(size_t begin_offset, size_t end_offset,
+                     const std::string& file_path, const std::vector<int>& path,
+                     absl::optional<Semantic> semantic) override {
     auto* annotation = annotation_proto_->add_annotation();
     for (int i = 0; i < path.size(); ++i) {
       annotation->add_path(path[i]);
@@ -109,6 +151,10 @@ class AnnotationProtoCollector : public AnnotationCollector {
     annotation->set_source_file(file_path);
     annotation->set_begin(begin_offset);
     annotation->set_end(end_offset);
+
+    if (semantic.has_value()) {
+      SetSemantic(annotation, *semantic, Rank1{});
+    }
   }
 
   void AddAnnotationNew(Annotation& a) override {
@@ -875,6 +921,7 @@ auto Printer::ValueImpl<owned>::ToStringOrCallback(Cb&& cb, Rank2)
 struct Printer::AnnotationRecord {
   std::vector<int> path;
   std::string file_path;
+  absl::optional<AnnotationCollector::Semantic> semantic;
 
   // AnnotationRecord's constructors are *not* marked as explicit,
   // specifically so that it is possible to construct a
@@ -887,15 +934,18 @@ struct Printer::AnnotationRecord {
       std::enable_if_t<std::is_convertible<const String&, std::string>::value,
                        int> = 0>
   AnnotationRecord(  // NOLINT(google-explicit-constructor)
-      const String& file_path)
-      : file_path(file_path) {}
+      const String& file_path,
+      absl::optional<AnnotationCollector::Semantic> semantic = absl::nullopt)
+      : file_path(file_path), semantic(semantic) {}
 
   template <typename Desc,
             // This SFINAE clause excludes char* from matching this
             // constructor.
             std::enable_if_t<std::is_class<Desc>::value, int> = 0>
-  AnnotationRecord(const Desc* desc)  // NOLINT(google-explicit-constructor)
-      : file_path(desc->file()->name()) {
+  AnnotationRecord(  // NOLINT(google-explicit-constructor)
+      const Desc* desc,
+      absl::optional<AnnotationCollector::Semantic> semantic = absl::nullopt)
+      : file_path(desc->file()->name()), semantic(semantic) {
     desc->GetLocationPath(&path);
   }
 };

src/google/protobuf/io/printer_unittest.cc

@@ -46,6 +46,7 @@
 #include "absl/log/absl_check.h"
 #include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
+#include "absl/types/optional.h"
 #include "google/protobuf/io/zero_copy_stream.h"
 #include "google/protobuf/io/zero_copy_stream_impl_lite.h"
 
@@ -189,6 +190,13 @@ class FakeAnnotationCollector : public AnnotationCollector {
         Record{begin_offset, end_offset, file_path, path});
   }
 
+  void AddAnnotation(size_t begin_offset, size_t end_offset,
+                     const std::string& file_path, const std::vector<int>& path,
+                     absl::optional<Semantic> semantic) override {
+    annotations_.emplace_back(
+        Record{begin_offset, end_offset, file_path, path, semantic});
+  }
+
   void AddAnnotationNew(Annotation& a) override {
     GeneratedCodeInfo::Annotation annotation;
     annotation.ParseFromString(a.second);
@@ -202,14 +210,17 @@ class FakeAnnotationCollector : public AnnotationCollector {
   }
 
   struct Record {
-    size_t start, end;
+    size_t start = 0;
+    size_t end = 0;
     std::string file_path;
     std::vector<int> path;
+    absl::optional<Semantic> semantic;
 
     friend std::ostream& operator<<(std::ostream& out, const Record& record) {
       return out << "Record{" << record.start << ", " << record.end << ", \""
                  << record.file_path << "\", ["
-                 << absl::StrJoin(record.path, ", ") << "]}";
+                 << absl::StrJoin(record.path, ", ") << "], "
+                 << record.semantic.value_or(kNone) << "}";
     }
   };
 
@@ -219,16 +230,18 @@ class FakeAnnotationCollector : public AnnotationCollector {
   std::vector<Record> annotations_;
 };
 
-template <typename Start, typename End, typename FilePath, typename Path>
-testing::Matcher<FakeAnnotationCollector::Record> Annotation(Start start,
-                                                             End end,
-                                                             FilePath file_path,
-                                                             Path path) {
-  return AllOf(Field("start", &FakeAnnotationCollector::Record::start, start),
-               Field("end", &FakeAnnotationCollector::Record::end, end),
-               Field("file_path", &FakeAnnotationCollector::Record::file_path,
-                     file_path),
-               Field("path", &FakeAnnotationCollector::Record::path, path));
+template <typename Start, typename End, typename FilePath, typename Path,
+          typename Semantic = absl::optional<AnnotationCollector::Semantic>>
+testing::Matcher<FakeAnnotationCollector::Record> Annotation(
+    Start start, End end, FilePath file_path, Path path,
+    Semantic semantic = absl::nullopt) {
+  return AllOf(
+      Field("start", &FakeAnnotationCollector::Record::start, start),
+      Field("end", &FakeAnnotationCollector::Record::end, end),
+      Field("file_path", &FakeAnnotationCollector::Record::file_path,
+            file_path),
+      Field("path", &FakeAnnotationCollector::Record::path, path),
+      Field("semantic", &FakeAnnotationCollector::Record::semantic, semantic));
 }
 
 TEST_F(PrinterTest, AnnotateMap) {
@@ -691,6 +704,32 @@ TEST_F(PrinterTest, EmitSameNameAnnotation) {
               ElementsAre(Annotation(6, 9, "file.proto", ElementsAre(33))));
 }
 
+TEST_F(PrinterTest, EmitSameNameAnnotationWithSemantic) {
+  FakeAnnotationCollector collector;
+  {
+    Printer printer(output(), '$', &collector);
+    FakeDescriptor descriptor{{"file.proto"}, {33}};
+    auto v = printer.WithVars({{"class", "Foo"}});
+    auto a = printer.WithAnnotations(
+        {{"class", {&descriptor, AnnotationCollector::kSet}}});
+
+    printer.Emit({{"f1", "x"}, {"f2", "y"}, {"f3", "z"}}, R"cc(
+      class $class$ {
+        int $f1$, $f2$, $f3$;
+      };
+    )cc");
+  }
+
+  EXPECT_EQ(written(),
+            "class Foo {\n"
+            "  int x, y, z;\n"
+            "};\n");
+
+  EXPECT_THAT(collector.Get(),
+              ElementsAre(Annotation(6, 9, "file.proto", ElementsAre(33),
+                                     AnnotationCollector::kSet)));
+}
+
 TEST_F(PrinterTest, EmitSameNameAnnotationFileNameOnly) {
   FakeAnnotationCollector collector;
   {