[yaml2obj] Remove --format option in favor of YAML tags

[beanz]

Summary

Our YAML library's handling of tags isn't perfect, but it is good enough to get rid of the need for the --format argument to yaml2obj. This patch does exactly that.

Instead of requiring --format, it infers the format based on the tags found in the object file. The supported tags are:

!ELF
!COFF
!mach-o
!fat-mach-o

I have a corresponding patch that is quite large that fixes up all the in-tree test cases.

[compnerd]

Nice! I think it would be a good idea to add a few tests. Basically, feed it a yaml document and ensure that the format detection happens correctly.

tools/yaml2obj/yaml2macho.cpp

450

Not your fault, but its hard to see that the code here is switching between a fat macho or a sliced/lipo'ed/thin one.

[silvas]

Makes sense. This is how we originally wanted to do it but didn't have tag support IIRC.

include/llvm/ObjectYAML/MachOYAML.h

 };
 
 struct UniversalBinary {
   FatHeader Header;
   std::vector<FatArch> FatArchs;
   std::vector<Object> Slices;
 };
 
-struct MachFile {
-  bool isFat;
-  UniversalBinary FatFile;
-  Object ThinFile;
-};
-
 } // namespace llvm::MachOYAML
 } // namespace llvm
 
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::MachOYAML::LoadCommand)
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::MachOYAML::Section)
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::Hex8)
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::Hex64)
 LLVM_YAML_IS_SEQUENCE_VECTOR(int64_t)

 template <> struct MappingTraits<MachOYAML::FatArch> {
   static void mapping(IO &IO, MachOYAML::FatArch &FatArch);
 };
 
 template <> struct MappingTraits<MachOYAML::UniversalBinary> {
   static void mapping(IO &IO, MachOYAML::UniversalBinary &UniversalBinary);
 };
 
-template <> struct MappingTraits<MachOYAML::MachFile> {
-  static void mapping(IO &IO, MachOYAML::MachFile &MachFile);
-};
-
 template <> struct MappingTraits<MachOYAML::LoadCommand> {
   static void mapping(IO &IO, MachOYAML::LoadCommand &LoadCommand);
 };
 
 template <> struct MappingTraits<MachOYAML::LinkEditData> {
   static void mapping(IO &IO, MachOYAML::LinkEditData &LinkEditData);
 };
 

include/llvm/ObjectYAML/ObjectYAML.h

+//===- ObjectYAML.h ---------------------------------------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_OBJECTYAML_OBJECTYAML_H
+#define LLVM_OBJECTYAML_OBJECTYAML_H
+
+#include "llvm/Support/YAMLTraits.h"
+#include "llvm/ObjectYAML/ELFYAML.h"
+#include "llvm/ObjectYAML/COFFYAML.h"
+#include "llvm/ObjectYAML/MachOYAML.h"
+
+namespace llvm {
+namespace yaml {
+
+struct YamlObjectFile {
+  std::unique_ptr<ELFYAML::Object> Elf;
+  std::unique_ptr<COFFYAML::Object> Coff;
+  std::unique_ptr<MachOYAML::Object> MachO;
+  std::unique_ptr<MachOYAML::UniversalBinary> FatMachO;
+};
+
+template <> struct MappingTraits<YamlObjectFile> {
+  static void mapping(IO &IO, YamlObjectFile &ObjectFile);
+};
+
+} // namespace yaml
+} // namespace llvm
+
+#endif

lib/ObjectYAML/CMakeLists.txt

 add_llvm_library(LLVMObjectYAML
   YAML.cpp
   COFFYAML.cpp
   ELFYAML.cpp
   MachOYAML.cpp
+  ObjectYAML.cpp
   )

lib/ObjectYAML/COFFYAML.cpp

void MappingTraits<COFFYAML::Section>::mapping(IO &IO, COFFYAML::Section &Sec) {
   IO.mapOptional("VirtualAddress", Sec.Header.VirtualAddress, 0U);
   IO.mapOptional("VirtualSize", Sec.Header.VirtualSize, 0U);
   IO.mapOptional("Alignment", Sec.Alignment, 0U);
   IO.mapRequired("SectionData", Sec.SectionData);
   IO.mapOptional("Relocations", Sec.Relocations);
 }
 
 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);
   IO.mapRequired("symbols", Obj.Symbols);
 }
 
 }
 }

