[Support] Teach YAMLIO about polymorphic types

[scott.linder]

Summary

Add support for "polymorphic" types to YAMLIO.

PolymorphicTraits can dynamically switch between other traits (Scalar, Map, or Sequence). When inputting, the PolymorphicTraits type is told which type to become, and when outputting the PolymorphicTraits type is asked which type it currently is.

Also adds support for TaggedScalarTraits to allow dynamically differentiating between multiple scalar types using YAML tags.

Some other minor changes, like serializing empty maps as "{}" and empty sequences as "[]", so that types are preserved when round-tripping PolymorphicTraits.

This patch is part of a series required for AMDGPU to support it's new MessagePack metadata.

[scott.linder]

Rebase and ping

[scott.linder]

Rebase and ping

[scott.linder]

In cleaning up this patch I noticed our YAML parser does not handle an explicit document containing only a block-level plain scalar, e.g.:

---
plain scalar
...

The YAMLTraits do not implement operator<< for the non-block scalar, presumably for this reason. As the new polymorphic trait can't determine if the value is scalar at compile-time, I added an assert to catch it dynamically for now. In the future modifying the parser to support this might be useful.

[scott.linder]

Default the test case destructor rather than give it an empty body.

[scott.linder]

Ping

[zturner]

I'm sorry this has taken so long to get a response. I'm not really sure who the best person to look at this is, but I don't think it's me unfortunately. @chandlerc , can you help find a reviewer for this?

[scott.linder]

Ping

[dblaikie]

I really don't know much about the YAMLIO code, but here's some high level review...

include/llvm/Support/YAMLTraits.h

301

