4 files changed, 94 insertions, 18 deletions

diff --git a/include/clang/Frontend/ASTUnit.h b/include/clang/Frontend/ASTUnit.h
index 2d38352..fa4bcf2 100644
--- a/include/clang/Frontend/ASTUnit.h
+++ b/include/clang/Frontend/ASTUnit.h
@@ -57,6 +57,7 @@ class FileManager;
 class HeaderSearch;
 class Preprocessor;
 class PCHContainerOperations;
+class PCHContainerReader;
 class SourceManager;
 class TargetInfo;
 class ASTFrontendAction;
@@ -725,8 +726,7 @@ public:
   ///
   /// \returns - The initialized ASTUnit or null if the AST failed to load.
   static std::unique_ptr<ASTUnit> LoadFromASTFile(
-      const std::string &Filename,
-      std::shared_ptr<PCHContainerOperations> PCHContainerOps,
+      const std::string &Filename, const PCHContainerReader &PCHContainerRdr,
       IntrusiveRefCntPtr<DiagnosticsEngine> Diags,
       const FileSystemOptions &FileSystemOpts, bool OnlyLocalDecls = false,
       ArrayRef<RemappedFile> RemappedFiles = None,
diff --git a/include/clang/Frontend/CompilerInstance.h b/include/clang/Frontend/CompilerInstance.h
index 2f3e1b6..45e5ed1 100644
--- a/include/clang/Frontend/CompilerInstance.h
+++ b/include/clang/Frontend/CompilerInstance.h
@@ -183,7 +183,7 @@ class CompilerInstance : public ModuleLoader {
 public:
   explicit CompilerInstance(
       std::shared_ptr<PCHContainerOperations> PCHContainerOps =
-          std::make_shared<RawPCHContainerOperations>(),
+          std::make_shared<PCHContainerOperations>(),
       bool BuildingModule = false);
   ~CompilerInstance() override;
 
@@ -508,6 +508,34 @@ public:
     return ThePCHContainerOperations;
   }
 
+  /// Return the appropriate PCHContainerWriter depending on the
+  /// current CodeGenOptions.
+  const PCHContainerWriter &getPCHContainerWriter() const {
+    assert(Invocation && "cannot determine module format without invocation");
+    StringRef Format = getHeaderSearchOpts().ModuleFormat;
+    auto *Writer = ThePCHContainerOperations->getWriterOrNull(Format);
+    if (!Writer) {
+      if (Diagnostics)
+        Diagnostics->Report(diag::err_module_format_unhandled) << Format;
+      llvm::report_fatal_error("unknown module format");
+    }
+    return *Writer;
+  }
+
+  /// Return the appropriate PCHContainerReader depending on the
+  /// current CodeGenOptions.
+  const PCHContainerReader &getPCHContainerReader() const {
+    assert(Invocation && "cannot determine module format without invocation");
+    StringRef Format = getHeaderSearchOpts().ModuleFormat;
+    auto *Reader = ThePCHContainerOperations->getReaderOrNull(Format);
+    if (!Reader) {
+      if (Diagnostics)
+        Diagnostics->Report(diag::err_module_format_unhandled) << Format;
+      llvm::report_fatal_error("unknown module format");
+    }
+    return *Reader;
+  }
+
   /// }
   /// @name Code Completion
   /// {
@@ -621,7 +649,7 @@ public:
   static IntrusiveRefCntPtr<ASTReader> createPCHExternalASTSource(
       StringRef Path, StringRef Sysroot, bool DisablePCHValidation,
       bool AllowPCHWithCompilerErrors, Preprocessor &PP, ASTContext &Context,
-      const PCHContainerOperations &PCHContainerOps,
+      const PCHContainerReader &PCHContainerRdr,
       void *DeserializationListener, bool OwnDeserializationListener,
       bool Preamble, bool UseGlobalModuleIndex);
 
diff --git a/include/clang/Frontend/PCHContainerOperations.h b/include/clang/Frontend/PCHContainerOperations.h
index 949ee63..868ea68 100644
--- a/include/clang/Frontend/PCHContainerOperations.h
+++ b/include/clang/Frontend/PCHContainerOperations.h
@@ -11,6 +11,7 @@
 #define LLVM_CLANG_PCH_CONTAINER_OPERATIONS_H
 
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include <memory>
 
@@ -19,6 +20,8 @@ class raw_pwrite_stream;
 class BitstreamReader;
 }
 
