Add dyn_cast<> support to YAML I/O's IO class

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@194655 91177308-0d34-0410-b5e6-96231b3b80d8

3 files changed, 80 insertions, 5 deletions

diff --git a/include/llvm/Support/YAMLTraits.h b/include/llvm/Support/YAMLTraits.h
index 6b6a8b7698..d6eeaace7b 100644
--- a/include/llvm/Support/YAMLTraits.h
+++ b/include/llvm/Support/YAMLTraits.h
@@ -18,6 +18,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/ADT/Twine.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/YAMLParser.h"
@@ -317,7 +318,7 @@ public:
   IO(void *Ctxt=NULL);
   virtual ~IO();
 
-  virtual bool outputting() = 0;
+  virtual bool outputting() const = 0;
 
   virtual unsigned beginSequence() = 0;
   virtual bool preflightElement(unsigned, void *&) = 0;
@@ -694,8 +695,10 @@ public:
   // To set alternate error reporting.
   void setDiagHandler(llvm::SourceMgr::DiagHandlerTy Handler, void *Ctxt = 0);
 
+  static bool classof(const IO *io) { return !io->outputting(); }
+
 private:
-  virtual bool outputting();
+  virtual bool outputting() const;
   virtual bool mapTag(StringRef, bool);
   virtual void beginMapping();
   virtual void endMapping();
@@ -819,7 +822,9 @@ public:
   Output(llvm::raw_ostream &, void *Ctxt=NULL);
   virtual ~Output();
 
-  virtual bool outputting();
+  static bool classof(const IO *io) { return io->outputting(); }
+  
+  virtual bool outputting() const;
   virtual bool mapTag(StringRef, bool);
   virtual void beginMapping();
   virtual void endMapping();
diff --git a/lib/Support/YAMLTraits.cpp b/lib/Support/YAMLTraits.cpp
index 415424ebed..b32ef400ba 100644
--- a/lib/Support/YAMLTraits.cpp
+++ b/lib/Support/YAMLTraits.cpp
@@ -59,7 +59,7 @@ void Input::setDiagHandler(SourceMgr::DiagHandlerTy Handler, void *Ctxt) {
   SrcMgr.setDiagHandler(Handler, Ctxt);
 }
 
-bool Input::outputting() {
+bool Input::outputting() const {
   return false;
 }
 
@@ -382,7 +382,7 @@ Output::Output(raw_ostream &yout, void *context)
 Output::~Output() {
 }
 
-bool Output::outputting() {
+bool Output::outputting() const {
   return true;
 }
 
diff --git a/unittests/Support/YAMLIOTest.cpp b/unittests/Support/YAMLIOTest.cpp
index db70b91ea1..f77df57e57 100644
--- a/unittests/Support/YAMLIOTest.cpp
+++ b/unittests/Support/YAMLIOTest.cpp
@@ -1074,6 +1074,76 @@ TEST(YAMLIO, TestTaggedDocumentsWriteAndRead) {
 }
 
 
+//===----------------------------------------------------------------------===//
+//  Test dyn_cast<> on IO object 
+//===----------------------------------------------------------------------===//
+
+struct DynCast {
+  int value;
+};
+typedef std::vector<DynCast> DynCastSequence;
+
+LLVM_YAML_IS_SEQUENCE_VECTOR(DynCast)
+
+namespace llvm {
+namespace yaml {
+  template <>
+  struct MappingTraits<DynCast> {
+    static void mapping(IO &io, DynCast& info) {
+      // Change 10 to 13 when writing yaml.
+      if (Output *output = dyn_cast<Output>(&io)) {
+        (void)output;
+        if (info.value == 10)
+          info.value = 13;
+      }
+      io.mapRequired("value", info.value);
+      // Change 20 to 23 when parsing yaml.
+      if (Input *input = dyn_cast<Input>(&io)) {
+        (void)input;
+        if (info.value == 20)
+          info.value = 23;
+      }
+    }
+  };
+}
+}
+
+//
+// Test writing then reading back a sequence of mappings
+//
+TEST(YAMLIO, TestDynCast) {
+  std::string intermediate;
+  {
+    DynCast entry1;
+    entry1.value = 10;
+    DynCast entry2;
+    entry2.value = 20;
+    DynCast entry3;
+    entry3.value = 30;
+    DynCastSequence seq;
+    seq.push_back(entry1);
+    seq.push_back(entry2);
+    seq.push_back(entry3);
+
+    llvm::raw_string_ostream ostr(intermediate);
+    Output yout(ostr);
+    yout << seq;
+  }
+
+  {
+    Input yin(intermediate);
+    DynCastSequence seq2;
+    yin >> seq2;
+
+    EXPECT_FALSE(yin.error());
+    EXPECT_EQ(seq2.size(), 3UL);
+    EXPECT_EQ(seq2[0].value, 13);   // Verify changed to 13.
+    EXPECT_EQ(seq2[1].value, 23);   // Verify changed to 23.
+    EXPECT_EQ(seq2[2].value, 30);   // Verify stays same.
+  }
+}
+
+
 
 //===----------------------------------------------------------------------===//
 //  Test error handling