[Remarks] Support parsing remark metadata in the YAML remark parser

This adds support to the yaml remark parser to be able to parse remarks
directly from the metadata.

This supports parsing separate metadata and following the external file
with the associated metadata, and also a standalone file containing
metadata + remarks all together.

llvm-svn: 367148

3 files changed, 129 insertions, 0 deletions

diff --git a/llvm/lib/Remarks/RemarkParser.cpp b/llvm/lib/Remarks/RemarkParser.cpp
index e7f174ce02c..1ff9cf626b6 100644
--- a/llvm/lib/Remarks/RemarkParser.cpp
+++ b/llvm/lib/Remarks/RemarkParser.cpp
@@ -80,6 +80,21 @@ llvm::remarks::createRemarkParser(Format ParserFormat, StringRef Buf,
   llvm_unreachable("unhandled ParseFormat");
 }
 
+Expected<std::unique_ptr<RemarkParser>>
+llvm::remarks::createRemarkParserFromMeta(Format ParserFormat, StringRef Buf,
+                                          Optional<ParsedStringTable> StrTab) {
+  switch (ParserFormat) {
+  // Depending on the metadata, the format can be either yaml or yaml-strtab,
+  // regardless of the input argument.
+  case Format::YAML:
+  case Format::YAMLStrTab:
+    return createYAMLParserFromMeta(Buf, std::move(StrTab));
+  case Format::Unknown:
+    return createStringError(std::make_error_code(std::errc::invalid_argument),
+                             "Unknown remark parser format.");
+  }
+}
+
 // Wrapper that holds the state needed to interact with the C API.
 struct CParser {
   std::unique_ptr<RemarkParser> TheParser;
diff --git a/llvm/lib/Remarks/YAMLRemarkParser.cpp b/llvm/lib/Remarks/YAMLRemarkParser.cpp
index f677bfc04ce..fdd059c7037 100644
--- a/llvm/lib/Remarks/YAMLRemarkParser.cpp
+++ b/llvm/lib/Remarks/YAMLRemarkParser.cpp
@@ -14,6 +14,7 @@
 #include "YAMLRemarkParser.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/Remarks/RemarkParser.h"
+#include "llvm/Support/Endian.h"
 
 using namespace llvm;
 using namespace llvm::remarks;
@@ -54,6 +55,110 @@ static SourceMgr setupSM(std::string &LastErrorMessage) {
   return SM;
 }
 
+// Parse the magic number. This function returns true if this represents remark
+// metadata, false otherwise.
+static Expected<bool> parseMagic(StringRef &Buf) {
+  if (!Buf.consume_front(remarks::Magic))
+    return false;
+
+  if (Buf.size() < 1 || !Buf.consume_front(StringRef("\0", 1)))
+    return createStringError(std::errc::illegal_byte_sequence,
+                             "Expecting \\0 after magic number.");
+  return true;
+}
+
+static Expected<uint64_t> parseVersion(StringRef &Buf) {
+  if (Buf.size() < sizeof(uint64_t))
+    return createStringError(std::errc::illegal_byte_sequence,
+                             "Expecting version number.");
+
+  uint64_t Version =
+      support::endian::read<uint64_t, support::little, support::unaligned>(
+          Buf.data());
+  if (Version != remarks::Version)
+    return createStringError(
+        std::errc::illegal_byte_sequence,
+        "Mismatching remark version. Got %u, expected %u.", Version,
+        remarks::Version);
+  Buf = Buf.drop_front(sizeof(uint64_t));
+  return Version;
+}
+
+static Expected<uint64_t> parseStrTabSize(StringRef &Buf) {
+  if (Buf.size() < sizeof(uint64_t))
+    return createStringError(std::errc::illegal_byte_sequence,
+                             "Expecting string table size.");
+  uint64_t StrTabSize =
+      support::endian::read<uint64_t, support::little, support::unaligned>(
+          Buf.data());
+  Buf = Buf.drop_front(sizeof(uint64_t));
+  return StrTabSize;
+}
+
+static Expected<ParsedStringTable> parseStrTab(StringRef &Buf,
+                                               uint64_t StrTabSize) {
+  if (Buf.size() < StrTabSize)
+    return createStringError(std::errc::illegal_byte_sequence,
+                             "Expecting string table.");
+
+  // Attach the string table to the parser.
+  ParsedStringTable Result(StringRef(Buf.data(), StrTabSize));
+  Buf = Buf.drop_front(StrTabSize);
+  return Result;
+}
+
+Expected<std::unique_ptr<YAMLRemarkParser>>
+remarks::createYAMLParserFromMeta(StringRef Buf,
+                                  Optional<ParsedStringTable> StrTab) {
+  // We now have a magic number. The metadata has to be correct.
+  Expected<bool> isMeta = parseMagic(Buf);
+  if (!isMeta)
+    return isMeta.takeError();
+  // If it's not recognized as metadata, roll back.
+  std::unique_ptr<MemoryBuffer> SeparateBuf;
+  if (*isMeta) {
+    Expected<uint64_t> Version = parseVersion(Buf);
+    if (!Version)
+      return Version.takeError();
+
+    Expected<uint64_t> StrTabSize = parseStrTabSize(Buf);
+    if (!StrTabSize)
+      return StrTabSize.takeError();
+
+    // If the size of string table is not 0, try to build one.
+    if (*StrTabSize != 0) {
+      if (StrTab)
+        return createStringError(std::errc::illegal_byte_sequence,
+                                 "String table already provided.");
+      Expected<ParsedStringTable> MaybeStrTab = parseStrTab(Buf, *StrTabSize);
+      if (!MaybeStrTab)
+        return MaybeStrTab.takeError();
+      StrTab = std::move(*MaybeStrTab);
+    }
+    // If it starts with "---", there is no external file.
+    if (!Buf.startswith("---")) {
+      // At this point, we expect Buf to contain the external file path.
+      // Try to open the file and start parsing from there.
+      ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
+          MemoryBuffer::getFile(Buf);
+      if (std::error_code EC = BufferOrErr.getError())
+        return errorCodeToError(EC);
+
+      // Keep the buffer alive.
+      SeparateBuf = std::move(*BufferOrErr);
+      Buf = SeparateBuf->getBuffer();
+    }
+  }
+
+  std::unique_ptr<YAMLRemarkParser> Result =
+      StrTab
+          ? llvm::make_unique<YAMLStrTabRemarkParser>(Buf, std::move(*StrTab))
+          : llvm::make_unique<YAMLRemarkParser>(Buf);
+  if (SeparateBuf)
+    Result->SeparateBuf = std::move(SeparateBuf);
+  return std::move(Result);
+}
+
 YAMLRemarkParser::YAMLRemarkParser(StringRef Buf)
     : YAMLRemarkParser(Buf, None) {}
 
diff --git a/llvm/lib/Remarks/YAMLRemarkParser.h b/llvm/lib/Remarks/YAMLRemarkParser.h
index e8f0edc50f4..ff03abbb576 100644
--- a/llvm/lib/Remarks/YAMLRemarkParser.h
+++ b/llvm/lib/Remarks/YAMLRemarkParser.h
@@ -18,6 +18,7 @@
 #include "llvm/Remarks/Remark.h"
 #include "llvm/Remarks/RemarkParser.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/YAMLParser.h"
 #include "llvm/Support/YAMLTraits.h"
@@ -58,6 +59,9 @@ struct YAMLRemarkParser : public RemarkParser {
   yaml::Stream Stream;
   /// Iterator in the YAML stream.
   yaml::document_iterator YAMLIt;
+  /// If we parse remark metadata in separate mode, we need to open a new file
+  /// and parse that.
+  std::unique_ptr<MemoryBuffer> SeparateBuf;
 
   YAMLRemarkParser(StringRef Buf);
 
@@ -104,6 +108,11 @@ protected:
   /// Parse one value to a string.
   Expected<StringRef> parseStr(yaml::KeyValueNode &Node) override;
 };
+
+Expected<std::unique_ptr<YAMLRemarkParser>>
+createYAMLParserFromMeta(StringRef Buf,
+                         Optional<ParsedStringTable> StrTab = None);
+
 } // end namespace remarks
 } // end namespace llvm