1 files changed, 8 insertions, 31 deletions
diff --git a/llvm/tools/yaml2obj/yaml2obj.cpp b/llvm/tools/yaml2obj/yaml2obj.cpp
index 533c8ea1446..76e6663c095 100644
--- a/llvm/tools/yaml2obj/yaml2obj.cpp
+++ b/llvm/tools/yaml2obj/yaml2obj.cpp
@@ -13,7 +13,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "yaml2obj.h"
+#include "llvm/ObjectYAML/yaml2obj.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ObjectYAML/ObjectYAML.h"
 #include "llvm/Support/CommandLine.h"
@@ -21,6 +21,7 @@
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/ToolOutputFile.h"
+#include "llvm/Support/WithColor.h"
 #include "llvm/Support/YAMLTraits.h"
 #include "llvm/Support/raw_ostream.h"
 #include <system_error>
@@ -42,32 +43,6 @@ LLVM_ATTRIBUTE_NORETURN static void error(Twine Message) {
   exit(1);
 }
 
-static int convertYAML(yaml::Input &YIn, raw_ostream &Out) {
-  unsigned CurDocNum = 0;
-  do {
-    if (++CurDocNum == DocNum) {
-      yaml::YamlObjectFile Doc;
-      YIn >> Doc;
-      if (YIn.error())
-        error("yaml2obj: Failed to parse YAML file!");
-      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);
-      if (Doc.Minidump)
-        return yaml2minidump(*Doc.Minidump, Out);
-      if (Doc.Wasm)
-        return yaml2wasm(*Doc.Wasm, Out);
-      error("yaml2obj: Unknown document type!");
-    }
-  } while (YIn.nextDocument());
-
-  error("yaml2obj: Cannot find the " + Twine(DocNum) +
-        llvm::getOrdinalSuffix(DocNum) + " document");
-}
-
 int main(int argc, char **argv) {
   InitLLVM X(argc, argv);
   cl::ParseCommandLineOptions(argc, argv);
@@ -87,10 +62,12 @@ int main(int argc, char **argv) {
     return 1;
 
   yaml::Input YIn(Buf.get()->getBuffer());
-  int Res = convertYAML(YIn, Out->os());
-  if (Res == 0)
-    Out->keep();
+  if (Error E = convertYAML(YIn, Out->os(), DocNum)) {
+    logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0]));
+    return 1;
+  }
 
+  Out->keep();
   Out->os().flush();
-  return Res;
+  return 0;
 }