[Remarks] Allow remarks::Format::YAML to take a string table

It should be allowed to take a string table in case all the strings in
the remarks point there, but it shouldn't use it during serialization.

llvm-svn: 372042

4 files changed, 41 insertions, 26 deletions

diff --git a/llvm/lib/Remarks/BitstreamRemarkSerializer.cpp b/llvm/lib/Remarks/BitstreamRemarkSerializer.cpp
index 03c5697..d02782c 100644
--- a/llvm/lib/Remarks/BitstreamRemarkSerializer.cpp
+++ b/llvm/lib/Remarks/BitstreamRemarkSerializer.cpp
@@ -326,7 +326,7 @@ StringRef BitstreamRemarkSerializerHelper::getBuffer() {
 
 BitstreamRemarkSerializer::BitstreamRemarkSerializer(raw_ostream &OS,
                                                      SerializerMode Mode)
-    : RemarkSerializer(OS, Mode),
+    : RemarkSerializer(Format::Bitstream, OS, Mode),
       Helper(BitstreamRemarkContainerType::SeparateRemarksFile) {
   assert(Mode == SerializerMode::Separate &&
          "For SerializerMode::Standalone, a pre-filled string table needs to "
@@ -338,7 +338,7 @@ BitstreamRemarkSerializer::BitstreamRemarkSerializer(raw_ostream &OS,
 BitstreamRemarkSerializer::BitstreamRemarkSerializer(raw_ostream &OS,
                                                      SerializerMode Mode,
                                                      StringTable StrTabIn)
-    : RemarkSerializer(OS, Mode),
+    : RemarkSerializer(Format::Bitstream, OS, Mode),
       Helper(Mode == SerializerMode::Separate
                  ? BitstreamRemarkContainerType::SeparateRemarksFile
                  : BitstreamRemarkContainerType::Standalone) {
diff --git a/llvm/lib/Remarks/RemarkSerializer.cpp b/llvm/lib/Remarks/RemarkSerializer.cpp
index caba7bb..ab19c84 100644
--- a/llvm/lib/Remarks/RemarkSerializer.cpp
+++ b/llvm/lib/Remarks/RemarkSerializer.cpp
@@ -42,15 +42,13 @@ remarks::createRemarkSerializer(Format RemarksFormat, SerializerMode Mode,
     return createStringError(std::errc::invalid_argument,
                              "Unknown remark serializer format.");
   case Format::YAML:
-    return createStringError(std::errc::invalid_argument,
-                             "Unable to use a string table with the yaml "
-                             "format. Use 'yaml-strtab' instead.");
+    return std::make_unique<YAMLRemarkSerializer>(OS, Mode, std::move(StrTab));
   case Format::YAMLStrTab:
     return std::make_unique<YAMLStrTabRemarkSerializer>(OS, Mode,
-                                                         std::move(StrTab));
+                                                        std::move(StrTab));
   case Format::Bitstream:
     return std::make_unique<BitstreamRemarkSerializer>(OS, Mode,
-                                                        std::move(StrTab));
+                                                       std::move(StrTab));
   }
   llvm_unreachable("Unknown remarks::Format enum");
 }
diff --git a/llvm/lib/Remarks/RemarkStringTable.cpp b/llvm/lib/Remarks/RemarkStringTable.cpp
index 0f8a95a..5115646 100644
--- a/llvm/lib/Remarks/RemarkStringTable.cpp
+++ b/llvm/lib/Remarks/RemarkStringTable.cpp
@@ -46,7 +46,6 @@ void StringTable::internalize(Remark &R) {
   if (R.Loc)
     Impl(R.Loc->SourceFilePath);
   for (Argument &Arg : R.Args) {
-    // We need to mutate elements from an ArrayRef here.
     Impl(Arg.Key);
     Impl(Arg.Val);
     if (Arg.Loc)
diff --git a/llvm/lib/Remarks/YAMLRemarkSerializer.cpp b/llvm/lib/Remarks/YAMLRemarkSerializer.cpp
index a4d67cc..66eb06b 100644
--- a/llvm/lib/Remarks/YAMLRemarkSerializer.cpp
+++ b/llvm/lib/Remarks/YAMLRemarkSerializer.cpp
@@ -56,11 +56,14 @@ template <> struct MappingTraits<remarks::Remark *> {
     else
       llvm_unreachable("Unknown remark type");
 
-    if (Optional<StringTable> &StrTab =
-            reinterpret_cast<YAMLRemarkSerializer *>(io.getContext())->StrTab) {
-      unsigned PassID = StrTab->add(Remark->PassName).first;
-      unsigned NameID = StrTab->add(Remark->RemarkName).first;
-      unsigned FunctionID = StrTab->add(Remark->FunctionName).first;
+    if (auto *Serializer = dyn_cast<YAMLStrTabRemarkSerializer>(
+            reinterpret_cast<RemarkSerializer *>(io.getContext()))) {
+      assert(Serializer->StrTab.hasValue() &&
+             "YAMLStrTabSerializer with no StrTab.");
+      StringTable &StrTab = *Serializer->StrTab;
+      unsigned PassID = StrTab.add(Remark->PassName).first;
+      unsigned NameID = StrTab.add(Remark->RemarkName).first;
+      unsigned FunctionID = StrTab.add(Remark->FunctionName).first;
       mapRemarkHeader(io, PassID, NameID, Remark->Loc, FunctionID,
                       Remark->Hotness, Remark->Args);
     } else {
@@ -78,9 +81,12 @@ template <> struct MappingTraits<RemarkLocation> {
     unsigned Line = RL.SourceLine;
     unsigned Col = RL.SourceColumn;
 
-    if (Optional<StringTable> &StrTab =
-            reinterpret_cast<YAMLRemarkSerializer *>(io.getContext())->StrTab) {
-      unsigned FileID = StrTab->add(File).first;
+    if (auto *Serializer = dyn_cast<YAMLStrTabRemarkSerializer>(
+            reinterpret_cast<RemarkSerializer *>(io.getContext()))) {
+      assert(Serializer->StrTab.hasValue() &&
+             "YAMLStrTabSerializer with no StrTab.");
+      StringTable &StrTab = *Serializer->StrTab;
+      unsigned FileID = StrTab.add(File).first;
       io.mapRequired("File", FileID);
     } else {
       io.mapRequired("File", File);
@@ -130,9 +136,12 @@ template <> struct MappingTraits<Argument> {
   static void mapping(IO &io, Argument &A) {
     assert(io.outputting() && "input not yet implemented");
 
-    if (Optional<StringTable> &StrTab =
-            reinterpret_cast<YAMLRemarkSerializer *>(io.getContext())->StrTab) {
-      auto ValueID = StrTab->add(A.Val).first;
+    if (auto *Serializer = dyn_cast<YAMLStrTabRemarkSerializer>(
+            reinterpret_cast<RemarkSerializer *>(io.getContext()))) {
+      assert(Serializer->StrTab.hasValue() &&
+             "YAMLStrTabSerializer with no StrTab.");
+      StringTable &StrTab = *Serializer->StrTab;
+      auto ValueID = StrTab.add(A.Val).first;
       io.mapRequired(A.Key.data(), ValueID);
     } else if (StringRef(A.Val).count('\n') > 1) {
       StringBlockVal S(A.Val);
@@ -149,8 +158,17 @@ template <> struct MappingTraits<Argument> {
 
 LLVM_YAML_IS_SEQUENCE_VECTOR(Argument)
 
-YAMLRemarkSerializer::YAMLRemarkSerializer(raw_ostream &OS, SerializerMode Mode)
-    : RemarkSerializer(OS, Mode), YAMLOutput(OS, reinterpret_cast<void *>(this)) {}
+YAMLRemarkSerializer::YAMLRemarkSerializer(raw_ostream &OS, SerializerMode Mode,
+                                           Optional<StringTable> StrTabIn)
+    : YAMLRemarkSerializer(Format::YAML, OS, Mode, std::move(StrTabIn)) {}
+
+YAMLRemarkSerializer::YAMLRemarkSerializer(Format SerializerFormat,
+                                           raw_ostream &OS, SerializerMode Mode,
+                                           Optional<StringTable> StrTabIn)
+    : RemarkSerializer(SerializerFormat, OS, Mode),
+      YAMLOutput(OS, reinterpret_cast<void *>(this)) {
+  StrTab = std::move(StrTabIn);
+}
 
 void YAMLRemarkSerializer::emit(const Remark &Remark) {
   // Again, YAMLTraits expect a non-const object for inputting, but we're not
@@ -183,7 +201,7 @@ std::unique_ptr<MetaSerializer> YAMLStrTabRemarkSerializer::metaSerializer(
     raw_ostream &OS, Optional<StringRef> ExternalFilename) {
   assert(StrTab);
   return std::make_unique<YAMLStrTabMetaSerializer>(OS, ExternalFilename,
-                                                     std::move(*StrTab));
+                                                    *StrTab);
 }
 
 static void emitMagic(raw_ostream &OS) {
@@ -200,9 +218,9 @@ static void emitVersion(raw_ostream &OS) {
   OS.write(Version.data(), Version.size());
 }
 
-static void emitStrTab(raw_ostream &OS, const Optional<StringTable> &StrTab) {
+static void emitStrTab(raw_ostream &OS, Optional<const StringTable *> StrTab) {
   // Emit the string table in the section.
-  uint64_t StrTabSize = StrTab ? StrTab->SerializedSize : 0;
+  uint64_t StrTabSize = StrTab ? (*StrTab)->SerializedSize : 0;
   // Emit the total size of the string table (the size itself excluded):
   // little-endian uint64_t.
   // Note: even if no string table is used, emit 0.
@@ -210,7 +228,7 @@ static void emitStrTab(raw_ostream &OS, const Optional<StringTable> &StrTab) {
   support::endian::write64le(StrTabSizeBuf.data(), StrTabSize);
   OS.write(StrTabSizeBuf.data(), StrTabSizeBuf.size());
   if (StrTab)
-    StrTab->serialize(OS);
+    (*StrTab)->serialize(OS);
 }
 
 static void emitExternalFile(raw_ostream &OS, StringRef Filename) {
@@ -233,7 +251,7 @@ void YAMLMetaSerializer::emit() {
 void YAMLStrTabMetaSerializer::emit() {
   emitMagic(OS);
   emitVersion(OS);
-  emitStrTab(OS, std::move(StrTab));
+  emitStrTab(OS, &StrTab);
   if (ExternalFilename)
     emitExternalFile(OS, *ExternalFilename);
 }