2 files changed, 81 insertions, 15 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);
 }
diff --git a/llvm/lib/CodeGen/MIRPrintingPass.cpp b/llvm/lib/CodeGen/MIRPrintingPass.cpp
index c66658cc977..5e0f4cdcbfd 100644
--- a/llvm/lib/CodeGen/MIRPrintingPass.cpp
+++ b/llvm/lib/CodeGen/MIRPrintingPass.cpp
@@ -15,6 +15,7 @@
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MIRYamlMapping.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
@@ -41,11 +42,30 @@ template <> struct BlockScalarTraits<Module> {
 
 namespace {
 
+/// This class prints out the machine functions using the MIR serialization
+/// format.
+class MIRPrinter {
+  raw_ostream &OS;
+
+public:
+  MIRPrinter(raw_ostream &OS) : OS(OS) {}
+
+  void print(const MachineFunction &MF);
+};
+
+void MIRPrinter::print(const MachineFunction &MF) {
+  yaml::MachineFunction YamlMF;
+  YamlMF.Name = MF.getName();
+  yaml::Output Out(OS);
+  Out << YamlMF;
+}
+
 /// This pass prints out the LLVM IR to an output stream using the MIR
 /// serialization format.
 struct MIRPrintingPass : public MachineFunctionPass {
   static char ID;
   raw_ostream &OS;
+  std::string MachineFunctions;
 
   MIRPrintingPass() : MachineFunctionPass(ID), OS(dbgs()) {}
   MIRPrintingPass(raw_ostream &OS) : MachineFunctionPass(ID), OS(OS) {}
@@ -58,13 +78,17 @@ struct MIRPrintingPass : public MachineFunctionPass {
   }
 
   virtual bool runOnMachineFunction(MachineFunction &MF) override {
-    // TODO: Print out the machine function.
+    std::string Str;
+    raw_string_ostream StrOS(Str);
+    MIRPrinter(StrOS).print(MF);
+    MachineFunctions.append(StrOS.str());
     return false;
   }
 
   virtual bool doFinalization(Module &M) override {
     yaml::Output Out(OS);
     Out << M;
+    OS << MachineFunctions;
     return false;
   }
 };