Perhaps this could/should be a declaration rather than a definition? (since it doesn't have any members anyway)

552

You can drop this since it's the default (all members are public in a struct)

lib/Support/YAMLTraits.cpp

455

Maybe have a temp for the StateStack element:

bool SequenceElement = StateStack.size() > 1;
if (SequenceElement) {
  auto &E = StateStack[StateStack.size() - 2];
  Sequence = E == inXXX || E = inYYY ....;
}

Or a quick utility function written somewhere , where you pass in StateStack[StateStack.size() - 2] (thus naming it, etc).

Some of the other conditions further down in the review seem similarly a bit verbose and repetitive.

716–718

Any particular reason for the 'this->' qualifiers? They're probably not needed here, I think?

unittests/Support/YAMLIOTest.cpp

2663

I probably wouldn't typedef this - it obscures the fact that this has memory ownership semantics.

2686

You can drop the "public" keywords here - they're implied/the default for structs. (similarly below)

2704–2705

Drop "else after return" (similarly below/etc)

2714

Probably prefer N = llvm::make_unique<Scalar>(); here and elsewhere/in similar contexts.

2805

auto node = llvm::make_unique<Scalar>(true);

(& similarly elsewhere)

[scott.linder]

I can submit a separate patch to clean up things like making the empty definitions into declarations, removing redundant public: qualifiers, and removing redundant references to this, but I don't want to mix those changes into this patch.

I will update this patch with the other changes, as well as full context.

include/llvm/Support/YAMLTraits.h

301

I agree, but I am just following the convention in the file, which is an empty definition (see e.g. MappingContextTraits). I mistakenly uploaded the patch without context, but I will include it in the next revision.

552

Again I'm trying to follow a convention found throughout the file for has_ structs, although all of these could be updated to drop the redundant qualifier.

[scott.linder]

Address feedback

[scott.linder]

Cleanup patch is at https://reviews.llvm.org/D53088

[scott.linder]

Rebase and remove redundant qualifiers

[scott.linder]

Ping

[scott.linder]

Ping

[nhaehnle]

One small nitpick, apart from that LGTM.

include/llvm/Support/YAMLTraits.h

205

Value instead of Val, I assume.

[scott.linder]

Address feedback

[ruiu]

This is not a comment to your code (so please commit if you already got an LGTM), but I have an experience of using YAMLIO in lld and found that the library is hard to use, at least for me, perhaps due to too much uses of template. YAML is intended to be designed as a language that is easy to parse, but if you use YAMLIO, it is sometimes very tricky to write code to parse/generate a YAML file. I feel like there's a better way of doing it, and the YAMLIO library could be improved drastically. Not sure if it is worth doing, given that we don't use YAMLIO that much in LLVM, but I'd like to share my thoughts...

[scott.linder]

The heavy use of templates does lead to some awkward usage, but it can also make some common use-cases very natural. I agree there is probably room to improve, but it is orthogonal to this patch.

If there are no objections I will be applying this patch (and the updates to some Clang tests at https://reviews.llvm.org/D51223) today.

include/llvm/Support/YAMLTraits.h

 namespace llvm {
 namespace yaml {
 
+enum class NodeKind : uint8_t {
+  Scalar,
+  Map,
+  Sequence,
+};
+
 struct EmptyContext {};
 
 /// This class should be specialized by any type that needs to be converted

   // Must provide:
   //
   // Function to write the value as a string:
-  //static void output(const T &value, void *ctxt, llvm::raw_ostream &out);
+  // static void output(const T &value, void *ctxt, 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);
+  // static StringRef input(StringRef scalar, void *ctxt, T &value);
   //
   // Function to determine if the value should be quoted.
-  //static QuotingType mustQuote(StringRef);
+  // static QuotingType mustQuote(StringRef);
 };
 
 /// This class should be specialized by type that requires custom conversion

 ///      static void output(const MyType &Value, void*, llvm::raw_ostream &Out)
 ///      {
 ///        // stream out custom formatting
-///        Out << Val;
+///        Out << Value;
 ///      }
 ///      static StringRef input(StringRef Scalar, void*, MyType &Value) {
 ///        // parse scalar and set `value`

   // 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);
+  //
+  // Optional:
+  // static StringRef inputTag(T &Val, std::string Tag)
+  // static void outputTag(const T &Val, raw_ostream &Out)
+};
+
+/// This class should be specialized by type that requires custom conversion
+/// to/from a YAML scalar with optional tags. For example:
+///
+///    template <>
+///    struct TaggedScalarTraits<MyType> {
+///      static void output(const MyType &Value, void*, llvm::raw_ostream
+///      &ScalarOut, llvm::raw_ostream &TagOut)
+///      {
+///        // stream out custom formatting including optional Tag
+///        Out << Value;
+///      }
+///      static StringRef input(StringRef Scalar, StringRef Tag, void*, MyType
+///      &Value) {
+///        // parse scalar and set `value`
+///        // return empty string on success, or error string
+///        return StringRef();
+///      }
+///      static QuotingType mustQuote(const MyType &Value, StringRef) {
+///        return QuotingType::Single;
+///      }
+///    };
+template <typename T> struct TaggedScalarTraits {
+  // Must provide:
+  //
+  // Function to write the value and tag as strings:
+  // static void output(const T &Value, void *ctx, llvm::raw_ostream &ScalarOut,
+  // llvm::raw_ostream &TagOut);
+  //
+  // 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, StringRef Tag, void *ctxt, T
+  // &Value);
+  //
+  // Function to determine if the value should be quoted.
+  // static QuotingType mustQuote(const T &Value, StringRef Scalar);
 };
 
 /// This class should be specialized by any type that needs to be converted

   // static void output(IO &io, T &elem);
 };
 
+/// This class should be specialized by any type that can be represented as
+/// a scalar, map, or sequence, decided dynamically. For example:
+///
+///    typedef std::unique_ptr<MyBase> MyPoly;
+///
+///    template<>
+///    struct PolymorphicTraits<MyPoly> {
+///      static NodeKind getKind(const MyPoly &poly) {
+///        return poly->getKind();
+///      }
+///      static MyScalar& getAsScalar(MyPoly &poly) {
+///        if (!poly || !isa<MyScalar>(poly))
+///          poly.reset(new MyScalar());
+///        return *cast<MyScalar>(poly.get());
+///      }
+///      // ...
+///    };
+template <typename T> struct PolymorphicTraits {
+  // Must provide:
+  // static NodeKind getKind(const T &poly);
+  // static scalar_type &getAsScalar(T &poly);
+  // static map_type &getAsMap(T &poly);
+  // static sequence_type &getAsSequence(T &poly);
+};
+
 // Only used for better diagnostics of missing traits
 template <typename T>
 struct MissingTrait;

       (sizeof(test<BlockScalarTraits<T>>(nullptr, nullptr)) == 1);
 };
 
