YAML: Add support for literal block scalar I/O.

This commit gives the users of the YAML Traits I/O library 
the ability to serialize scalars using the YAML literal block 
scalar notation by allowing them to implement a specialization 
of the `BlockScalarTraits` struct for their custom types.

Reviewers: Duncan P. N. Exon Smith

Differential Revision: http://reviews.llvm.org/D9613

llvm-svn: 237404

4 files changed, 313 insertions, 1 deletions

diff --git a/llvm/docs/YamlIO.rst b/llvm/docs/YamlIO.rst
index ab8b19ac2a5..aa4bae35d1a 100644
--- a/llvm/docs/YamlIO.rst
+++ b/llvm/docs/YamlIO.rst
@@ -467,6 +467,56 @@ looks like:
       // Determine if this scalar needs quotes.
       static bool mustQuote(StringRef) { return true; }
     };
+
+Block Scalars
+-------------
+
+YAML block scalars are string literals that are represented in YAML using the
+literal block notation, just like the example shown below:
+
+.. code-block:: yaml
+
+    text: |
+      First line
+      Second line
+
+The YAML I/O library provides support for translating between YAML block scalars
+and specific C++ types by allowing you to specialize BlockScalarTraits<> on
+your data type. The library doesn't provide any built-in support for block
+scalar I/O for types like std::string and llvm::StringRef as they are already
+supported by YAML I/O and use the ordinary scalar notation by default.
+
+BlockScalarTraits specializations are very similar to the
+ScalarTraits specialization - YAML I/O will provide the native type and your
+specialization must create a temporary llvm::StringRef when writing, and
+it will also provide an llvm::StringRef that has the value of that block scalar
+and your specialization must convert that to your native data type when reading.
+An example of a custom type with an appropriate specialization of
+BlockScalarTraits is shown below:
+
+.. code-block:: c++
+
+    using llvm::yaml::BlockScalarTraits;
+    using llvm::yaml::IO;
+
+    struct MyStringType {
+      std::string Str;
+    };
+
+    template <>
+    struct BlockScalarTraits<MyStringType> {
+      static void output(const MyStringType &Value, void *Ctxt,
+                         llvm::raw_ostream &OS) {
+        OS << Value.Str;
+      }
+
+      static StringRef input(StringRef Scalar, void *Ctxt,
+                             MyStringType &Value) {
+        Value.Str = Scalar.str();
+        return StringRef();
+      }
+    };
+
     
 
 Mappings
diff --git a/llvm/include/llvm/Support/YAMLTraits.h b/llvm/include/llvm/Support/YAMLTraits.h
index b48c58e4352..3bdff209180 100644
--- a/llvm/include/llvm/Support/YAMLTraits.h
+++ b/llvm/include/llvm/Support/YAMLTraits.h
@@ -121,6 +121,35 @@ struct ScalarTraits {
 };
 
 
+/// This class should be specialized by type that requires custom conversion
+/// to/from a YAML literal block scalar. For example:
+///
+///    template <>
+///    struct BlockScalarTraits<MyType> {
+///      static void output(const MyType &Value, void*, llvm::raw_ostream &Out)
+///      {
+///        // stream out custom formatting
+///        Out << Val;
+///      }
+///      static StringRef input(StringRef Scalar, void*, MyType &Value) {
+///        // parse scalar and set `value`
+///        // return empty string on success, or error string
+///        return StringRef();
+///      }
+///    };
+template <typename T>
+struct BlockScalarTraits {
+  // Must provide:
+  //
+  // Function to write the value as a string:
+  // static void output(const T &Value, void *ctx, llvm::raw_ostream &Out);
+  //
+  // Function to convert a string to a value.  Returns the empty
+  // StringRef on success or an error string if string is malformed:
+  // static StringRef input(StringRef Scalar, void *ctxt, T &Value);
+};
+
+
 /// This class should be specialized by any type that needs to be converted
 /// to/from a YAML sequence.  For example:
 ///
@@ -224,6 +253,26 @@ public:
 };
 
 