+using llvm::StringRef;
+
 namespace clang {
 
 class ASTConsumer;
@@ -33,14 +36,16 @@ struct PCHBuffer {
   bool IsComplete;
   llvm::SmallVector<char, 0> Data;
 };
+  
+/// This abstract interface provides operations for creating
+/// containers for serialized ASTs (precompiled headers and clang
+/// modules).
+class PCHContainerWriter {
+public: 
+  virtual ~PCHContainerWriter() = 0;
+  virtual StringRef getFormat() const = 0;
 
-/// \brief This abstract interface provides operations for creating
-/// and unwrapping containers for serialized ASTs (precompiled headers
-/// and clang modules).
-class PCHContainerOperations {
-public:
-  virtual ~PCHContainerOperations();
-  /// \brief Return an ASTConsumer that can be chained with a
+  /// Return an ASTConsumer that can be chained with a
   /// PCHGenerator that produces a wrapper file format containing a
   /// serialized AST bitstream.
   virtual std::unique_ptr<ASTConsumer> CreatePCHContainerGenerator(
@@ -49,16 +54,28 @@ public:
       const LangOptions &LO, const std::string &MainFileName,
       const std::string &OutputFileName, llvm::raw_pwrite_stream *OS,
       std::shared_ptr<PCHBuffer> Buffer) const = 0;
+};
 
-  /// \brief Initialize an llvm::BitstreamReader with the serialized AST inside
+/// This abstract interface provides operations for unwrapping
+/// containers for serialized ASTs (precompiled headers and clang
+/// modules).
+class PCHContainerReader {
+public: 
+  virtual ~PCHContainerReader() = 0;
+  /// Equivalent to the format passed to -fmodule-format=
+  virtual StringRef getFormat() const = 0;
+
+  /// Initialize an llvm::BitstreamReader with the serialized AST inside
   /// the PCH container Buffer.
   virtual void ExtractPCH(llvm::MemoryBufferRef Buffer,
                           llvm::BitstreamReader &StreamFile) const = 0;
 };
 
-/// \brief Implements a raw pass-through PCH container.
-class RawPCHContainerOperations : public PCHContainerOperations {
-  /// \brief Return an ASTConsumer that can be chained with a
+/// Implements write operations for a raw pass-through PCH container.
+class RawPCHContainerWriter : public PCHContainerWriter {
+  StringRef getFormat() const override { return "raw"; }
+
+  /// Return an ASTConsumer that can be chained with a
   /// PCHGenerator that writes the module to a flat file.
   std::unique_ptr<ASTConsumer> CreatePCHContainerGenerator(
       DiagnosticsEngine &Diags, const HeaderSearchOptions &HSO,
@@ -66,11 +83,42 @@ class RawPCHContainerOperations : public PCHContainerOperations {
       const LangOptions &LO, const std::string &MainFileName,
       const std::string &OutputFileName, llvm::raw_pwrite_stream *OS,
       std::shared_ptr<PCHBuffer> Buffer) const override;
+};
 
-  /// \brief Initialize an llvm::BitstreamReader with Buffer.
+/// Implements read operations for a raw pass-through PCH container.
+class RawPCHContainerReader : public PCHContainerReader {
+  StringRef getFormat() const override { return "raw"; }
+
+  /// Initialize an llvm::BitstreamReader with Buffer.
   void ExtractPCH(llvm::MemoryBufferRef Buffer,
                   llvm::BitstreamReader &StreamFile) const override;
 };
+
+/// A registry of PCHContainerWriter and -Reader objects for different formats.
+class PCHContainerOperations {
+  llvm::StringMap<std::unique_ptr<PCHContainerWriter>> Writers;
+  llvm::StringMap<std::unique_ptr<PCHContainerReader>> Readers;
+public:
+  /// Automatically registers a RawPCHContainerWriter and
+  /// RawPCHContainerReader.
+  PCHContainerOperations();
+  void registerWriter(std::unique_ptr<PCHContainerWriter> Writer) {
+    Writers[Writer->getFormat()] = std::move(Writer);
+  }
+  void registerReader(std::unique_ptr<PCHContainerReader> Reader) {
+    Readers[Reader->getFormat()] = std::move(Reader);
+  }  
+  const PCHContainerWriter *getWriterOrNull(StringRef Format) {
+    return Writers[Format].get();
+  }
+  const PCHContainerReader *getReaderOrNull(StringRef Format) {
+    return Readers[Format].get();
+  }
+  const PCHContainerReader &getRawReader() {
+    return *getReaderOrNull("raw");
+  }
+};
+
 }
 
 #endif
diff --git a/include/clang/Frontend/Utils.h b/include/clang/Frontend/Utils.h
index 6399653..aa567b4 100644
--- a/include/clang/Frontend/Utils.h
+++ b/include/clang/Frontend/Utils.h
@@ -45,7 +45,7 @@ class HeaderSearch;
 class HeaderSearchOptions;
 class IdentifierTable;
 class LangOptions;
-class PCHContainerOperations;
+class PCHContainerReader;
 class Preprocessor;
 class PreprocessorOptions;
 class PreprocessorOutputOptions;
@@ -63,7 +63,7 @@ void ApplyHeaderSearchOptions(HeaderSearch &HS,
 /// InitializePreprocessor - Initialize the preprocessor getting it and the
 /// environment ready to process a single file.
 void InitializePreprocessor(Preprocessor &PP, const PreprocessorOptions &PPOpts,
-                            const PCHContainerOperations &PCHContainerOps,
+                            const PCHContainerReader &PCHContainerRdr,
                             const FrontendOptions &FEOpts);
 
 /// DoPrintPreprocessedInput - Implement -E mode.