1 files changed, 20 insertions, 42 deletions

diff --git a/llvm/tools/yaml2obj/yaml2obj.cpp b/llvm/tools/yaml2obj/yaml2obj.cpp
index a2a8e0b8b83..f746d84a389 100644
--- a/llvm/tools/yaml2obj/yaml2obj.cpp
+++ b/llvm/tools/yaml2obj/yaml2obj.cpp
@@ -16,6 +16,7 @@
 
 #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"
@@ -32,29 +33,6 @@ 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)"));
@@ -62,14 +40,26 @@ DocNum("docnum", cl::init(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
@@ -99,21 +89,9 @@ int main(int argc, char **argv) {
   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();