+// Test if TaggedScalarTraits<T> is defined on type T.
+template <class T> struct has_TaggedScalarTraits {
+  using Signature_input = StringRef (*)(StringRef, StringRef, void *, T &);
+  using Signature_output = void (*)(const T &, void *, raw_ostream &,
+                                    raw_ostream &);
+  using Signature_mustQuote = QuotingType (*)(const T &, StringRef);
+
+  template <typename U>
+  static char test(SameType<Signature_input, &U::input> *,
+                   SameType<Signature_output, &U::output> *,
+                   SameType<Signature_mustQuote, &U::mustQuote> *);
+
+  template <typename U> static double test(...);
+
+  static bool const value =
+      (sizeof(test<TaggedScalarTraits<T>>(nullptr, nullptr, nullptr)) == 1);
+};
+
 // Test if MappingContextTraits<T> is defined on type T.
 template <class T, class Context> struct has_MappingTraits {
   using Signature_mapping = void (*)(class IO &, T &, Context &);

   static bool const value = (sizeof(test<DocumentListTraits<T>>(nullptr))==1);
 };
 
+template <class T> struct has_PolymorphicTraits {
+  using Signature_getKind = NodeKind (*)(const T &);
+
+  template <typename U>
+  static char test(SameType<Signature_getKind, &U::getKind> *);
+
+  template <typename U> static double test(...);
+
+  static bool const value = (sizeof(test<PolymorphicTraits<T>>(nullptr)) == 1);
+};
+
 inline bool isNumeric(StringRef S) {
   const static auto skipDigits = [](StringRef Input) {
     return Input.drop_front(

                                         !has_ScalarBitSetTraits<T>::value &&
                                         !has_ScalarTraits<T>::value &&
                                         !has_BlockScalarTraits<T>::value &&
+                                        !has_TaggedScalarTraits<T>::value &&
                                         !has_MappingTraits<T, Context>::value &&
                                         !has_SequenceTraits<T>::value &&
                                         !has_CustomMappingTraits<T>::value &&
-                                        !has_DocumentListTraits<T>::value> {};
+                                        !has_DocumentListTraits<T>::value &&
+                                        !has_PolymorphicTraits<T>::value> {};
 
 template <typename T, typename Context>
 struct validatedMappingTraits

 
   virtual void scalarString(StringRef &, QuotingType) = 0;
   virtual void blockScalarString(StringRef &) = 0;
+  virtual void scalarTag(std::string &) = 0;
+
+  virtual NodeKind getNodeKind() = 0;
 
   virtual void setError(const Twine &) = 0;
 

   }
 }
 
+template <typename T>
+typename std::enable_if<has_TaggedScalarTraits<T>::value, void>::type
+yamlize(IO &io, T &Val, bool, EmptyContext &Ctx) {
+  if (io.outputting()) {
+    std::string ScalarStorage, TagStorage;
+    raw_string_ostream ScalarBuffer(ScalarStorage), TagBuffer(TagStorage);
+    TaggedScalarTraits<T>::output(Val, io.getContext(), ScalarBuffer,
+                                  TagBuffer);
+    io.scalarTag(TagBuffer.str());
+    StringRef ScalarStr = ScalarBuffer.str();
+    io.scalarString(ScalarStr,
+                    TaggedScalarTraits<T>::mustQuote(Val, ScalarStr));
+  } else {
+    std::string Tag;
+    io.scalarTag(Tag);
+    StringRef Str;
+    io.scalarString(Str, QuotingType::None);
+    StringRef Result =
+        TaggedScalarTraits<T>::input(Str, Tag, io.getContext(), Val);
+    if (!Result.empty()) {
+      io.setError(Twine(Result));
+    }
+  }
+}
+
 template <typename T, typename Context>
 typename std::enable_if<validatedMappingTraits<T, Context>::value, void>::type
 yamlize(IO &io, T &Val, bool, Context &Ctx) {

 }
 
 template <typename T>
+typename std::enable_if<has_PolymorphicTraits<T>::value, void>::type
+yamlize(IO &io, T &Val, bool, EmptyContext &Ctx) {
+  switch (io.outputting() ? PolymorphicTraits<T>::getKind(Val)
+                          : io.getNodeKind()) {
+  case NodeKind::Scalar:
+    return yamlize(io, PolymorphicTraits<T>::getAsScalar(Val), true, Ctx);
+  case NodeKind::Map:
+    return yamlize(io, PolymorphicTraits<T>::getAsMap(Val), true, Ctx);
+  case NodeKind::Sequence:
+    return yamlize(io, PolymorphicTraits<T>::getAsSequence(Val), true, Ctx);
+  }
+}
+
+template <typename T>
 typename std::enable_if<missingTraits<T, EmptyContext>::value, void>::type
 yamlize(IO &io, T &Val, bool, EmptyContext &Ctx) {
   char missing_yaml_trait_for_type[sizeof(MissingTrait<T>)];

   void endBitSetScalar() override;
   void scalarString(StringRef &, QuotingType) override;
   void blockScalarString(StringRef &) override;
+  void scalarTag(std::string &) override;
+  NodeKind getNodeKind() override;
   void setError(const Twine &message) override;
   bool canElideEmptySequence() override;
 

   void endBitSetScalar() override;
   void scalarString(StringRef &, QuotingType) override;
   void blockScalarString(StringRef &) override;
+  void scalarTag(std::string &) override;
+  NodeKind getNodeKind() override;
   void setError(const Twine &message) override;
   bool canElideEmptySequence() override;
 

   void flowKey(StringRef Key);
 
   enum InState {
-    inSeq,
-    inFlowSeq,
+    inSeqFirstElement,
+    inSeqOtherElement,
+    inFlowSeqFirstElement,
+    inFlowSeqOtherElement,
     inMapFirstKey,
     inMapOtherKey,
     inFlowMapFirstKey,
     inFlowMapOtherKey
   };
 
+  static bool inSeqAnyElement(InState State);
+  static bool inFlowSeqAnyElement(InState State);
+  static bool inMapAnyKey(InState State);
+  static bool inFlowMapAnyKey(InState State);
+
   raw_ostream &Out;
   int WrapColumn;
   SmallVector<InState, 8> StateStack;

   return In;
 }
 
+// Define non-member operator>> so that Input can stream in a polymorphic type.
+template <typename T>
+inline typename std::enable_if<has_PolymorphicTraits<T>::value, Input &>::type
+operator>>(Input &In, T &Val) {
+  EmptyContext Ctx;
+  if (In.setCurrentDocument())
+    yamlize(In, Val, true, Ctx);
+  return In;
+}
+
 // Provide better error message about types missing a trait specialization
 template <typename T>
 inline typename std::enable_if<missingTraits<T, EmptyContext>::value,

   return Out;
 }
 
