4 files changed, 121 insertions, 2 deletions

diff --git a/docs/YamlIO.rst b/docs/YamlIO.rst
index 79e07cd989..3ecd03afb2 100644
--- a/docs/YamlIO.rst
+++ b/docs/YamlIO.rst
@@ -633,6 +633,20 @@ This works for both reading and writing. For example:
     };
 
 
+Tags
+----
+
+The YAML syntax supports tags as a way to specify the type of a node before
+it is parsed. This allows dynamic types of nodes.  But the YAML I/O model uses
+static typing, so there are limits to how you can use tags with the YAML I/O
+model. Recently, we added support to YAML I/O for checking/setting the optional 
+tag on a map. Using this functionality it is even possbile to support differnt 
+mappings, as long as they are convertable.  
+
+To check a tag, inside your mapping() method you can use io.mapTag() to specify
+what the tag should be.  This will also add that tag when writing yaml.
+
+
 Sequence
 ========
 
diff --git a/include/llvm/Support/YAMLTraits.h b/include/llvm/Support/YAMLTraits.h
index 98964fc002..6b6a8b7698 100644
--- a/include/llvm/Support/YAMLTraits.h
+++ b/include/llvm/Support/YAMLTraits.h
@@ -330,6 +330,7 @@ public:
   virtual void postflightFlowElement(void*) = 0;
   virtual void endFlowSequence() = 0;
 
+  virtual bool mapTag(StringRef Tag, bool Default=false) = 0;
   virtual void beginMapping() = 0;
   virtual void endMapping() = 0;
   virtual bool preflightKey(const char*, bool, bool, bool &, void *&) = 0;
@@ -404,8 +405,7 @@ public:
   void mapOptional(const char* Key, T& Val, const T& Default) {
     this->processKeyWithDefault(Key, Val, Default, false);
   }
-
-
+  
 private:
   template <typename T>
   void processKeyWithDefault(const char *Key, T &Val, const T& DefaultValue,
@@ -696,6 +696,7 @@ public:
 
 private:
   virtual bool outputting();
+  virtual bool mapTag(StringRef, bool);
   virtual void beginMapping();
   virtual void endMapping();
   virtual bool preflightKey(const char *, bool, bool, bool &, void *&);
@@ -819,6 +820,7 @@ public:
   virtual ~Output();
 
   virtual bool outputting();
+  virtual bool mapTag(StringRef, bool);
   virtual void beginMapping();
   virtual void endMapping();
   virtual bool preflightKey(const char *key, bool, bool, bool &, void *&);
diff --git a/lib/Support/YAMLTraits.cpp b/lib/Support/YAMLTraits.cpp
index 19eaed1ac7..415424ebed 100644
--- a/lib/Support/YAMLTraits.cpp
+++ b/lib/Support/YAMLTraits.cpp
@@ -81,6 +81,16 @@ bool Input::setCurrentDocument() {
 void Input::nextDocument() {
   ++DocIterator;
 }
+  
+bool Input::mapTag(StringRef Tag, bool Default) {
+  StringRef foundTag = CurrentNode->_node->getVerbatimTag();
+  if (foundTag.empty()) {
+    // If no tag found and 'Tag' is the default, say it was found.
+    return Default;
+  }
+  // Return true iff found tag matches supplied tag.
+  return Tag.equals(foundTag);
+}
 
 void Input::beginMapping() {
   if (EC)
@@ -381,6 +391,14 @@ void Output::beginMapping() {
   NeedsNewLine = true;
 }
 
+bool Output::mapTag(StringRef Tag, bool Use) {
+  if (Use) {
+    this->output(" ");
+    this->output(Tag);
+  }
+  return Use;
+}
+
 void Output::endMapping() {
   StateStack.pop_back();
 }
diff --git a/unittests/Support/YAMLIOTest.cpp b/unittests/Support/YAMLIOTest.cpp
index 8ae05f4b60..db70b91ea1 100644
--- a/unittests/Support/YAMLIOTest.cpp
+++ b/unittests/Support/YAMLIOTest.cpp
@@ -989,6 +989,91 @@ TEST(YAMLIO, TestSequenceDocListWriteAndRead) {
   }
 }
 
+//===----------------------------------------------------------------------===//
+//  Test document tags
+//===----------------------------------------------------------------------===//
+
+struct MyDouble {
+  MyDouble() : value(0.0) { }
+  MyDouble(double x) : value(x) { }
+  double value;
+};
+
+LLVM_YAML_IS_DOCUMENT_LIST_VECTOR(MyDouble);
+
+
+namespace llvm {
+namespace yaml {
+  template <>
+  struct MappingTraits<MyDouble> {
+    static void mapping(IO &io, MyDouble &d) {
+      if (io.mapTag("!decimal", true)) {
+        mappingDecimal(io, d);
+      } else if (io.mapTag("!fraction")) {
+        mappingFraction(io, d);
+      }
+    }
+    static void mappingDecimal(IO &io, MyDouble &d) {
+      io.mapRequired("value", d.value);
+    }
+    static void mappingFraction(IO &io, MyDouble &d) {
+        double num, denom;
+        io.mapRequired("numerator",      num);
+        io.mapRequired("denominator",    denom);
+        // convert fraction to double
+        d.value = num/denom;
+    }
+  };
+ }
+}
+
+
+//
+// Test the reading of two different tagged yaml documents.
+//
+TEST(YAMLIO, TestTaggedDocuments) {
+  std::vector<MyDouble> docList;
+  Input yin("--- !decimal\nvalue:  3.0\n"
+            "--- !fraction\nnumerator:  9.0\ndenominator:  2\n...\n");
+  yin >> docList;
+  EXPECT_FALSE(yin.error());
+  EXPECT_EQ(docList.size(), 2UL);
+  EXPECT_EQ(docList[0].value, 3.0);
+  EXPECT_EQ(docList[1].value, 4.5);
+}
+
+
+
+//
+// Test writing then reading back tagged documents
+//
+TEST(YAMLIO, TestTaggedDocumentsWriteAndRead) {
+  std::string intermediate;
+  {
+    MyDouble a(10.25);
+    MyDouble b(-3.75);
+    std::vector<MyDouble> docList;
+    docList.push_back(a);
+    docList.push_back(b);
+
+    llvm::raw_string_ostream ostr(intermediate);
+    Output yout(ostr);
+    yout << docList;
+  }
+
+  {
+    Input yin(intermediate);
+    std::vector<MyDouble> docList2;
+    yin >> docList2;
+
+    EXPECT_FALSE(yin.error());
+    EXPECT_EQ(docList2.size(), 2UL);
+    EXPECT_EQ(docList2[0].value, 10.25);
+    EXPECT_EQ(docList2[1].value, -3.75);
+  }
+}
+
+
 
 //===----------------------------------------------------------------------===//
 //  Test error handling