lib/ObjectYAML/ELFYAML.cpp

if (Object->Header.Machine == ELFYAML::ELF_EM(ELF::EM_MIPS) &&
     IO.mapRequired("Type", Rel.Type);
 
   IO.mapOptional("Addend", Rel.Addend, (int64_t)0);
 }
 
 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);
   IO.setContext(nullptr);
 }
 
 LLVM_YAML_STRONG_TYPEDEF(uint8_t, MIPS_AFL_REG)
 LLVM_YAML_STRONG_TYPEDEF(uint8_t, MIPS_ABI_FP)
 LLVM_YAML_STRONG_TYPEDEF(uint32_t, MIPS_AFL_EXT)
 LLVM_YAML_STRONG_TYPEDEF(uint32_t, MIPS_AFL_ASE)
 LLVM_YAML_STRONG_TYPEDEF(uint32_t, MIPS_AFL_FLAGS1)
 
 } // end namespace yaml
 } // end namespace llvm

lib/ObjectYAML/MachOYAML.cpp

void MappingTraits<MachOYAML::UniversalBinary>::mapping(
   IO.mapRequired("FatHeader", UniversalBinary.Header);
   IO.mapRequired("FatArchs", UniversalBinary.FatArchs);
   IO.mapRequired("Slices", UniversalBinary.Slices);
 
   if (IO.getContext() == &UniversalBinary)
     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);
   IO.mapOptional("BindOpcodes", LinkEditData.BindOpcodes);
   IO.mapOptional("WeakBindOpcodes", LinkEditData.WeakBindOpcodes);
   IO.mapOptional("LazyBindOpcodes", LinkEditData.LazyBindOpcodes);
   IO.mapOptional("ExportTrie", LinkEditData.ExportTrie);
   IO.mapOptional("NameList", LinkEditData.NameList);

lib/ObjectYAML/ObjectYAML.cpp

+//===- 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);
+    }
+  }
+}

tools/obj2yaml/macho2yaml.cpp

 //===------ macho2yaml.cpp - obj2yaml conversion tool -----------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Error.h"
 #include "obj2yaml.h"
 #include "llvm/Object/MachOUniversal.h"