+// Define non-member operator<< so that Output can stream out a polymorphic
+// type.
+template <typename T>
+inline typename std::enable_if<has_PolymorphicTraits<T>::value, Output &>::type
+operator<<(Output &Out, T &Val) {
+  EmptyContext Ctx;
+  Out.beginDocuments();
+  if (Out.preflightDocument(0)) {
+    // FIXME: The parser does not support explicit documents terminated with a
+    // plain scalar; the end-marker is included as part of the scalar token.
+    assert(PolymorphicTraits<T>::getKind(Val) != NodeKind::Scalar && "plain scalar documents are not supported");
+    yamlize(Out, Val, true, Ctx);
+    Out.postflightDocument();
+  }
+  Out.endDocuments();
+  return Out;
+}
+
 // Provide better error message about types missing a trait specialization
 template <typename T>
 inline typename std::enable_if<missingTraits<T, EmptyContext>::value,

lib/Support/YAMLTraits.cpp

 
 void Input::blockScalarString(StringRef &S) { scalarString(S, QuotingType::None); }
 
+void Input::scalarTag(std::string &Tag) {
+  Tag = CurrentNode->_node->getVerbatimTag();
+}
+
 void Input::setError(HNode *hnode, const Twine &message) {
   assert(hnode && "HNode must not be NULL");
   setError(hnode->_node, message);
 }
 
