3 files changed, 58 insertions, 0 deletions

diff --git a/include/llvm/Support/YAMLParser.h b/include/llvm/Support/YAMLParser.h
index 12958fa..e3f4f00 100644
--- a/include/llvm/Support/YAMLParser.h
+++ b/include/llvm/Support/YAMLParser.h
@@ -77,7 +77,11 @@ std::string escape(StringRef Input);
 ///        documents.
 class Stream {
 public:
+  /// @brief This keeps a reference to the string referenced by \p Input.
   Stream(StringRef Input, SourceMgr &);
+
+  /// @brief This takes ownership of \p InputBuffer.
+  Stream(MemoryBuffer *InputBuffer, SourceMgr &);
   ~Stream();
 
   document_iterator begin();
diff --git a/lib/Support/YAMLParser.cpp b/lib/Support/YAMLParser.cpp
index 34df636..2ae33f5 100644
--- a/lib/Support/YAMLParser.cpp
+++ b/lib/Support/YAMLParser.cpp
@@ -252,6 +252,7 @@ namespace yaml {
 class Scanner {
 public:
   Scanner(const StringRef Input, SourceMgr &SM);
+  Scanner(MemoryBuffer *Buffer, SourceMgr &SM_);
 
   /// @brief Parse the next token and return it without popping it.
   Token &peekNext();
@@ -708,6 +709,21 @@ Scanner::Scanner(StringRef Input, SourceMgr &sm)
   End = InputBuffer->getBufferEnd();
 }
 
+Scanner::Scanner(MemoryBuffer *Buffer, SourceMgr &SM_)
+  : SM(SM_)
+  , InputBuffer(Buffer)
+  , Current(InputBuffer->getBufferStart())
+  , End(InputBuffer->getBufferEnd())
+  , Indent(-1)
+  , Column(0)
+  , Line(0)
+  , FlowLevel(0)
+  , IsStartOfStream(true)
+  , IsSimpleKeyAllowed(true)
+  , Failed(false) {
+    SM.AddNewSourceBuffer(InputBuffer, SMLoc());
+}
+
 Token &Scanner::peekNext() {
   // If the current token is a possible simple key, keep parsing until we
   // can confirm.
@@ -1532,6 +1548,10 @@ Stream::Stream(StringRef Input, SourceMgr &SM)
   : scanner(new Scanner(Input, SM))
   , CurrentDoc(0) {}
 
+Stream::Stream(MemoryBuffer *InputBuffer, SourceMgr &SM)
+  : scanner(new Scanner(InputBuffer, SM))
+  , CurrentDoc(0) {}
+
 Stream::~Stream() {}
 
 bool Stream::failed() { return scanner->failed(); }
diff --git a/unittests/Support/YAMLParserTest.cpp b/unittests/Support/YAMLParserTest.cpp
index 480a573..e983935 100644
--- a/unittests/Support/YAMLParserTest.cpp
+++ b/unittests/Support/YAMLParserTest.cpp
@@ -10,6 +10,7 @@
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/Casting.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/YAMLParser.h"
 #include "gtest/gtest.h"
@@ -21,6 +22,12 @@ static void SuppressDiagnosticsOutput(const SMDiagnostic &, void *) {
   // to reduce noise in unit test runs.
 }
 
+// Assumes Ctx is an SMDiagnostic where Diag can be stored.
+static void CollectDiagnosticsOutput(const SMDiagnostic &Diag, void *Ctx) {
+  SMDiagnostic* DiagOut = static_cast<SMDiagnostic*>(Ctx);
+  *DiagOut = Diag;
+}
+
 // Checks that the given input gives a parse error. Makes sure that an error
 // text is available and the parse fails.
 static void ExpectParseError(StringRef Message, StringRef Input) {
@@ -182,4 +189,31 @@ TEST(YAMLParser, WorksWithIteratorAlgorithms) {
   EXPECT_EQ(6, std::distance(Array->begin(), Array->end()));
 }
 
+TEST(YAMLParser, DefaultDiagnosticFilename) {
+  SourceMgr SM;
+
+  SMDiagnostic GeneratedDiag;
+  SM.setDiagHandler(CollectDiagnosticsOutput, &GeneratedDiag);
+
+  // When we construct a YAML stream over an unnamed string,
+  // the filename is hard-coded as "YAML".
+  yaml::Stream UnnamedStream("[]", SM);
+  UnnamedStream.printError(UnnamedStream.begin()->getRoot(), "Hello, World!");
+  EXPECT_EQ("YAML", GeneratedDiag.getFilename());
+}
+
+TEST(YAMLParser, DiagnosticFilenameFromBufferID) {
+  SourceMgr SM;
+
+  SMDiagnostic GeneratedDiag;
+  SM.setDiagHandler(CollectDiagnosticsOutput, &GeneratedDiag);
+
+  // When we construct a YAML stream over a named buffer,
+  // we get its ID as filename in diagnostics.
+  MemoryBuffer* Buffer = MemoryBuffer::getMemBuffer("[]", "buffername.yaml");
+  yaml::Stream Stream(Buffer, SM);
+  Stream.printError(Stream.begin()->getRoot(), "Hello, World!");
+  EXPECT_EQ("buffername.yaml", GeneratedDiag.getFilename());
+}
+
 } // end namespace llvm