-#include "llvm/ObjectYAML/MachOYAML.h"
+#include "llvm/ObjectYAML/ObjectYAML.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/LEB128.h"
 
 #include <string.h> // for memcpy
 
 using namespace llvm;
 
 class MachODumper {

 }
 
 Error macho2yaml(raw_ostream &Out, const object::MachOObjectFile &Obj) {
   MachODumper Dumper(Obj);
   Expected<std::unique_ptr<MachOYAML::Object>> YAML = Dumper.dump();
   if (!YAML)
     return YAML.takeError();
 
+  yaml::YamlObjectFile YAMLFile;
+  YAMLFile.MachO = std::move(YAML.get());
+
   yaml::Output Yout(Out);
-  Yout << *(YAML.get());
+  Yout << YAMLFile;
   return Error::success();
 }
 
 Error macho2yaml(raw_ostream &Out, const object::MachOUniversalBinary &Obj) {
-  MachOYAML::MachFile YAMLFile;
-  YAMLFile.isFat = true;
-  MachOYAML::UniversalBinary &YAML = YAMLFile.FatFile;
+  yaml::YamlObjectFile YAMLFile;
+  YAMLFile.FatMachO.reset(new MachOYAML::UniversalBinary());
+  MachOYAML::UniversalBinary &YAML = *YAMLFile.FatMachO;
   YAML.Header.magic = Obj.getMagic();
   YAML.Header.nfat_arch = Obj.getNumberOfObjects();
 
   for (auto Slice : Obj.objects()) {
     MachOYAML::FatArch arch;
     arch.cputype = Slice.getCPUType();
     arch.cpusubtype = Slice.getCPUSubType();
     arch.offset = Slice.getOffset();

tools/yaml2obj/yaml2coff.cpp

 //===----------------------------------------------------------------------===//
 
 #include "yaml2obj.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/Object/COFF.h"
-#include "llvm/ObjectYAML/COFFYAML.h"
+#include "llvm/ObjectYAML/ObjectYAML.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
 #include <vector>
 
 using namespace llvm;
 
static bool writeCOFF(COFFParser &CP, raw_ostream &OS) {
   }
 
   // Output string table.
   if (CP.Obj.Header.PointerToSymbolTable)
     OS.write(&CP.StringTable[0], CP.StringTable.size());
   return true;
 }
 
-int yaml2coff(yaml::Input &YIn, raw_ostream &Out) {
-  COFFYAML::Object Doc;
-  YIn >> Doc;
-  if (YIn.error()) {
-    errs() << "yaml2obj: Failed to parse YAML file!\n";
-    return 1;
-  }
-
+int yaml2coff(llvm::COFFYAML::Object &Doc, raw_ostream &Out) {
   COFFParser CP(Doc);
   if (!CP.parse()) {
     errs() << "yaml2obj: Failed to parse YAML file!\n";
     return 1;
   }
 
   if (!layoutOptionalHeader(CP)) {
     errs() << "yaml2obj: Failed to layout optional header for COFF file!\n";

tools/yaml2obj/yaml2elf.cpp

 static bool is64Bit(const ELFYAML::Object &Doc) {
   return Doc.Header.Class == ELFYAML::ELF_ELFCLASS(ELF::ELFCLASS64);
 }
 
 static bool isLittleEndian(const ELFYAML::Object &Doc) {
   return Doc.Header.Data == ELFYAML::ELF_ELFDATA(ELF::ELFDATA2LSB);
 }
 
-int yaml2elf(yaml::Input &YIn, raw_ostream &Out) {
-  ELFYAML::Object Doc;
-  YIn >> Doc;
-  if (YIn.error()) {
-    errs() << "yaml2obj: Failed to parse YAML file!\n";
-    return 1;
-  }
+int yaml2elf(llvm::ELFYAML::Object &Doc, raw_ostream &Out) {
   using object::ELFType;
   typedef ELFType<support::little, true> LE64;
   typedef ELFType<support::big, true> BE64;
   typedef ELFType<support::little, false> LE32;
   typedef ELFType<support::big, false> BE32;
   if (is64Bit(Doc)) {
     if (isLittleEndian(Doc))
       return ELFState<LE64>::writeELF(Out, Doc);

tools/yaml2obj/yaml2macho.cpp

 //===- yaml2macho - Convert YAML to a Mach object file --------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 ///
 /// \file
 /// \brief The Mach component of yaml2obj.
 ///
 //===----------------------------------------------------------------------===//
 
 #include "yaml2obj.h"
-#include "llvm/ObjectYAML/MachOYAML.h"
+#include "llvm/ObjectYAML/ObjectYAML.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/LEB128.h"
 #include "llvm/Support/MachO.h"
 #include "llvm/Support/YAMLTraits.h"
 #include "llvm/Support/raw_ostream.h"
 
 #include "llvm/Support/Format.h"
 
for (auto Str : Obj.LinkEdit.StringTable) {
     OS.write(Str.data(), Str.size());
     OS.write('\0');
   }
   return Error::success();
 }
 
 class UniversalWriter {
 public:
-  UniversalWriter(MachOYAML::MachFile &MachFile)
-      : MachFile(MachFile), fileStart(0) {}
+  UniversalWriter(yaml::YamlObjectFile &ObjectFile)
+      : ObjectFile(ObjectFile), fileStart(0) {}
 
   Error writeMachO(raw_ostream &OS);
 
 private:
   Error writeFatHeader(raw_ostream &OS);
   Error writeFatArchs(raw_ostream &OS);
 
   void ZeroToOffset(raw_ostream &OS, size_t offset);
 
-  MachOYAML::MachFile &MachFile;
+  yaml::YamlObjectFile &ObjectFile;
   uint64_t fileStart;
 };
 
 Error UniversalWriter::writeMachO(raw_ostream &OS) {
   fileStart = OS.tell();
-  if (!MachFile.isFat) {
-    MachOWriter Writer(MachFile.ThinFile);
+  if (ObjectFile.MachO) {
+    MachOWriter Writer(*ObjectFile.MachO);
     return Writer.writeMachO(OS);
   }
   if (auto Err = writeFatHeader(OS))
     return Err;
   if (auto Err = writeFatArchs(OS))
     return Err;
-  auto &FatFile = MachFile.FatFile;
+  auto &FatFile = *ObjectFile.FatMachO;
   assert(FatFile.FatArchs.size() == FatFile.Slices.size());
   for (size_t i = 0; i < FatFile.Slices.size(); i++) {
     ZeroToOffset(OS, FatFile.FatArchs[i].offset);
     MachOWriter Writer(FatFile.Slices[i]);
     if (auto Err = Writer.writeMachO(OS))
       return Err;
     auto SliceEnd = FatFile.FatArchs[i].offset + FatFile.FatArchs[i].size;
     ZeroToOffset(OS, SliceEnd);
   }
   return Error::success();
 }
 
 Error UniversalWriter::writeFatHeader(raw_ostream &OS) {
-  auto &FatFile = MachFile.FatFile;
+  auto &FatFile = *ObjectFile.FatMachO;
   MachO::fat_header header;
   header.magic = FatFile.Header.magic;
   header.nfat_arch = FatFile.Header.nfat_arch;
   if (sys::IsLittleEndianHost)
     swapStruct(header);
   OS.write(reinterpret_cast<const char *>(&header), sizeof(MachO::fat_header));
   return Error::success();
 }
void writeFatArch<MachO::fat_arch_64>(MachOYAML::FatArch &Arch,
   FatArch.reserved = Arch.reserved;
   if (sys::IsLittleEndianHost)
     swapStruct(FatArch);
   OS.write(reinterpret_cast<const char *>(&FatArch),
            sizeof(MachO::fat_arch_64));
 }
 
 Error UniversalWriter::writeFatArchs(raw_ostream &OS) {
-  auto &FatFile = MachFile.FatFile;
+  auto &FatFile = *ObjectFile.FatMachO;
   bool is64Bit = FatFile.Header.magic == MachO::FAT_MAGIC_64;
   for (auto Arch : FatFile.FatArchs) {
     if (is64Bit)
       writeFatArch<MachO::fat_arch_64>(Arch, OS);
     else
       writeFatArch<MachO::fat_arch>(Arch, OS);
   }
 
   return Error::success();
 }
 
 void UniversalWriter::ZeroToOffset(raw_ostream &OS, size_t Offset) {
   auto currOffset = OS.tell() - fileStart;
   if (currOffset < Offset)
     ZeroFillBytes(OS, Offset - currOffset);
 }
 
 } // end anonymous namespace
 
-int yaml2macho(yaml::Input &YIn, raw_ostream &Out) {
-  MachOYAML::MachFile Doc;
-  YIn >> Doc;
-  if (YIn.error()) {
-    errs() << "yaml2obj: Failed to parse YAML file!\n";
-    return 1;
-  }
-
+int yaml2macho(yaml::YamlObjectFile &Doc, raw_ostream &Out) {
   UniversalWriter Writer(Doc);
   if (auto Err = Writer.writeMachO(Out)) {
     errs() << toString(std::move(Err));
     return 1;
   }
   return 0;
 }

tools/yaml2obj/yaml2obj.h

 //===--- yaml2obj.h - -------------------------------------------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 /// \file
 /// \brief Common declarations for yaml2obj
 //===----------------------------------------------------------------------===//
 #ifndef LLVM_TOOLS_YAML2OBJ_YAML2OBJ_H
 #define LLVM_TOOLS_YAML2OBJ_YAML2OBJ_H
 
 namespace llvm {
 class raw_ostream;
+
+namespace COFFYAML {
+struct Object;
+}
+
+namespace ELFYAML {
+struct Object;
+}
+
 namespace yaml {
 class Input;
+struct YamlObjectFile;
 }
 }
-int yaml2coff(llvm::yaml::Input &YIn, llvm::raw_ostream &Out);
-int yaml2elf(llvm::yaml::Input &YIn, llvm::raw_ostream &Out);
-int yaml2macho(llvm::yaml::Input &YIn, llvm::raw_ostream &Out);
+
+int yaml2coff(llvm::COFFYAML::Object &Doc, llvm::raw_ostream &Out);
+int yaml2elf(llvm::ELFYAML::Object &Doc, llvm::raw_ostream &Out);
+int yaml2macho(llvm::yaml::YamlObjectFile &Doc, llvm::raw_ostream &Out);
 
 #endif

tools/yaml2obj/yaml2obj.cpp

 // binary equivalent.
 //
 // This is used for writing tests that require binary files.
 //
 //===----------------------------------------------------------------------===//
 
 #include "yaml2obj.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/ObjectYAML/ObjectYAML.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/PrettyStackTrace.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/ToolOutputFile.h"
 #include "llvm/Support/YAMLTraits.h"
 #include "llvm/Support/raw_ostream.h"
 #include <system_error>
 
 using namespace llvm;
 
 static cl::opt<std::string>
   Input(cl::Positional, cl::desc("<input>"), cl::init("-"));
 
-// TODO: The "right" way to tell what kind of object file a given YAML file
-// corresponds to is to look at YAML "tags" (e.g. `!Foo`). Then, different
-// tags (`!ELF`, `!COFF`, etc.) would be used to discriminate between them.
-// Interpreting the tags is needed eventually for when writing test cases,
-// so that we can e.g. have `!Archive` contain a sequence of `!ELF`, and
-// just Do The Right Thing. However, interpreting these tags and acting on
-// them appropriately requires some work in the YAML parser and the YAMLIO
-// library.
-enum YAMLObjectFormat {
-  YOF_COFF,
-  YOF_ELF,
-  YOF_MACHO
-};
-
-cl::opt<YAMLObjectFormat> Format(
-  "format",
-  cl::desc("Interpret input as this type of object file"),
-  cl::values(
-    clEnumValN(YOF_COFF, "coff", "COFF object file format"),
-    clEnumValN(YOF_ELF, "elf", "ELF object file format"),
-    clEnumValN(YOF_MACHO, "macho", "Mach-O object file format"),
-  clEnumValEnd));
-
 cl::opt<unsigned>
 DocNum("docnum", cl::init(1),
        cl::desc("Read specified document from input (default = 1)"));
 
 static cl::opt<std::string> OutputFilename("o", cl::desc("Output filename"),
                                            cl::value_desc("filename"));
 
-typedef int (*ConvertFuncPtr)(yaml::Input & YIn, raw_ostream &Out);
-
-static int convertYAML(yaml::Input &YIn, raw_ostream &Out,
-                       ConvertFuncPtr Convert) {
+static int convertYAML(yaml::Input &YIn, raw_ostream &Out) {
   unsigned CurDocNum = 0;
   do {
-    if (++CurDocNum == DocNum)
-      return Convert(YIn, Out);
+    if (++CurDocNum == DocNum) {
+      yaml::YamlObjectFile Doc;
+      YIn >> Doc;
+      if (YIn.error()) {
+        errs() << "yaml2obj: Failed to parse YAML file!\n";
+        return 1;
+      }
+
+      if (Doc.Elf)
+        return yaml2elf(*Doc.Elf, Out);
+      if (Doc.Coff)
+        return yaml2coff(*Doc.Coff, Out);
+      if (Doc.MachO || Doc.FatMachO)
+        return yaml2macho(Doc, Out);
+      errs() << "yaml2obj: Unknown document type!\n";
+      return 1;
+    }
   } while (YIn.nextDocument());
 
   errs() << "yaml2obj: Cannot find the " << DocNum
          << llvm::getOrdinalSuffix(DocNum) << " document\n";
   return 1;
 }
 
 int main(int argc, char **argv) {
if (EC) {
     return 1;
   }
 
   ErrorOr<std::unique_ptr<MemoryBuffer>> Buf =
       MemoryBuffer::getFileOrSTDIN(Input);
   if (!Buf)
     return 1;
 
-  ConvertFuncPtr Convert = nullptr;
-  if (Format == YOF_COFF)
-    Convert = yaml2coff;
-  else if (Format == YOF_ELF)
-    Convert = yaml2elf;
-  else if (Format == YOF_MACHO)
-    Convert = yaml2macho;
-  else {
-    errs() << "Not yet implemented\n";
-    return 1;
-  }
-
   yaml::Input YIn(Buf.get()->getBuffer());
 
-  int Res = convertYAML(YIn, Out->os(), Convert);
+  int Res = convertYAML(YIn, Out->os());
   if (Res == 0)
     Out->keep();
 
   return Res;
 }