+NodeKind Input::getNodeKind() {
+  if (isa<ScalarHNode>(CurrentNode))
+    return NodeKind::Scalar;
+  else if (isa<MapHNode>(CurrentNode))
+    return NodeKind::Map;
+  else if (isa<SequenceHNode>(CurrentNode))
+    return NodeKind::Sequence;
+  llvm_unreachable("Unsupported node kind");
+}
+
 void Input::setError(Node *node, const Twine &message) {
   Strm->printError(node, message);
   EC = make_error_code(errc::invalid_argument);

     // If this tag is being written inside a sequence we should write the start
     // of the sequence before writing the tag, otherwise the tag won't be
     // attached to the element in the sequence, but rather the sequence itself.
-    bool SequenceElement =
-        StateStack.size() > 1 && (StateStack[StateStack.size() - 2] == inSeq ||
-          StateStack[StateStack.size() - 2] == inFlowSeq);
+    bool SequenceElement = false;
+    if (StateStack.size() > 1) {
+      auto &E = StateStack[StateStack.size() - 2];
+      SequenceElement = inSeqAnyElement(E) || inFlowSeqAnyElement(E);
+    }
     if (SequenceElement && StateStack.back() == inMapFirstKey) {
       newLineCheck();
     } else {

 }
 
 void Output::endMapping() {
+  // If we did not map anything, we should explicitly emit an empty map
+  if (StateStack.back() == inMapFirstKey)
+    output("{}");
   StateStack.pop_back();
 }
 

 }
 
 unsigned Output::beginSequence() {
-  StateStack.push_back(inSeq);
+  StateStack.push_back(inSeqFirstElement);
   NeedsNewLine = true;
   return 0;
 }
 
 void Output::endSequence() {
+  // If we did not emit anything, we should explicitly emit an empty sequence
+  if (StateStack.back() == inSeqFirstElement)
+    output("[]");
   StateStack.pop_back();
 }
 

 }
 
 void Output::postflightElement(void *) {
+  if (StateStack.back() == inSeqFirstElement) {
+    StateStack.pop_back();
+    StateStack.push_back(inSeqOtherElement);
+  } else if (StateStack.back() == inFlowSeqFirstElement) {
+    StateStack.pop_back();
+    StateStack.push_back(inFlowSeqOtherElement);
+  }
 }
 
 unsigned Output::beginFlowSequence() {
-  StateStack.push_back(inFlowSeq);
+  StateStack.push_back(inFlowSeqFirstElement);
   newLineCheck();
   ColumnAtFlowStart = Column;
   output("[ ");

   }
 }
 
+void Output::scalarTag(std::string &Tag) {
+  if (Tag.empty())
+    return;
+  newLineCheck();
+  output(Tag);
+  output(" ");
+}
+
 void Output::setError(const Twine &message) {
 }
 

     return true;
   if (StateStack.back() != inMapFirstKey)
     return true;
