MIR Serialization: Connect the machine function analysis pass to the MIR parser.

This commit connects the machine function analysis pass (which creates machine
functions) to the MIR parser, which will initialize the machine functions 
with the state from the MIR file and reconstruct the machine IR.

This commit introduces a new interface called 'MachineFunctionInitializer',
which can be used to provide custom initialization for the machine functions.

This commit also introduces a new diagnostic class called 
'DiagnosticInfoMIRParser' which is used for MIR parsing errors.
This commit modifies the default diagnostic handling in LLVMContext - now the
the diagnostics are printed directly into llvm::errs() so that the MIR parsing 
errors can be printed with colours.  

Reviewers: Justin Bogner

Differential Revision: http://reviews.llvm.org/D9928

llvm-svn: 239753
diff --git a/llvm/lib/CodeGen/LLVMTargetMachine.cpp b/llvm/lib/CodeGen/LLVMTargetMachine.cpp
index ef44fda..b7cac3b 100644
--- a/llvm/lib/CodeGen/LLVMTargetMachine.cpp
+++ b/llvm/lib/CodeGen/LLVMTargetMachine.cpp
@@ -87,11 +87,11 @@
 }
 
 /// addPassesToX helper drives creation and initialization of TargetPassConfig.
-static MCContext *addPassesToGenerateCode(LLVMTargetMachine *TM,
-                                          PassManagerBase &PM,
-                                          bool DisableVerify,
-                                          AnalysisID StartAfter,
-                                          AnalysisID StopAfter) {
+static MCContext *
+addPassesToGenerateCode(LLVMTargetMachine *TM, PassManagerBase &PM,
+                        bool DisableVerify, AnalysisID StartAfter,
+                        AnalysisID StopAfter,
+                        MachineFunctionInitializer *MFInitializer = nullptr) {
 
   // Add internal analysis passes from the target machine.
   PM.add(createTargetTransformInfoWrapperPass(TM->getTargetIRAnalysis()));
@@ -121,7 +121,7 @@
   PM.add(MMI);
 
   // Set up a MachineFunction for the rest of CodeGen to work on.
-  PM.add(new MachineFunctionAnalysis(*TM));
+  PM.add(new MachineFunctionAnalysis(*TM, MFInitializer));
 
   // Enable FastISel with -fast, but allow that to be overridden.
   if (EnableFastISelOption == cl::BOU_TRUE ||
@@ -142,10 +142,11 @@
 
 bool LLVMTargetMachine::addPassesToEmitFile(
     PassManagerBase &PM, raw_pwrite_stream &Out, CodeGenFileType FileType,
-    bool DisableVerify, AnalysisID StartAfter, AnalysisID StopAfter) {
+    bool DisableVerify, AnalysisID StartAfter, AnalysisID StopAfter,
+    MachineFunctionInitializer *MFInitializer) {
   // Add common CodeGen passes.
-  MCContext *Context = addPassesToGenerateCode(this, PM, DisableVerify,
-                                               StartAfter, StopAfter);
+  MCContext *Context = addPassesToGenerateCode(
+      this, PM, DisableVerify, StartAfter, StopAfter, MFInitializer);
   if (!Context)
     return true;
 
diff --git a/llvm/lib/CodeGen/MIRParser/MIRParser.cpp b/llvm/lib/CodeGen/MIRParser/MIRParser.cpp
index 7a51b38..69ecb92 100644
--- a/llvm/lib/CodeGen/MIRParser/MIRParser.cpp
+++ b/llvm/lib/CodeGen/MIRParser/MIRParser.cpp
@@ -14,9 +14,13 @@
 
 #include "llvm/CodeGen/MIRParser/MIRParser.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/AsmParser/Parser.h"
+#include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MIRYamlMapping.h"
+#include "llvm/IR/DiagnosticInfo.h"
+#include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/LineIterator.h"
 #include "llvm/Support/SMLoc.h"
@@ -27,7 +31,7 @@
 
 using namespace llvm;
 
-namespace {
+namespace llvm {
 
 /// This class implements the parsing of LLVM IR that's embedded inside a MIR
 /// file.
@@ -35,29 +39,43 @@
   SourceMgr SM;
   StringRef Filename;
   LLVMContext &Context;
+  StringMap<std::unique_ptr<yaml::MachineFunction>> Functions;
 
 public:
   MIRParserImpl(std::unique_ptr<MemoryBuffer> Contents, StringRef Filename,
                 LLVMContext &Context);
 
+  void reportDiagnostic(const SMDiagnostic &Diag);
+
+  /// Report an error with the given message at unknown location.
+  ///
+  /// Always returns true.
+  bool error(const Twine &Message);
+
   /// Try to parse the optional LLVM module and the machine functions in the MIR
   /// file.
   ///
   /// Return null if an error occurred.
-  std::unique_ptr<Module> parse(SMDiagnostic &Error);
+  std::unique_ptr<Module> parse();
 
   /// Parse the machine function in the current YAML document.
   ///
   /// Return true if an error occurred.
   bool parseMachineFunction(yaml::Input &In);
 
+  /// Initialize the machine function to the state that's described in the MIR
+  /// file.
+  ///
+  /// Return true if error occurred.
+  bool initializeMachineFunction(MachineFunction &MF);
+
 private:
   /// Return a MIR diagnostic converted from an LLVM assembly diagnostic.
   SMDiagnostic diagFromLLVMAssemblyDiag(const SMDiagnostic &Error,
                                         SMRange SourceRange);
 };
 
-} // end anonymous namespace
+} // end namespace llvm
 
 MIRParserImpl::MIRParserImpl(std::unique_ptr<MemoryBuffer> Contents,
                              StringRef Filename, LLVMContext &Context)
@@ -65,16 +83,38 @@
   SM.AddNewSourceBuffer(std::move(Contents), SMLoc());
 }
 
-static void handleYAMLDiag(const SMDiagnostic &Diag, void *Context) {
-  *reinterpret_cast<SMDiagnostic *>(Context) = Diag;
+bool MIRParserImpl::error(const Twine &Message) {
+  Context.diagnose(DiagnosticInfoMIRParser(
+      DS_Error, SMDiagnostic(Filename, SourceMgr::DK_Error, Message.str())));
+  return true;
 }
 
-std::unique_ptr<Module> MIRParserImpl::parse(SMDiagnostic &Error) {
+void MIRParserImpl::reportDiagnostic(const SMDiagnostic &Diag) {
+  DiagnosticSeverity Kind;
+  switch (Diag.getKind()) {
+  case SourceMgr::DK_Error:
+    Kind = DS_Error;
+    break;
+  case SourceMgr::DK_Warning:
+    Kind = DS_Warning;
+    break;
+  case SourceMgr::DK_Note:
+    Kind = DS_Note;
+    break;
+  }
+  Context.diagnose(DiagnosticInfoMIRParser(Kind, Diag));
+}
+
+static void handleYAMLDiag(const SMDiagnostic &Diag, void *Context) {
+  reinterpret_cast<MIRParserImpl *>(Context)->reportDiagnostic(Diag);
+}
+
+std::unique_ptr<Module> MIRParserImpl::parse() {
   yaml::Input In(SM.getMemoryBuffer(SM.getMainFileID())->getBuffer(),
-                 /*Ctxt=*/nullptr, handleYAMLDiag, &Error);
+                 /*Ctxt=*/nullptr, handleYAMLDiag, this);
 
   if (!In.setCurrentDocument()) {
-    if (!Error.getMessage().empty())
+    if (In.error())
       return nullptr;
     // Create an empty module when the MIR file is empty.
     return llvm::make_unique<Module>(Filename, Context);
@@ -85,10 +125,11 @@
   // without having to go trough YAML traits.
   if (const auto *BSN =
           dyn_cast_or_null<yaml::BlockScalarNode>(In.getCurrentNode())) {
+    SMDiagnostic Error;
     M = parseAssembly(MemoryBufferRef(BSN->getValue(), Filename), Error,
                       Context);
     if (!M) {
-      Error = diagFromLLVMAssemblyDiag(Error, BSN->getSourceRange());
+      reportDiagnostic(diagFromLLVMAssemblyDiag(Error, BSN->getSourceRange()));
       return M;
     }
     In.nextDocument();
@@ -110,12 +151,21 @@
 }
 
 bool MIRParserImpl::parseMachineFunction(yaml::Input &In) {
-  yaml::MachineFunction MF;
-  yaml::yamlize(In, MF, false);
+  auto MF = llvm::make_unique<yaml::MachineFunction>();
+  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.
+  auto FunctionName = MF->Name;
+  Functions.insert(std::make_pair(FunctionName, std::move(MF)));
+  return false;
+}
+
+bool MIRParserImpl::initializeMachineFunction(MachineFunction &MF) {
+  auto It = Functions.find(MF.getName());
+  if (It == Functions.end())
+    return error(Twine("no machine function information for function '") +
+                 MF.getName() + "' in the MIR file");
+  // TODO: Recreate the machine function.
   return false;
 }
 
@@ -150,22 +200,33 @@
                       Error.getFixIts());
 }
 
-std::unique_ptr<Module> llvm::parseMIRFile(StringRef Filename,
-                                           SMDiagnostic &Error,
-                                           LLVMContext &Context) {
+MIRParser::MIRParser(std::unique_ptr<MIRParserImpl> Impl)
+    : Impl(std::move(Impl)) {}
+
+MIRParser::~MIRParser() {}
+
+std::unique_ptr<Module> MIRParser::parseLLVMModule() { return Impl->parse(); }
+
+bool MIRParser::initializeMachineFunction(MachineFunction &MF) {
+  return Impl->initializeMachineFunction(MF);
+}
+
+std::unique_ptr<MIRParser> llvm::createMIRParserFromFile(StringRef Filename,
+                                                         SMDiagnostic &Error,
+                                                         LLVMContext &Context) {
   auto FileOrErr = MemoryBuffer::getFile(Filename);
   if (std::error_code EC = FileOrErr.getError()) {
     Error = SMDiagnostic(Filename, SourceMgr::DK_Error,
                          "Could not open input file: " + EC.message());
-    return std::unique_ptr<Module>();
+    return nullptr;
   }
-  return parseMIR(std::move(FileOrErr.get()), Error, Context);
+  return createMIRParser(std::move(FileOrErr.get()), Context);
 }
 
-std::unique_ptr<Module> llvm::parseMIR(std::unique_ptr<MemoryBuffer> Contents,
-                                       SMDiagnostic &Error,
-                                       LLVMContext &Context) {
+std::unique_ptr<MIRParser>
+llvm::createMIRParser(std::unique_ptr<MemoryBuffer> Contents,
+                      LLVMContext &Context) {
   auto Filename = Contents->getBufferIdentifier();
-  MIRParserImpl Parser(std::move(Contents), Filename, Context);
-  return Parser.parse(Error);
+  return llvm::make_unique<MIRParser>(
+      llvm::make_unique<MIRParserImpl>(std::move(Contents), Filename, Context));
 }
diff --git a/llvm/lib/CodeGen/MachineFunction.cpp b/llvm/lib/CodeGen/MachineFunction.cpp
index addb9e8..960e06a 100644
--- a/llvm/lib/CodeGen/MachineFunction.cpp
+++ b/llvm/lib/CodeGen/MachineFunction.cpp
@@ -19,6 +19,7 @@
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/CodeGen/MachineConstantPool.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
+#include "llvm/CodeGen/MachineFunctionInitializer.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineJumpTableInfo.h"
@@ -41,6 +42,8 @@
 
 #define DEBUG_TYPE "codegen"
 
+void MachineFunctionInitializer::anchor() {}
+
 //===----------------------------------------------------------------------===//
 // MachineFunction implementation
 //===----------------------------------------------------------------------===//
diff --git a/llvm/lib/CodeGen/MachineFunctionAnalysis.cpp b/llvm/lib/CodeGen/MachineFunctionAnalysis.cpp
index f6f34ba..338cd1e 100644
--- a/llvm/lib/CodeGen/MachineFunctionAnalysis.cpp
+++ b/llvm/lib/CodeGen/MachineFunctionAnalysis.cpp
@@ -15,12 +15,14 @@
 #include "llvm/CodeGen/GCMetadata.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/CodeGen/MachineFunctionInitializer.h"
 using namespace llvm;
 
 char MachineFunctionAnalysis::ID = 0;
 
-MachineFunctionAnalysis::MachineFunctionAnalysis(const TargetMachine &tm) :
-  FunctionPass(ID), TM(tm), MF(nullptr) {
+MachineFunctionAnalysis::MachineFunctionAnalysis(
+    const TargetMachine &tm, MachineFunctionInitializer *MFInitializer)
+    : FunctionPass(ID), TM(tm), MF(nullptr), MFInitializer(MFInitializer) {
   initializeMachineModuleInfoPass(*PassRegistry::getPassRegistry());
 }
 
@@ -47,6 +49,8 @@
   assert(!MF && "MachineFunctionAnalysis already initialized!");
   MF = new MachineFunction(&F, TM, NextFnNum++,
                            getAnalysis<MachineModuleInfo>());
+  if (MFInitializer)
+    MFInitializer->initializeMachineFunction(*MF);
   return false;
 }