1 files changed, 56 insertions, 14 deletions

diff --git a/llvm/lib/CodeGen/MIRParser/MIRParser.cpp b/llvm/lib/CodeGen/MIRParser/MIRParser.cpp
index 8a048c1dfb8..a0de3715bb5 100644
--- a/llvm/lib/CodeGen/MIRParser/MIRParser.cpp
+++ b/llvm/lib/CodeGen/MIRParser/MIRParser.cpp
@@ -16,6 +16,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/AsmParser/Parser.h"
+#include "llvm/CodeGen/MIRYamlMapping.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/SMLoc.h"
 #include "llvm/Support/SourceMgr.h"
@@ -38,10 +39,16 @@ public:
   MIRParserImpl(std::unique_ptr<MemoryBuffer> Contents, StringRef Filename,
                 LLVMContext &Context);
 
-  /// Try to parse the optional LLVM module in the MIR file.
+  /// Try to parse the optional LLVM module and the machine functions in the MIR
+  /// file.
   ///
-  /// Return null if an error occurred while parsing the LLVM module.
-  std::unique_ptr<Module> parseLLVMModule(SMDiagnostic &Error);
+  /// Return null if an error occurred.
+  std::unique_ptr<Module> parse(SMDiagnostic &Error);
+
+  /// Parse the machine function in the current YAML document.
+  ///
+  /// Return true if an error occurred.
+  bool parseMachineFunction(yaml::Input &In);
 };
 
 } // end anonymous namespace
@@ -52,21 +59,56 @@ MIRParserImpl::MIRParserImpl(std::unique_ptr<MemoryBuffer> Contents,
   SM.AddNewSourceBuffer(std::move(Contents), SMLoc());
 }
 
-std::unique_ptr<Module> MIRParserImpl::parseLLVMModule(SMDiagnostic &Error) {
-  yaml::Input In(SM.getMemoryBuffer(SM.getMainFileID())->getBuffer());
+static void handleYAMLDiag(const SMDiagnostic &Diag, void *Context) {
+  *reinterpret_cast<SMDiagnostic *>(Context) = Diag;
+}
 
+std::unique_ptr<Module> MIRParserImpl::parse(SMDiagnostic &Error) {
+  yaml::Input In(SM.getMemoryBuffer(SM.getMainFileID())->getBuffer(),
+                 /*Ctxt=*/nullptr, handleYAMLDiag, &Error);
+
+  if (!In.setCurrentDocument()) {
+    if (!Error.getMessage().empty())
+      return nullptr;
+    // Create an empty module when the MIR file is empty.
+    return llvm::make_unique<Module>(Filename, Context);
+  }
+
+  std::unique_ptr<Module> M;
   // Parse the block scalar manually so that we can return unique pointer
   // without having to go trough YAML traits.
-  if (In.setCurrentDocument()) {
-    if (const auto *BSN =
-            dyn_cast_or_null<yaml::BlockScalarNode>(In.getCurrentNode())) {
-      return parseAssembly(MemoryBufferRef(BSN->getValue(), Filename), Error,
-                           Context);
-    }
+  if (const auto *BSN =
+          dyn_cast_or_null<yaml::BlockScalarNode>(In.getCurrentNode())) {
+    M = parseAssembly(MemoryBufferRef(BSN->getValue(), Filename), Error,
+                      Context);
+    if (!M)
+      return M;
+    In.nextDocument();
+    if (!In.setCurrentDocument())
+      return M;
+  } else {
+    // Create an new, empty module.
+    M = llvm::make_unique<Module>(Filename, Context);
   }
 
-  // Create an new, empty module.
-  return llvm::make_unique<Module>(Filename, Context);
+  // Parse the machine functions.
+  do {
+    if (parseMachineFunction(In))
+      return nullptr;
+    In.nextDocument();
+  } while (In.setCurrentDocument());
+
+  return M;
+}
+
+bool MIRParserImpl::parseMachineFunction(yaml::Input &In) {
+  yaml::MachineFunction MF;
+  yaml::yamlize(In, MF, false);
+  if (In.error())
+    return true;
+  // TODO: Initialize the real machine function with the state in the yaml
+  // machine function later on.
+  return false;
 }
 
 std::unique_ptr<Module> llvm::parseMIRFile(StringRef Filename,
@@ -86,5 +128,5 @@ std::unique_ptr<Module> llvm::parseMIR(std::unique_ptr<MemoryBuffer> Contents,
                                        LLVMContext &Context) {
   auto Filename = Contents->getBufferIdentifier();
   MIRParserImpl Parser(std::move(Contents), Filename, Context);
-  return Parser.parseLLVMModule(Error);
+  return Parser.parse(Error);
 }