Introduce module file extensions to piggy-back data onto module files.

Introduce the notion of a module file extension, which introduces
additional information into a module file at the time it is built that
can then be queried when the module file is read. Module file
extensions are identified by a block name (which must be unique to the
extension) and can write any bitstream records into their own
extension block within the module file. When a module file is loaded,
any extension blocks are matched up with module file extension
readers, that are per-module-file and are given access to the input
bitstream.

Note that module file extensions can only be introduced by
programmatic clients that have access to the CompilerInvocation. There
is only one such extension at the moment, which is used for testing
the module file extension harness. As a future direction, one could
imagine allowing the plugin mechanism to introduce new module file
extensions.

llvm-svn: 251955

8 files changed, 287 insertions, 14 deletions

diff --git a/clang/lib/Frontend/ASTUnit.cpp b/clang/lib/Frontend/ASTUnit.cpp
index d3d7185ea3a..6b89447305f 100644
--- a/clang/lib/Frontend/ASTUnit.cpp
+++ b/clang/lib/Frontend/ASTUnit.cpp
@@ -186,7 +186,7 @@ struct ASTUnit::ASTWriterData {
   llvm::BitstreamWriter Stream;
   ASTWriter Writer;
 
-  ASTWriterData() : Stream(Buffer), Writer(Stream) { }
+  ASTWriterData() : Stream(Buffer), Writer(Stream, { }) { }
 };
 
 void ASTUnit::clearFileLevelDecls() {
@@ -709,7 +709,7 @@ std::unique_ptr<ASTUnit> ASTUnit::LoadFromASTFile(
   bool disableValid = false;
   if (::getenv("LIBCLANG_DISABLE_PCH_VALIDATION"))
     disableValid = true;
-  AST->Reader = new ASTReader(PP, Context, PCHContainerRdr,
+  AST->Reader = new ASTReader(PP, Context, PCHContainerRdr, { },
                               /*isysroot=*/"",
                               /*DisableValidation=*/disableValid,
                               AllowPCHWithCompilerErrors);
@@ -927,6 +927,7 @@ public:
                              const Preprocessor &PP, StringRef isysroot,
                              raw_ostream *Out)
       : PCHGenerator(PP, "", nullptr, isysroot, std::make_shared<PCHBuffer>(),
+                     ArrayRef<llvm::IntrusiveRefCntPtr<ModuleFileExtension>>(),
                      /*AllowASTWithErrors=*/true),
         Unit(Unit), Hash(Unit.getCurrentTopLevelHashValue()), Action(Action),
         Out(Out) {
@@ -2500,7 +2501,7 @@ bool ASTUnit::serialize(raw_ostream &OS) {
 
   SmallString<128> Buffer;
   llvm::BitstreamWriter Stream(Buffer);
-  ASTWriter Writer(Stream);
+  ASTWriter Writer(Stream, { });
   return serializeUnit(Writer, Buffer, getSema(), hasErrors, OS);
 }
 
diff --git a/clang/lib/Frontend/CMakeLists.txt b/clang/lib/Frontend/CMakeLists.txt
index 9a3e459640a..af42a905b38 100644
--- a/clang/lib/Frontend/CMakeLists.txt
+++ b/clang/lib/Frontend/CMakeLists.txt
@@ -35,6 +35,7 @@ add_clang_library(clangFrontend
   PrintPreprocessedOutput.cpp
   SerializedDiagnosticPrinter.cpp
   SerializedDiagnosticReader.cpp
+  TestModuleFileExtension.cpp
   TextDiagnostic.cpp
   TextDiagnosticBuffer.cpp
   TextDiagnosticPrinter.cpp
diff --git a/clang/lib/Frontend/ChainedIncludesSource.cpp b/clang/lib/Frontend/ChainedIncludesSource.cpp
index d36162b8dab..1c1081fbe08 100644
--- a/clang/lib/Frontend/ChainedIncludesSource.cpp
+++ b/clang/lib/Frontend/ChainedIncludesSource.cpp
@@ -82,6 +82,7 @@ createASTReader(CompilerInstance &CI, StringRef pchFile,
   std::unique_ptr<ASTReader> Reader;
   Reader.reset(new ASTReader(PP, CI.getASTContext(),
                              CI.getPCHContainerReader(),
+                             /*Extensions=*/{ },
                              /*isysroot=*/"", /*DisableValidation=*/true));
   for (unsigned ti = 0; ti < bufNames.size(); ++ti) {
     StringRef sr(bufNames[ti]);
@@ -160,8 +161,10 @@ IntrusiveRefCntPtr<ExternalSemaSource> clang::createChainedIncludesSource(
     Clang->createASTContext();
 
     auto Buffer = std::make_shared<PCHBuffer>();
+    ArrayRef<llvm::IntrusiveRefCntPtr<ModuleFileExtension>> Extensions;
     auto consumer = llvm::make_unique<PCHGenerator>(
-        Clang->getPreprocessor(), "-", nullptr, /*isysroot=*/"", Buffer);
+        Clang->getPreprocessor(), "-", nullptr, /*isysroot=*/"", Buffer,
+        Extensions);
     Clang->getASTContext().setASTMutationListener(
                                             consumer->GetASTMutationListener());
     Clang->setASTConsumer(std::move(consumer));
diff --git a/clang/lib/Frontend/CompilerInstance.cpp b/clang/lib/Frontend/CompilerInstance.cpp
index 8220a6e4925..c3f19a3a00d 100644
--- a/clang/lib/Frontend/CompilerInstance.cpp
+++ b/clang/lib/Frontend/CompilerInstance.cpp
@@ -408,7 +408,9 @@ void CompilerInstance::createPCHExternalASTSource(
   ModuleManager = createPCHExternalASTSource(
       Path, getHeaderSearchOpts().Sysroot, DisablePCHValidation,
       AllowPCHWithCompilerErrors, getPreprocessor(), getASTContext(),
-      getPCHContainerReader(), DeserializationListener,
+      getPCHContainerReader(),
+      getFrontendOpts().ModuleFileExtensions,
+      DeserializationListener,
       OwnDeserializationListener, Preamble,
       getFrontendOpts().UseGlobalModuleIndex);
 }
@@ -417,15 +419,16 @@ IntrusiveRefCntPtr<ASTReader> CompilerInstance::createPCHExternalASTSource(
     StringRef Path, StringRef Sysroot, bool DisablePCHValidation,
     bool AllowPCHWithCompilerErrors, Preprocessor &PP, ASTContext &Context,
     const PCHContainerReader &PCHContainerRdr,
+    ArrayRef<IntrusiveRefCntPtr<ModuleFileExtension>> Extensions,
     void *DeserializationListener, bool OwnDeserializationListener,
     bool Preamble, bool UseGlobalModuleIndex) {
   HeaderSearchOptions &HSOpts = PP.getHeaderSearchInfo().getHeaderSearchOpts();
 
   IntrusiveRefCntPtr<ASTReader> Reader(new ASTReader(
-      PP, Context, PCHContainerRdr, Sysroot.empty() ? "" : Sysroot.data(),
-      DisablePCHValidation, AllowPCHWithCompilerErrors,
-      /*AllowConfigurationMismatch*/ false, HSOpts.ModulesValidateSystemHeaders,
-      UseGlobalModuleIndex));
+      PP, Context, PCHContainerRdr, Extensions,
+      Sysroot.empty() ? "" : Sysroot.data(), DisablePCHValidation,
+      AllowPCHWithCompilerErrors, /*AllowConfigurationMismatch*/ false,
+      HSOpts.ModulesValidateSystemHeaders, UseGlobalModuleIndex));
 
   // We need the external source to be set up before we read the AST, because
   // eagerly-deserialized declarations may use it.
@@ -1267,6 +1270,7 @@ void CompilerInstance::createModuleManager() {
                                                  *FrontendTimerGroup);
     ModuleManager = new ASTReader(
         getPreprocessor(), getASTContext(), getPCHContainerReader(),
+        getFrontendOpts().ModuleFileExtensions,
         Sysroot.empty() ? "" : Sysroot.c_str(), PPOpts.DisablePCHValidation,
         /*AllowASTWithCompilerErrors=*/false,
         /*AllowConfigurationMismatch=*/false,
diff --git a/clang/lib/Frontend/CompilerInvocation.cpp b/clang/lib/Frontend/CompilerInvocation.cpp
index 0c1c4eaefa9..468cc5746ce 100644
--- a/clang/lib/Frontend/CompilerInvocation.cpp
+++ b/clang/lib/Frontend/CompilerInvocation.cpp
@@ -7,6 +7,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "TestModuleFileExtension.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Basic/FileManager.h"
 #include "clang/Basic/Version.h"
@@ -19,6 +20,7 @@
 #include "clang/Frontend/Utils.h"
 #include "clang/Lex/HeaderSearchOptions.h"
 #include "clang/Serialization/ASTReader.h"
+#include "clang/Serialization/ModuleFileExtension.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
@@ -832,6 +834,30 @@ static void ParseFileSystemArgs(FileSystemOptions &Opts, ArgList &Args) {
   Opts.WorkingDir = Args.getLastArgValue(OPT_working_directory);
 }
 
+/// Parse the argument to the -ftest-module-file-extension
+/// command-line argument.
+///
+/// \returns true on error, false on success.
+static bool parseTestModuleFileExtensionArg(StringRef Arg,
+                                            std::string &BlockName,
+                                            unsigned &MajorVersion,
+                                            unsigned &MinorVersion,
+                                            bool &Hashed,
+                                            std::string &UserInfo) {
+  SmallVector<StringRef, 5> Args;
+  Arg.split(Args, ':', 5);
+  if (Args.size() < 5)
+    return true;
+
+  BlockName = Args[0];
+  if (Args[1].getAsInteger(10, MajorVersion)) return true;
+  if (Args[2].getAsInteger(10, MinorVersion)) return true;
+  if (Args[3].getAsInteger(2, Hashed)) return true;
+  if (Args.size() > 4)
+    UserInfo = Args[4];
+  return false;
+}
+
 static InputKind ParseFrontendArgs(FrontendOptions &Opts, ArgList &Args,
                                    DiagnosticsEngine &Diags) {
   using namespace options;
@@ -924,6 +950,26 @@ static InputKind ParseFrontendArgs(FrontendOptions &Opts, ArgList &Args,
       if (A->getValue(0) == Opts.AddPluginActions[i])
         Opts.AddPluginArgs[i].emplace_back(A->getValue(1));
 
+  for (const std::string &Arg :
+         Args.getAllArgValues(OPT_ftest_module_file_extension_EQ)) {
+    std::string BlockName;
+    unsigned MajorVersion;
+    unsigned MinorVersion;
+    bool Hashed;
+    std::string UserInfo;
+    if (parseTestModuleFileExtensionArg(Arg, BlockName, MajorVersion,
+                                        MinorVersion, Hashed, UserInfo)) {
+      Diags.Report(diag::err_test_module_file_extension_format) << Arg;
+
+      continue;
+    }
+
+    // Add the testing module file extension.
+    Opts.ModuleFileExtensions.push_back(
+      new TestModuleFileExtension(BlockName, MajorVersion, MinorVersion,
+                                  Hashed, UserInfo));
+  }
+
   if (const Arg *A = Args.getLastArg(OPT_code_completion_at)) {
     Opts.CodeCompletionAt =
       ParsedSourceLocation::FromString(A->getValue());
@@ -2076,6 +2122,12 @@ std::string CompilerInvocation::getModuleHash() const {
   // Extend the signature with the user build path.
   code = hash_combine(code, hsOpts.ModuleUserBuildPath);
 
+  // Extend the signature with the module file extensions.
+  const FrontendOptions &frontendOpts = getFrontendOpts();
+  for (auto ext : frontendOpts.ModuleFileExtensions) {
+    code = ext->hashExtension(code);
+  }
+
   // Darwin-specific hack: if we have a sysroot, use the contents and
   // modification time of
   //   $sysroot/System/Library/CoreServices/SystemVersion.plist
diff --git a/clang/lib/Frontend/FrontendActions.cpp b/clang/lib/Frontend/FrontendActions.cpp
index 6ba4b134c38..865fb474020 100644
--- a/clang/lib/Frontend/FrontendActions.cpp
+++ b/clang/lib/Frontend/FrontendActions.cpp
@@ -91,7 +91,8 @@ GeneratePCHAction::CreateASTConsumer(CompilerInstance &CI, StringRef InFile) {
   auto Buffer = std::make_shared<PCHBuffer>();
   std::vector<std::unique_ptr<ASTConsumer>> Consumers;
   Consumers.push_back(llvm::make_unique<PCHGenerator>(
-      CI.getPreprocessor(), OutputFile, nullptr, Sysroot, Buffer));
+                        CI.getPreprocessor(), OutputFile, nullptr, Sysroot,
+                        Buffer, CI.getFrontendOpts().ModuleFileExtensions));
   Consumers.push_back(CI.getPCHContainerWriter().CreatePCHContainerGenerator(
       CI, InFile, OutputFile, OS, Buffer));
 
@@ -133,10 +134,13 @@ GenerateModuleAction::CreateASTConsumer(CompilerInstance &CI,
 
   auto Buffer = std::make_shared<PCHBuffer>();
   std::vector<std::unique_ptr<ASTConsumer>> Consumers;
+
   Consumers.push_back(llvm::make_unique<PCHGenerator>(
-      CI.getPreprocessor(), OutputFile, Module, Sysroot, Buffer,
-      /*AllowASTWithErrors*/false,
-      /*IncludeTimestamps*/+CI.getFrontendOpts().BuildingImplicitModule));
+                        CI.getPreprocessor(), OutputFile, Module, Sysroot,
+                        Buffer, CI.getFrontendOpts().ModuleFileExtensions,
+                        /*AllowASTWithErrors=*/false,
+                        /*IncludeTimestamps=*/
+                          +CI.getFrontendOpts().BuildingImplicitModule));
   Consumers.push_back(CI.getPCHContainerWriter().CreatePCHContainerGenerator(
       CI, InFile, OutputFile, OS, Buffer));
   return llvm::make_unique<MultiplexConsumer>(std::move(Consumers));
@@ -421,6 +425,7 @@ void VerifyPCHAction::ExecuteAction() {
   const std::string &Sysroot = CI.getHeaderSearchOpts().Sysroot;
   std::unique_ptr<ASTReader> Reader(new ASTReader(
       CI.getPreprocessor(), CI.getASTContext(), CI.getPCHContainerReader(),
+      CI.getFrontendOpts().ModuleFileExtensions,
       Sysroot.empty() ? "" : Sysroot.c_str(),
       /*DisableValidation*/ false,
       /*AllowPCHWithCompilerErrors*/ false,
@@ -564,6 +569,20 @@ namespace {
       }
       return false;
     }
+
+    /// Indicates that a particular module file extension has been read.
+    void readModuleFileExtension(
+           const ModuleFileExtensionMetadata &Metadata) override {
+      Out.indent(2) << "Module file extension '"
+                    << Metadata.BlockName << "' " << Metadata.MajorVersion
+                    << "." << Metadata.MinorVersion;
+      if (!Metadata.UserInfo.empty()) {
+        Out << ": ";
+        Out.write_escaped(Metadata.UserInfo);
+      }
+
+      Out << "\n";
+    }
 #undef DUMP_BOOLEAN
   };
 }
@@ -583,7 +602,8 @@ void DumpModuleInfoAction::ExecuteAction() {
   DumpModuleInfoListener Listener(Out);
   ASTReader::readASTFileControlBlock(
       getCurrentFile(), getCompilerInstance().getFileManager(),
-      getCompilerInstance().getPCHContainerReader(), Listener);
+      getCompilerInstance().getPCHContainerReader(),
+      /*FindModuleFileExtensions=*/true, Listener);
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/clang/lib/Frontend/TestModuleFileExtension.cpp b/clang/lib/Frontend/TestModuleFileExtension.cpp
new file mode 100644
index 00000000000..e17dc3c1ba6
--- /dev/null
+++ b/clang/lib/Frontend/TestModuleFileExtension.cpp
@@ -0,0 +1,121 @@
+//===-- TestModuleFileExtension.cpp - Module Extension Tester -------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+#include "TestModuleFileExtension.h"
+#include "clang/Frontend/FrontendDiagnostic.h"
+#include "clang/Serialization/ASTReader.h"
+#include "llvm/ADT/Hashing.h"
+#include "llvm/Bitcode/BitstreamWriter.h"
+#include "llvm/Support/raw_ostream.h"
+using namespace clang;
+using namespace clang::serialization;
+
+TestModuleFileExtension::Writer::~Writer() { }
+
+void TestModuleFileExtension::Writer::writeExtensionContents(
+       llvm::BitstreamWriter &Stream) {
+  using namespace llvm;
+
+  // Write an abbreviation for this record.
+  BitCodeAbbrev *Abv = new llvm::BitCodeAbbrev();
+  Abv->Add(BitCodeAbbrevOp(FIRST_EXTENSION_RECORD_ID));
+  Abv->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::VBR, 6)); // # of characters
+  Abv->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::Blob));   // message
+  auto Abbrev = Stream.EmitAbbrev(Abv);
+
+  // Write a message into the extension block.
+  SmallString<64> Message;
+  {
+    auto Ext = static_cast<TestModuleFileExtension *>(getExtension());
+    raw_svector_ostream OS(Message);
+    OS << "Hello from " << Ext->BlockName << " v" << Ext->MajorVersion << "."
+       << Ext->MinorVersion;
+  }
+  SmallVector<uint64_t, 4> Record;
+  Record.push_back(FIRST_EXTENSION_RECORD_ID);
+  Record.push_back(Message.size());
+  Stream.EmitRecordWithBlob(Abbrev, Record, Message);
+}
+
+TestModuleFileExtension::Reader::Reader(ModuleFileExtension *Ext,
+                                        const llvm::BitstreamCursor &InStream)
+  : ModuleFileExtensionReader(Ext), Stream(InStream)
+{
+  // Read the extension block.
+  SmallVector<uint64_t, 4> Record;
+  while (true) {
+    llvm::BitstreamEntry Entry = Stream.advanceSkippingSubblocks();
+    switch (Entry.Kind) {
+    case llvm::BitstreamEntry::SubBlock:
+    case llvm::BitstreamEntry::EndBlock:
+    case llvm::BitstreamEntry::Error:
+      return;
+
+    case llvm::BitstreamEntry::Record:
+      break;
+    }
+
+    Record.clear();
+    StringRef Blob;
+    unsigned RecCode = Stream.readRecord(Entry.ID, Record, &Blob);
+    switch (RecCode) {
+    case FIRST_EXTENSION_RECORD_ID: {
+      StringRef Message = Blob.substr(0, Record[0]);
+      fprintf(stderr, "Read extension block message: %s\n",
+              Message.str().c_str());
+      break;
+    }
+    }
+  }
+}
+
+TestModuleFileExtension::Reader::~Reader() { }
+
+TestModuleFileExtension::~TestModuleFileExtension() { }
+
+ModuleFileExtensionMetadata
+TestModuleFileExtension::getExtensionMetadata() const {
+  return { BlockName, MajorVersion, MinorVersion, UserInfo };
+}
+
+llvm::hash_code TestModuleFileExtension::hashExtension(
+                  llvm::hash_code Code) const {
+  if (Hashed) {
+    Code = llvm::hash_combine(Code, BlockName);
+    Code = llvm::hash_combine(Code, MajorVersion);
+    Code = llvm::hash_combine(Code, MinorVersion);
+    Code = llvm::hash_combine(Code, UserInfo);
+  }
+  
+  return Code;
+}
+
+std::unique_ptr<ModuleFileExtensionWriter>
+TestModuleFileExtension::createExtensionWriter(ASTWriter &) {
+  return std::unique_ptr<ModuleFileExtensionWriter>(new Writer(this));
+}
+
+std::unique_ptr<ModuleFileExtensionReader>
+TestModuleFileExtension::createExtensionReader(
+  const ModuleFileExtensionMetadata &Metadata,
+  ASTReader &Reader, serialization::ModuleFile &Mod,
+  const llvm::BitstreamCursor &Stream)
+{
+  assert(Metadata.BlockName == BlockName && "Wrong block name");
+  if (std::make_pair(Metadata.MajorVersion, Metadata.MinorVersion) !=
+        std::make_pair(MajorVersion, MinorVersion)) {
+    Reader.getDiags().Report(Mod.ImportLoc,
+                             diag::err_test_module_file_extension_version)
+      << BlockName << Metadata.MajorVersion << Metadata.MinorVersion
+      << MajorVersion << MinorVersion;
+    return nullptr;
+  }
+
+  return std::unique_ptr<ModuleFileExtensionReader>(
+                                                    new TestModuleFileExtension::Reader(this, Stream));
+}
diff --git a/clang/lib/Frontend/TestModuleFileExtension.h b/clang/lib/Frontend/TestModuleFileExtension.h
new file mode 100644
index 00000000000..2a7245463cf
--- /dev/null
+++ b/clang/lib/Frontend/TestModuleFileExtension.h
@@ -0,0 +1,71 @@
+//===-- TestModuleFileExtension.h - Module Extension Tester -----*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+#ifndef LLVM_CLANG_FRONTEND_TESTMODULEFILEEXTENSION_H
+#define LLVM_CLANG_FRONTEND_TESTMODULEFILEEXTENSION_H
+
+#include "clang/Serialization/ModuleFileExtension.h"
+#include "clang/Basic/LLVM.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Bitcode/BitstreamReader.h"
+#include <string>
+
+namespace clang {
+
+/// A module file extension used for testing purposes.
+class TestModuleFileExtension : public ModuleFileExtension {
+  std::string BlockName;
+  unsigned MajorVersion;
+  unsigned MinorVersion;
+  bool Hashed;
+  std::string UserInfo;
+
+  class Writer : public ModuleFileExtensionWriter {
+  public:
+    Writer(ModuleFileExtension *Ext) : ModuleFileExtensionWriter(Ext) { }
+    ~Writer() override;
+
+    void writeExtensionContents(llvm::BitstreamWriter &Stream) override;
+  };
+
+  class Reader : public ModuleFileExtensionReader {
+    llvm::BitstreamCursor Stream;
+
+  public:
+    ~Reader() override;
+
+    Reader(ModuleFileExtension *Ext, const llvm::BitstreamCursor &InStream);
+  };
+
+public:
+  TestModuleFileExtension(StringRef BlockName,
+                          unsigned MajorVersion,
+                          unsigned MinorVersion,
+                          bool Hashed,
+                          StringRef UserInfo)
+    : BlockName(BlockName),
+      MajorVersion(MajorVersion), MinorVersion(MinorVersion),
+      Hashed(Hashed), UserInfo(UserInfo) { }
+  ~TestModuleFileExtension() override;
+
+  ModuleFileExtensionMetadata getExtensionMetadata() const override;
+
+  llvm::hash_code hashExtension(llvm::hash_code Code) const override;
+
+  std::unique_ptr<ModuleFileExtensionWriter>
+  createExtensionWriter(ASTWriter &Writer) override;
+
+  std::unique_ptr<ModuleFileExtensionReader>
+  createExtensionReader(const ModuleFileExtensionMetadata &Metadata,
+                        ASTReader &Reader, serialization::ModuleFile &Mod,
+                        const llvm::BitstreamCursor &Stream) override;
+};
+
+} // end namespace clang
+
+#endif // LLVM_CLANG_FRONTEND_TESTMODULEFILEEXTENSION_H