+// Test if BlockScalarTraits<T> is defined on type T.
+template <class T>
+struct has_BlockScalarTraits
+{
+  typedef StringRef (*Signature_input)(StringRef, void *, T &);
+  typedef void (*Signature_output)(const T &, void *, llvm::raw_ostream &);
+
+  template <typename U>
+  static char test(SameType<Signature_input, &U::input> *,
+                   SameType<Signature_output, &U::output> *);
+
+  template <typename U>
+  static double test(...);
+
+public:
+  static bool const value =
+      (sizeof(test<BlockScalarTraits<T>>(nullptr, nullptr)) == 1);
+};
+
+
 // Test if MappingTraits<T> is defined on type T.
 template <class T>
 struct has_MappingTraits
@@ -410,6 +459,7 @@ struct missingTraits : public std::integral_constant<bool,
                                          !has_ScalarEnumerationTraits<T>::value
                                       && !has_ScalarBitSetTraits<T>::value
                                       && !has_ScalarTraits<T>::value
+                                      && !has_BlockScalarTraits<T>::value
                                       && !has_MappingTraits<T>::value
                                       && !has_SequenceTraits<T>::value
                                       && !has_DocumentListTraits<T>::value >  {};
@@ -462,6 +512,7 @@ public:
   virtual void endBitSetScalar() = 0;
 
   virtual void scalarString(StringRef &, bool) = 0;
+  virtual void blockScalarString(StringRef &) = 0;
 
   virtual void setError(const Twine &) = 0;
 
@@ -646,6 +697,24 @@ yamlize(IO &io, T &Val, bool) {
   }
 }
 
+template <typename T>
+typename std::enable_if<has_BlockScalarTraits<T>::value, void>::type
+yamlize(IO &YamlIO, T &Val, bool) {
+  if (YamlIO.outputting()) {
+    std::string Storage;
+    llvm::raw_string_ostream Buffer(Storage);
+    BlockScalarTraits<T>::output(Val, YamlIO.getContext(), Buffer);
+    StringRef Str = Buffer.str();
+    YamlIO.blockScalarString(Str);
+  } else {
+    StringRef Str;
+    YamlIO.blockScalarString(Str);
+    StringRef Result =
+        BlockScalarTraits<T>::input(Str, YamlIO.getContext(), Val);
+    if (!Result.empty())
+      YamlIO.setError(llvm::Twine(Result));
+  }
+}
 
 template<typename T>
 typename std::enable_if<validatedMappingTraits<T>::value, void>::type
@@ -937,6 +1006,7 @@ private:
   bool bitSetMatch(const char *, bool ) override;
   void endBitSetScalar() override;
   void scalarString(StringRef &, bool) override;
+  void blockScalarString(StringRef &) override;
   void setError(const Twine &message) override;
   bool canElideEmptySequence() override;
 
@@ -968,7 +1038,8 @@ private:
     StringRef value() const { return _value; }
 
     static inline bool classof(const HNode *n) {
-      return ScalarNode::classof(n->_node);
+      return ScalarNode::classof(n->_node) ||
+             BlockScalarNode::classof(n->_node);
     }
     static inline bool classof(const ScalarHNode *) { return true; }
   protected:
@@ -1067,6 +1138,7 @@ public:
   bool bitSetMatch(const char *, bool ) override;
   void endBitSetScalar() override;
   void scalarString(StringRef &, bool) override;
+  void blockScalarString(StringRef &) override;
   void setError(const Twine &message) override;
   bool canElideEmptySequence() override;
 public:
@@ -1208,6 +1280,16 @@ operator>>(Input &yin, T &docSeq) {
   return yin;
 }
 
+// Define non-member operator>> so that Input can stream in a block scalar.
+template <typename T>
+inline
+typename std::enable_if<has_BlockScalarTraits<T>::value, Input &>::type
+operator>>(Input &In, T &Val) {
+  if (In.setCurrentDocument())
+    yamlize(In, Val, true);
+  return In;
+}
+
 // Provide better error message about types missing a trait specialization
 template <typename T>
 inline
@@ -1263,6 +1345,20 @@ operator<<(Output &yout, T &seq) {
   return yout;
 }
 
