1 files changed, 37 insertions, 21 deletions

diff --git a/lib/Support/YAMLTraits.cpp b/lib/Support/YAMLTraits.cpp
index b9bbee7883c6..09eb36943de9 100644
--- a/lib/Support/YAMLTraits.cpp
+++ b/lib/Support/YAMLTraits.cpp
@@ -1,9 +1,8 @@
 //===- lib/Support/YAMLTraits.cpp -----------------------------------------===//
 //
-//                             The LLVM Linker
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
@@ -114,6 +113,11 @@ const Node *Input::getCurrentNode() const {
 }
 
 bool Input::mapTag(StringRef Tag, bool Default) {
+  // CurrentNode can be null if setCurrentDocument() was unable to
+  // parse the document because it was invalid or empty.
+  if (!CurrentNode)
+    return false;
+
   std::string foundTag = CurrentNode->_node->getVerbatimTag();
   if (foundTag.empty()) {
     // If no tag found and 'Tag' is the default, say it was found.
@@ -442,7 +446,8 @@ bool Output::outputting() {
 
 void Output::beginMapping() {
   StateStack.push_back(inMapFirstKey);
-  NeedsNewLine = true;
+  PaddingBeforeContainer = Padding;
+  Padding = "\n";
 }
 
 bool Output::mapTag(StringRef Tag, bool Use) {
@@ -470,7 +475,7 @@ bool Output::mapTag(StringRef Tag, bool Use) {
       }
       // Tags inside maps in sequences should act as keys in the map from a
       // formatting perspective, so we always want a newline in a sequence.
-      NeedsNewLine = true;
+      Padding = "\n";
     }
   }
   return Use;
@@ -478,8 +483,12 @@ bool Output::mapTag(StringRef Tag, bool Use) {
 
 void Output::endMapping() {
   // If we did not map anything, we should explicitly emit an empty map
-  if (StateStack.back() == inMapFirstKey)
+  if (StateStack.back() == inMapFirstKey) {
+    Padding = PaddingBeforeContainer;
+    newLineCheck();
     output("{}");
+    Padding = "\n";
+  }
   StateStack.pop_back();
 }
 
@@ -544,14 +553,19 @@ void Output::endDocuments() {
 
 unsigned Output::beginSequence() {
   StateStack.push_back(inSeqFirstElement);
-  NeedsNewLine = true;
+  PaddingBeforeContainer = Padding;
+  Padding = "\n";
   return 0;
 }
 
 void Output::endSequence() {
   // If we did not emit anything, we should explicitly emit an empty sequence
-  if (StateStack.back() == inSeqFirstElement)
+  if (StateStack.back() == inSeqFirstElement) {
+    Padding = PaddingBeforeContainer;
+    newLineCheck();
     output("[]");
+    Padding = "\n";
+  }
   StateStack.pop_back();
 }
 
@@ -661,11 +675,6 @@ void Output::scalarString(StringRef &S, QuotingType MustQuote) {
     return;
   }
 
-  unsigned i = 0;
-  unsigned j = 0;
-  unsigned End = S.size();
-  const char *Base = S.data();
-
   const char *const Quote = MustQuote == QuotingType::Single ? "'" : "\"";
   output(Quote); // Starting quote.
 
@@ -673,11 +682,16 @@ void Output::scalarString(StringRef &S, QuotingType MustQuote) {
   // present, and will be escaped using a variety of unicode-scalar and special short-form
   // escapes. This is handled in yaml::escape.
   if (MustQuote == QuotingType::Double) {
-    output(yaml::escape(Base, /* EscapePrintable= */ false));
+    output(yaml::escape(S, /* EscapePrintable= */ false));
     outputUpToEndOfLine(Quote);
     return;
   }
 
+  unsigned i = 0;
+  unsigned j = 0;
+  unsigned End = S.size();
+  const char *Base = S.data();
+
   // When using single-quoted strings, any single quote ' must be doubled to be escaped.
   while (j < End) {
     if (S[j] == '\'') {                    // Escape quotes.
@@ -742,7 +756,7 @@ void Output::outputUpToEndOfLine(StringRef s) {
   output(s);
   if (StateStack.empty() || (!inFlowSeqAnyElement(StateStack.back()) &&
                              !inFlowMapAnyKey(StateStack.back())))
-    NeedsNewLine = true;
+    Padding = "\n";
 }
 
 void Output::outputNewLine() {
@@ -755,11 +769,13 @@ void Output::outputNewLine() {
 //
 
 void Output::newLineCheck() {
-  if (!NeedsNewLine)
+  if (Padding != "\n") {
+    output(Padding);
+    Padding = {};
     return;
-  NeedsNewLine = false;
-
+  }
   outputNewLine();
+  Padding = {};
 
   if (StateStack.size() == 0)
     return;
@@ -793,9 +809,9 @@ void Output::paddedKey(StringRef key) {
   output(":");
   const char *spaces = "                ";
   if (key.size() < strlen(spaces))
-    output(&spaces[key.size()]);
+    Padding = &spaces[key.size()];
   else
-    output(" ");
+    Padding = " ";
 }
 
 void Output::flowKey(StringRef Key) {