-  return (StateStack[StateStack.size()-2] != inSeq);
+  return !inSeqAnyElement(StateStack[StateStack.size() - 2]);
 }
 
 void Output::output(StringRef s) {

 
 void Output::outputUpToEndOfLine(StringRef s) {
   output(s);
-  if (StateStack.empty() || (StateStack.back() != inFlowSeq &&
-                             StateStack.back() != inFlowMapFirstKey &&
-                             StateStack.back() != inFlowMapOtherKey))
+  if (StateStack.empty() || (!inFlowSeqAnyElement(StateStack.back()) &&
+                             !inFlowMapAnyKey(StateStack.back())))
     NeedsNewLine = true;
 }
 

 
   outputNewLine();
 
-  assert(StateStack.size() > 0);
+  if (StateStack.size() == 0)
+    return;
+
   unsigned Indent = StateStack.size() - 1;
   bool OutputDash = false;
 
-  if (StateStack.back() == inSeq) {
+  if (StateStack.back() == inSeqFirstElement ||
+      StateStack.back() == inSeqOtherElement) {
     OutputDash = true;
-  } else if ((StateStack.size() > 1) && ((StateStack.back() == inMapFirstKey) ||
-             (StateStack.back() == inFlowSeq) ||
-             (StateStack.back() == inFlowMapFirstKey)) &&
-             (StateStack[StateStack.size() - 2] == inSeq)) {
+  } else if ((StateStack.size() > 1) &&
+             ((StateStack.back() == inMapFirstKey) ||
+              inFlowSeqAnyElement(StateStack.back()) ||
+              (StateStack.back() == inFlowMapFirstKey)) &&
+             inSeqAnyElement(StateStack[StateStack.size() - 2])) {
     --Indent;
     OutputDash = true;
   }

   output(": ");
 }
 
+NodeKind Output::getNodeKind() { report_fatal_error("invalid call"); }
+
+bool Output::inSeqAnyElement(InState State) {
+  return State == inSeqFirstElement || State == inSeqOtherElement;
+}
+
+bool Output::inFlowSeqAnyElement(InState State) {
+  return State == inFlowSeqFirstElement || State == inFlowSeqOtherElement;
+}
+
+bool Output::inMapAnyKey(InState State) {
+  return State == inMapFirstKey || State == inMapOtherKey;
+}
+
+bool Output::inFlowMapAnyKey(InState State) {
+  return State == inFlowMapFirstKey || State == inFlowMapOtherKey;
+}
+
 //===----------------------------------------------------------------------===//
 //  traits for built-in types
 //===----------------------------------------------------------------------===//

unittests/Support/YAMLIOTest.cpp

 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/Casting.h"

   EXPECT_FALSE(isNumeric("-inf"));
   EXPECT_FALSE(isNumeric("1,230.15"));
 }