+// Define non-member operator<< so that Output can stream out a block scalar.
+template <typename T>
+inline
+typename std::enable_if<has_BlockScalarTraits<T>::value, Output &>::type
+operator<<(Output &Out, T &Val) {
+  Out.beginDocuments();
+  if (Out.preflightDocument(0)) {
+    yamlize(Out, Val, true);
+    Out.postflightDocument();
+  }
+  Out.endDocuments();
+  return Out;
+}
+
 // Provide better error message about types missing a trait specialization
 template <typename T>
 inline
diff --git a/llvm/lib/Support/YAMLTraits.cpp b/llvm/lib/Support/YAMLTraits.cpp
index 8ef36a33280..90f34f6e232 100644
--- a/llvm/lib/Support/YAMLTraits.cpp
+++ b/llvm/lib/Support/YAMLTraits.cpp
@@ -14,6 +14,7 @@
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/Format.h"
+#include "llvm/Support/LineIterator.h"
 #include "llvm/Support/YAMLParser.h"
 #include "llvm/Support/raw_ostream.h"
 #include <cctype>
@@ -309,6 +310,8 @@ void Input::scalarString(StringRef &S, bool) {
   }
 }
 
+void Input::blockScalarString(StringRef &S) { scalarString(S, false); }
+
 void Input::setError(HNode *hnode, const Twine &message) {
   assert(hnode && "HNode must not be NULL");
   this->setError(hnode->_node, message);
@@ -331,6 +334,11 @@ std::unique_ptr<Input::HNode> Input::createHNodes(Node *N) {
       KeyStr = StringRef(Buf, Len);
     }
     return llvm::make_unique<ScalarHNode>(N, KeyStr);
+  } else if (BlockScalarNode *BSN = dyn_cast<BlockScalarNode>(N)) {
+    StringRef Value = BSN->getValue();
+    char *Buf = StringAllocator.Allocate<char>(Value.size());
+    memcpy(Buf, Value.data(), Value.size());
+    return llvm::make_unique<ScalarHNode>(N, StringRef(Buf, Value.size()));
   } else if (SequenceNode *SQ = dyn_cast<SequenceNode>(N)) {
     auto SQHNode = llvm::make_unique<SequenceHNode>(N);
     for (Node &SN : *SQ) {
@@ -609,6 +617,24 @@ void Output::scalarString(StringRef &S, bool MustQuote) {
   this->outputUpToEndOfLine("'"); // Ending single quote.
 }
 
+void Output::blockScalarString(StringRef &S) {
+  if (!StateStack.empty())
+    newLineCheck();
+  output(" |");
+  outputNewLine();
+
+  unsigned Indent = StateStack.empty() ? 1 : StateStack.size();
+
+  auto Buffer = MemoryBuffer::getMemBuffer(S, "", false);
+  for (line_iterator Lines(*Buffer, false); !Lines.is_at_end(); ++Lines) {
+    for (unsigned I = 0; I < Indent; ++I) {
+      output("  ");
+    }
+    output(*Lines);
+    outputNewLine();
+  }
+}
+
 void Output::setError(const Twine &message) {
 }
 
diff --git a/llvm/unittests/Support/YAMLIOTest.cpp b/llvm/unittests/Support/YAMLIOTest.cpp
index 0a7060cde02..7248124992b 100644
--- a/llvm/unittests/Support/YAMLIOTest.cpp
+++ b/llvm/unittests/Support/YAMLIOTest.cpp
@@ -780,6 +780,146 @@ TEST(YAMLIO, TestReadWriteMyCustomType) {
 
 
 //===----------------------------------------------------------------------===//
+//  Test BlockScalarTraits
+//===----------------------------------------------------------------------===//
+
+struct MultilineStringType {
+  std::string str;
+};
+
+struct MultilineStringTypeMap {
+  MultilineStringType name;
+  MultilineStringType description;
+  MultilineStringType ingredients;
+  MultilineStringType recipes;
+  MultilineStringType warningLabels;
+  MultilineStringType documentation;
+  int price;
+};
+
+namespace llvm {
+namespace yaml {
+  template <>
+  struct MappingTraits<MultilineStringTypeMap> {
+    static void mapping(IO &io, MultilineStringTypeMap& s) {
+      io.mapRequired("name", s.name);
+      io.mapRequired("description", s.description);
+      io.mapRequired("ingredients", s.ingredients);
+      io.mapRequired("recipes", s.recipes);
+      io.mapRequired("warningLabels", s.warningLabels);
+      io.mapRequired("documentation", s.documentation);
+      io.mapRequired("price", s.price);
+     }
+  };
+
+  // MultilineStringType is formatted as a yaml block literal scalar. A value of
+  // "Hello\nWorld" would be represented in yaml as
+  //  |
+  //    Hello
+  //    World
+  template <>
+  struct BlockScalarTraits<MultilineStringType> {
+    static void output(const MultilineStringType &value, void *ctxt,
+                       llvm::raw_ostream &out) {
+      out << value.str;
+    }
+    static StringRef input(StringRef scalar, void *ctxt,
+                           MultilineStringType &value) {
+      value.str = scalar.str();
+      return StringRef();
+    }
+  };
+}
+}
+
+LLVM_YAML_IS_DOCUMENT_LIST_VECTOR(MultilineStringType)
+
+//
+// Test writing then reading back custom values
+//
+TEST(YAMLIO, TestReadWriteMultilineStringType) {
+  std::string intermediate;
+  {
+    MultilineStringTypeMap map;
+    map.name.str = "An Item";
+    map.description.str = "Hello\nWorld";
+    map.ingredients.str = "SubItem 1\nSub Item 2\n\nSub Item 3\n";
+    map.recipes.str = "\n\nTest 1\n\n\n";
+    map.warningLabels.str = "";
+    map.documentation.str = "\n\n";
+    map.price = 350;
+
+    llvm::raw_string_ostream ostr(intermediate);
+    Output yout(ostr);
+    yout << map;
+  }
+  {
+    Input yin(intermediate);
+    MultilineStringTypeMap map2;
+    yin >> map2;
+
+    EXPECT_FALSE(yin.error());
+    EXPECT_EQ(map2.name.str, "An Item\n");
+    EXPECT_EQ(map2.description.str, "Hello\nWorld\n");
+    EXPECT_EQ(map2.ingredients.str, "SubItem 1\nSub Item 2\n\nSub Item 3\n");
+    EXPECT_EQ(map2.recipes.str, "\n\nTest 1\n");
+    EXPECT_TRUE(map2.warningLabels.str.empty());
+    EXPECT_TRUE(map2.documentation.str.empty());
+    EXPECT_EQ(map2.price, 350);
+  }
+}
+
+//
+// Test writing then reading back custom values
+//
+TEST(YAMLIO, TestReadWriteBlockScalarDocuments) {
+  std::string intermediate;
+  {
+    std::vector<MultilineStringType> documents;
+    MultilineStringType doc;
+    doc.str = "Hello\nWorld";
+    documents.push_back(doc);
+
+    llvm::raw_string_ostream ostr(intermediate);
+    Output yout(ostr);
+    yout << documents;
+
+    // Verify that the block scalar header was written out on the same line
+    // as the document marker.
+    EXPECT_NE(llvm::StringRef::npos, llvm::StringRef(ostr.str()).find("--- |"));
+  }
+  {
+    Input yin(intermediate);
+    std::vector<MultilineStringType> documents2;
+    yin >> documents2;
+
+    EXPECT_FALSE(yin.error());
+    EXPECT_EQ(documents2.size(), size_t(1));
+    EXPECT_EQ(documents2[0].str, "Hello\nWorld\n");
+  }
+}
+
+TEST(YAMLIO, TestReadWriteBlockScalarValue) {
+  std::string intermediate;
+  {
+    MultilineStringType doc;
+    doc.str = "Just a block\nscalar doc";
+
+    llvm::raw_string_ostream ostr(intermediate);
+    Output yout(ostr);
+    yout << doc;
+  }
+  {
+    Input yin(intermediate);
+    MultilineStringType doc;
+    yin >> doc;
+
+    EXPECT_FALSE(yin.error());
+    EXPECT_EQ(doc.str, "Just a block\nscalar doc\n");
+  }
+}
+
+//===----------------------------------------------------------------------===//
 //  Test flow sequences
 //===----------------------------------------------------------------------===//