1 files changed, 43 insertions, 33 deletions

diff --git a/llvm/lib/ObjectYAML/yaml2obj.cpp b/llvm/lib/ObjectYAML/yaml2obj.cpp
index 8ca490bf9ca..c18fa5cfdb5 100644
--- a/llvm/lib/ObjectYAML/yaml2obj.cpp
+++ b/llvm/lib/ObjectYAML/yaml2obj.cpp
@@ -8,59 +8,69 @@
 
 #include "llvm/ObjectYAML/yaml2obj.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/Twine.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/ObjectYAML/ObjectYAML.h"
 #include "llvm/Support/Errc.h"
+#include "llvm/Support/WithColor.h"
 #include "llvm/Support/YAMLTraits.h"
 
 namespace llvm {
 namespace yaml {
 
-Error convertYAML(yaml::Input &YIn, raw_ostream &Out, unsigned DocNum) {
-  auto BoolToErr = [](bool Ret) -> Error {
-    if (!Ret)
-      return createStringError(errc::invalid_argument, "yaml2obj failed");
-    return Error::success();
-  };
-
+bool convertYAML(yaml::Input &YIn, raw_ostream &Out, ErrorHandler ErrHandler,
+                 unsigned DocNum) {
   unsigned CurDocNum = 0;
   do {
-    if (++CurDocNum == DocNum) {
-      yaml::YamlObjectFile Doc;
-      YIn >> Doc;
-      if (std::error_code EC = YIn.error())
-        return createStringError(EC, "Failed to parse YAML input!");
-      if (Doc.Elf)
-        return BoolToErr(yaml2elf(*Doc.Elf, Out));
-      if (Doc.Coff)
-        return BoolToErr(yaml2coff(*Doc.Coff, Out));
-      if (Doc.MachO || Doc.FatMachO)
-        return BoolToErr(yaml2macho(Doc, Out));
-      if (Doc.Minidump)
-        return BoolToErr(yaml2minidump(*Doc.Minidump, Out));
-      if (Doc.Wasm)
-        return BoolToErr(yaml2wasm(*Doc.Wasm, Out));
-      return createStringError(errc::invalid_argument,
-                               "Unknown document type!");
+    if (++CurDocNum != DocNum)
+      continue;
+
+    yaml::YamlObjectFile Doc;
+    YIn >> Doc;
+    if (std::error_code EC = YIn.error()) {
+      ErrHandler("failed to parse YAML input: " + EC.message());
+      return false;
     }
+
+    if (Doc.Elf)
+      return yaml2elf(*Doc.Elf, Out, ErrHandler);
+    if (Doc.Coff)
+      return yaml2coff(*Doc.Coff, Out, ErrHandler);
+    if (Doc.MachO || Doc.FatMachO)
+      return yaml2macho(Doc, Out, ErrHandler);
+    if (Doc.Minidump)
+      return yaml2minidump(*Doc.Minidump, Out, ErrHandler);
+    if (Doc.Wasm)
+      return yaml2wasm(*Doc.Wasm, Out, ErrHandler);
+
+    ErrHandler("unknown document type");
+    return false;
+
   } while (YIn.nextDocument());
 
-  return createStringError(errc::invalid_argument,
-                           "Cannot find the %u%s document", DocNum,
-                           getOrdinalSuffix(DocNum).data());
+  ErrHandler("cannot find the " + Twine(DocNum) +
+             getOrdinalSuffix(DocNum).data() + " document");
+  return false;
 }
 
-Expected<std::unique_ptr<object::ObjectFile>>
-yaml2ObjectFile(SmallVectorImpl<char> &Storage, StringRef Yaml) {
+std::unique_ptr<object::ObjectFile>
+yaml2ObjectFile(SmallVectorImpl<char> &Storage, StringRef Yaml,
+                ErrorHandler ErrHandler) {
   Storage.clear();
   raw_svector_ostream OS(Storage);
 
   yaml::Input YIn(Yaml);
-  if (Error E = convertYAML(YIn, OS))
-    return std::move(E);
+  if (!convertYAML(YIn, OS, ErrHandler))
+    return {};
+
+  Expected<std::unique_ptr<object::ObjectFile>> ObjOrErr =
+      object::ObjectFile::createObjectFile(
+          MemoryBufferRef(OS.str(), "YamlObject"));
+  if (ObjOrErr)
+    return std::move(*ObjOrErr);
 
-  return object::ObjectFile::createObjectFile(
-      MemoryBufferRef(OS.str(), "YamlObject"));
+  ErrHandler(toString(ObjOrErr.takeError()));
+  return {};
 }
 
 } // namespace yaml