+
+//===----------------------------------------------------------------------===//
+//  Test PolymorphicTraits and TaggedScalarTraits
+//===----------------------------------------------------------------------===//
+
+struct Poly {
+  enum NodeKind {
+    NK_Scalar,
+    NK_Seq,
+    NK_Map,
+  } Kind;
+
+  Poly(NodeKind Kind) : Kind(Kind) {}
+
+  virtual ~Poly() = default;
+
+  NodeKind getKind() const { return Kind; }
+};
+
+struct Scalar : Poly {
+  enum ScalarKind {
+    SK_Unknown,
+    SK_Double,
+    SK_Bool,
+  } SKind;
+
+  union {
+    double DoubleValue;
+    bool BoolValue;
+  };
+
+  Scalar() : Poly(NK_Scalar), SKind(SK_Unknown) {}
+  Scalar(double DoubleValue)
+      : Poly(NK_Scalar), SKind(SK_Double), DoubleValue(DoubleValue) {}
+  Scalar(bool BoolValue)
+      : Poly(NK_Scalar), SKind(SK_Bool), BoolValue(BoolValue) {}
+
+  static bool classof(const Poly *N) { return N->getKind() == NK_Scalar; }
+};
+
+struct Seq : Poly, std::vector<std::unique_ptr<Poly>> {
+  Seq() : Poly(NK_Seq) {}
+
+  static bool classof(const Poly *N) { return N->getKind() == NK_Seq; }
+};
+
+struct Map : Poly, llvm::StringMap<std::unique_ptr<Poly>> {
+  Map() : Poly(NK_Map) {}
+
+  static bool classof(const Poly *N) { return N->getKind() == NK_Map; }
+};
+
+namespace llvm {
+namespace yaml {
+
+template <> struct PolymorphicTraits<std::unique_ptr<Poly>> {
+  static NodeKind getKind(const std::unique_ptr<Poly> &N) {
+    if (isa<Scalar>(*N))
+      return NodeKind::Scalar;
+    if (isa<Seq>(*N))
+      return NodeKind::Sequence;
+    if (isa<Map>(*N))
+      return NodeKind::Map;
+    llvm_unreachable("unsupported node type");
+  }
+
+  static Scalar &getAsScalar(std::unique_ptr<Poly> &N) {
+    if (!N || !isa<Scalar>(*N))
+      N = llvm::make_unique<Scalar>();
+    return *cast<Scalar>(N.get());
+  }
+
+  static Seq &getAsSequence(std::unique_ptr<Poly> &N) {
+    if (!N || !isa<Seq>(*N))
+      N = llvm::make_unique<Seq>();
+    return *cast<Seq>(N.get());
+  }
+
+  static Map &getAsMap(std::unique_ptr<Poly> &N) {
+    if (!N || !isa<Map>(*N))
+      N = llvm::make_unique<Map>();
+    return *cast<Map>(N.get());
+  }
+};
+
+template <> struct TaggedScalarTraits<Scalar> {
+  static void output(const Scalar &S, void *Ctxt, raw_ostream &ScalarOS,
+                     raw_ostream &TagOS) {
+    switch (S.SKind) {
+    case Scalar::SK_Unknown:
+      report_fatal_error("output unknown scalar");
+      break;
+    case Scalar::SK_Double:
+      TagOS << "!double";
+      ScalarTraits<double>::output(S.DoubleValue, Ctxt, ScalarOS);
+      break;
+    case Scalar::SK_Bool:
+      TagOS << "!bool";
+      ScalarTraits<bool>::output(S.BoolValue, Ctxt, ScalarOS);
+      break;
+    }
+  }
+
+  static StringRef input(StringRef ScalarStr, StringRef Tag, void *Ctxt,
+                         Scalar &S) {
+    S.SKind = StringSwitch<Scalar::ScalarKind>(Tag)
+                  .Case("!double", Scalar::SK_Double)
+                  .Case("!bool", Scalar::SK_Bool)
+                  .Default(Scalar::SK_Unknown);
+    switch (S.SKind) {
+    case Scalar::SK_Unknown:
+      return StringRef("unknown scalar tag");
+    case Scalar::SK_Double:
+      return ScalarTraits<double>::input(ScalarStr, Ctxt, S.DoubleValue);
+    case Scalar::SK_Bool:
+      return ScalarTraits<bool>::input(ScalarStr, Ctxt, S.BoolValue);
+    }
+    llvm_unreachable("unknown scalar kind");
+  }
+
+  static QuotingType mustQuote(const Scalar &S, StringRef Str) {
+    switch (S.SKind) {
+    case Scalar::SK_Unknown:
+      report_fatal_error("quote unknown scalar");
+    case Scalar::SK_Double:
+      return ScalarTraits<double>::mustQuote(Str);
+    case Scalar::SK_Bool:
+      return ScalarTraits<bool>::mustQuote(Str);
+    }
+    llvm_unreachable("unknown scalar kind");
+  }
+};
+
+template <> struct CustomMappingTraits<Map> {
+  static void inputOne(IO &IO, StringRef Key, Map &M) {
+    IO.mapRequired(Key.str().c_str(), M[Key]);
+  }
+
+  static void output(IO &IO, Map &M) {
+    for (auto &N : M)
+      IO.mapRequired(N.getKey().str().c_str(), N.getValue());
+  }
+};
+
+template <> struct SequenceTraits<Seq> {
+  static size_t size(IO &IO, Seq &A) { return A.size(); }
+
+  static std::unique_ptr<Poly> &element(IO &IO, Seq &A, size_t Index) {
+    if (Index >= A.size())
+      A.resize(Index + 1);
+    return A[Index];
+  }
+};
+
+} // namespace yaml
+} // namespace llvm
+
+TEST(YAMLIO, TestReadWritePolymorphicScalar) {
+  std::string intermediate;
+  std::unique_ptr<Poly> node = llvm::make_unique<Scalar>(true);
+
+  llvm::raw_string_ostream ostr(intermediate);
+  Output yout(ostr);
+#ifdef GTEST_HAS_DEATH_TEST
+#ifndef NDEBUG
+  EXPECT_DEATH(yout << node, "plain scalar documents are not supported");
+#endif
+#endif
+}
+
+TEST(YAMLIO, TestReadWritePolymorphicSeq) {
+  std::string intermediate;
+  {
+    auto seq = llvm::make_unique<Seq>();
+    seq->push_back(llvm::make_unique<Scalar>(true));
+    seq->push_back(llvm::make_unique<Scalar>(1.0));
+    auto node = llvm::unique_dyn_cast<Poly>(seq);
+
+    llvm::raw_string_ostream ostr(intermediate);
+    Output yout(ostr);
+    yout << node;
+  }
+  {
+    Input yin(intermediate);
+    std::unique_ptr<Poly> node;
+    yin >> node;
+
+    EXPECT_FALSE(yin.error());
+    auto seq = llvm::dyn_cast<Seq>(node.get());
+    ASSERT_TRUE(seq);
+    ASSERT_EQ(seq->size(), 2u);
+    auto first = llvm::dyn_cast<Scalar>((*seq)[0].get());
+    ASSERT_TRUE(first);
+    EXPECT_EQ(first->SKind, Scalar::SK_Bool);
+    EXPECT_TRUE(first->BoolValue);
+    auto second = llvm::dyn_cast<Scalar>((*seq)[1].get());
+    ASSERT_TRUE(second);
+    EXPECT_EQ(second->SKind, Scalar::SK_Double);
+    EXPECT_EQ(second->DoubleValue, 1.0);
+  }
+}
+
+TEST(YAMLIO, TestReadWritePolymorphicMap) {
+  std::string intermediate;
+  {
+    auto map = llvm::make_unique<Map>();
+    (*map)["foo"] = llvm::make_unique<Scalar>(false);
+    (*map)["bar"] = llvm::make_unique<Scalar>(2.0);
+    std::unique_ptr<Poly> node = llvm::unique_dyn_cast<Poly>(map);
+
+    llvm::raw_string_ostream ostr(intermediate);
+    Output yout(ostr);
+    yout << node;
+  }
+  {
+    Input yin(intermediate);
+    std::unique_ptr<Poly> node;
+    yin >> node;
+
+    EXPECT_FALSE(yin.error());
+    auto map = llvm::dyn_cast<Map>(node.get());
+    ASSERT_TRUE(map);
+    auto foo = llvm::dyn_cast<Scalar>((*map)["foo"].get());
+    ASSERT_TRUE(foo);
+    EXPECT_EQ(foo->SKind, Scalar::SK_Bool);
+    EXPECT_FALSE(foo->BoolValue);
+    auto bar = llvm::dyn_cast<Scalar>((*map)["bar"].get());
+    ASSERT_TRUE(bar);
+    EXPECT_EQ(bar->SKind, Scalar::SK_Double);
+    EXPECT_EQ(bar->DoubleValue, 2.0);
+  }
+}