5 files changed, 60 insertions, 28 deletions
diff --git a/llvm/lib/ObjectYAML/CMakeLists.txt b/llvm/lib/ObjectYAML/CMakeLists.txt
index 7e1363d8dac..77370908046 100644
--- a/llvm/lib/ObjectYAML/CMakeLists.txt
+++ b/llvm/lib/ObjectYAML/CMakeLists.txt
@@ -3,4 +3,5 @@ add_llvm_library(LLVMObjectYAML
   COFFYAML.cpp
   ELFYAML.cpp
   MachOYAML.cpp
+  ObjectYAML.cpp
   )
diff --git a/llvm/lib/ObjectYAML/COFFYAML.cpp b/llvm/lib/ObjectYAML/COFFYAML.cpp
index 02ae0535615..7f9f4c1f8c2 100644
--- a/llvm/lib/ObjectYAML/COFFYAML.cpp
+++ b/llvm/lib/ObjectYAML/COFFYAML.cpp
@@ -493,6 +493,7 @@ void MappingTraits<COFFYAML::Section>::mapping(IO &IO, COFFYAML::Section &Sec) {
 }
 
 void MappingTraits<COFFYAML::Object>::mapping(IO &IO, COFFYAML::Object &Obj) {
+  IO.mapTag("!COFF", true);
   IO.mapOptional("OptionalHeader", Obj.OptionalHeader);
   IO.mapRequired("header", Obj.Header);
   IO.mapRequired("sections", Obj.Sections);
diff --git a/llvm/lib/ObjectYAML/ELFYAML.cpp b/llvm/lib/ObjectYAML/ELFYAML.cpp
index a480783e44c..521e2375d08 100644
--- a/llvm/lib/ObjectYAML/ELFYAML.cpp
+++ b/llvm/lib/ObjectYAML/ELFYAML.cpp
@@ -820,6 +820,7 @@ void MappingTraits<ELFYAML::Relocation>::mapping(IO &IO,
 void MappingTraits<ELFYAML::Object>::mapping(IO &IO, ELFYAML::Object &Object) {
   assert(!IO.getContext() && "The IO context is initialized already");
   IO.setContext(&Object);
+  IO.mapTag("!ELF", true);
   IO.mapRequired("FileHeader", Object.Header);
   IO.mapOptional("Sections", Object.Sections);
   IO.mapOptional("Symbols", Object.Symbols);
diff --git a/llvm/lib/ObjectYAML/MachOYAML.cpp b/llvm/lib/ObjectYAML/MachOYAML.cpp
index 15ff4efb5c3..0a303a6e739 100644
--- a/llvm/lib/ObjectYAML/MachOYAML.cpp
+++ b/llvm/lib/ObjectYAML/MachOYAML.cpp
@@ -132,34 +132,6 @@ void MappingTraits<MachOYAML::UniversalBinary>::mapping(
     IO.setContext(nullptr);
 }
 
-void MappingTraits<MachOYAML::MachFile>::mapping(
-    IO &IO, MachOYAML::MachFile &MachFile) {
-  if (!IO.getContext()) {
-    IO.setContext(&MachFile);
-  }
-  if (IO.outputting()) {
-    if (MachFile.isFat) {
-      IO.mapTag("!fat-mach-o", true);
-      MappingTraits<MachOYAML::UniversalBinary>::mapping(IO, MachFile.FatFile);
-    } else {
-      IO.mapTag("!mach-o", true);
-      MappingTraits<MachOYAML::Object>::mapping(IO, MachFile.ThinFile);
-    }
-  } else {
-    if (IO.mapTag("!fat-mach-o")) {
-      MachFile.isFat = true;
-      MappingTraits<MachOYAML::UniversalBinary>::mapping(IO, MachFile.FatFile);
-    } else if (IO.mapTag("!mach-o")) {
-      MachFile.isFat = false;
-      MappingTraits<MachOYAML::Object>::mapping(IO, MachFile.ThinFile);
-    } else {
-      assert(false && "No tag found in YAML, cannot identify file type!");
-    }
-  }
-  if (IO.getContext() == &MachFile)
-    IO.setContext(nullptr);
-}
-
 void MappingTraits<MachOYAML::LinkEditData>::mapping(
     IO &IO, MachOYAML::LinkEditData &LinkEditData) {
   IO.mapOptional("RebaseOpcodes", LinkEditData.RebaseOpcodes);
diff --git a/llvm/lib/ObjectYAML/ObjectYAML.cpp b/llvm/lib/ObjectYAML/ObjectYAML.cpp
new file mode 100644
index 00000000000..97741b5ec8b
--- /dev/null
+++ b/llvm/lib/ObjectYAML/ObjectYAML.cpp
@@ -0,0 +1,57 @@
+//===- ObjectYAML.cpp - YAML utilities for object files -------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines a wrapper class for handling tagged YAML input
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ObjectYAML/YAML.h"
+#include "llvm/ObjectYAML/ObjectYAML.h"
+
+using namespace llvm;
+using namespace yaml;
+
+void MappingTraits<YamlObjectFile>::mapping(IO &IO,
+                                            YamlObjectFile &ObjectFile) {
+  if (IO.outputting()) {
+    if (ObjectFile.Elf)
+      MappingTraits<ELFYAML::Object>::mapping(IO, *ObjectFile.Elf);
+    if (ObjectFile.Coff)
+      MappingTraits<COFFYAML::Object>::mapping(IO, *ObjectFile.Coff);
+    if (ObjectFile.MachO)
+      MappingTraits<MachOYAML::Object>::mapping(IO, *ObjectFile.MachO);
+    if (ObjectFile.FatMachO)
+      MappingTraits<MachOYAML::UniversalBinary>::mapping(IO,
+                                                         *ObjectFile.FatMachO);
+  } else {
+    if (IO.mapTag("!ELF")) {
+      ObjectFile.Elf.reset(new ELFYAML::Object());
+      MappingTraits<ELFYAML::Object>::mapping(IO, *ObjectFile.Elf);
+    } else if (IO.mapTag("!COFF")) {
+      ObjectFile.Coff.reset(new COFFYAML::Object());
+      MappingTraits<COFFYAML::Object>::mapping(IO, *ObjectFile.Coff);
+    } else if (IO.mapTag("!mach-o")) {
+      ObjectFile.MachO.reset(new MachOYAML::Object());
+      MappingTraits<MachOYAML::Object>::mapping(IO, *ObjectFile.MachO);
+    } else if (IO.mapTag("!fat-mach-o")) {
+      ObjectFile.FatMachO.reset(new MachOYAML::UniversalBinary());
+      MappingTraits<MachOYAML::UniversalBinary>::mapping(IO,
+                                                         *ObjectFile.FatMachO);
+    } else {
+      Input &In = (Input &)IO;
+      std::string Tag = In.getCurrentNode()->getRawTag();
+      if (Tag.empty())
+        IO.setError("YAML Object File missing document type tag!");
+      else
+        IO.setError(
+            llvm::Twine("YAML Object File unsupported document type tag '") +
+            llvm::Twine(Tag.c_str()) + llvm::Twine("'!"));
+    }
+  }
+}