Add #pragma clang module begin/end pragmas and generate them when preprocessing a module.

These pragmas are intended to simulate the effect of entering or leaving a file
with an associated module. This is not completely implemented yet: declarations
between the pragmas will not be attributed to the correct module, but macro
visibility is already functional.

Modules named by #pragma clang module begin must already be known to clang (in
some module map that's either loaded or on the search path).


git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@302098 91177308-0d34-0410-b5e6-96231b3b80d8

5 files changed, 176 insertions, 66 deletions

diff --git a/lib/Lex/PPDirectives.cpp b/lib/Lex/PPDirectives.cpp
index 4826e399af..06fee8e5b0 100644
--- a/lib/Lex/PPDirectives.cpp
+++ b/lib/Lex/PPDirectives.cpp
@@ -2049,12 +2049,12 @@ void Preprocessor::HandleIncludeDirective(SourceLocation HashLoc,
         M->getTopLevelModuleName() == getLangOpts().CurrentModule)
       return;
 
-    assert(!CurSubmodule && "should not have marked this as a module yet");
-    CurSubmodule = M;
+    assert(!CurLexerSubmodule && "should not have marked this as a module yet");
+    CurLexerSubmodule = M;
 
     // Let the macro handling code know that any future macros are within
     // the new submodule.
-    EnterSubmodule(M, HashLoc);
+    EnterSubmodule(M, HashLoc, /*ForPragma*/false);
 
     // Let the parser know that any future declarations are within the new
     // submodule.
@@ -2082,7 +2082,7 @@ void Preprocessor::HandleIncludeNextDirective(SourceLocation HashLoc,
   } else if (isInPrimaryFile()) {
     Lookup = nullptr;
     Diag(IncludeNextTok, diag::pp_include_next_in_primary);
-  } else if (CurSubmodule) {
+  } else if (CurLexerSubmodule) {
     // Start looking up in the directory *after* the one in which the current
     // file would be found, if any.
     assert(CurPPLexer && "#include_next directive in macro?");
diff --git a/lib/Lex/PPLexerChange.cpp b/lib/Lex/PPLexerChange.cpp
index fcc49b3870..1938328c90 100644
--- a/lib/Lex/PPLexerChange.cpp
+++ b/lib/Lex/PPLexerChange.cpp
@@ -117,7 +117,7 @@ void Preprocessor::EnterSourceFileWithLexer(Lexer *TheLexer,
   CurLexer.reset(TheLexer);
   CurPPLexer = TheLexer;
   CurDirLookup = CurDir;
-  CurSubmodule = nullptr;
+  CurLexerSubmodule = nullptr;
   if (CurLexerKind != CLK_LexAfterModuleImport)
     CurLexerKind = CLK_Lexer;
 
@@ -142,7 +142,7 @@ void Preprocessor::EnterSourceFileWithPTH(PTHLexer *PL,
   CurDirLookup = CurDir;
   CurPTHLexer.reset(PL);
   CurPPLexer = CurPTHLexer.get();
-  CurSubmodule = nullptr;
+  CurLexerSubmodule = nullptr;
   if (CurLexerKind != CLK_LexAfterModuleImport)
     CurLexerKind = CLK_PTHLexer;
   
@@ -337,6 +337,26 @@ bool Preprocessor::HandleEndOfFile(Token &Result, bool isEndOfMacro) {
   assert(!CurTokenLexer &&
          "Ending a file when currently in a macro!");
 
+  // If we have an unclosed module region from a pragma at the end of a
+  // module, complain and close it now.
+  // FIXME: This is not correct if we are building a module from PTH.
+  const bool LeavingSubmodule = CurLexer && CurLexerSubmodule;
+  if ((LeavingSubmodule || IncludeMacroStack.empty()) &&
+      !BuildingSubmoduleStack.empty() &&
+      BuildingSubmoduleStack.back().IsPragma) {
+    Diag(BuildingSubmoduleStack.back().ImportLoc,
+         diag::err_pp_module_begin_without_module_end);
+    Module *M = LeaveSubmodule(/*ForPragma*/true);
+
+    Result.startToken();
+    const char *EndPos = getCurLexerEndPos();
+    CurLexer->BufferPtr = EndPos;
+    CurLexer->FormTokenWithChars(Result, EndPos, tok::annot_module_end);
+    Result.setAnnotationEndLoc(Result.getLocation());
+    Result.setAnnotationValue(M);
+    return true;
+  }
+
   // See if this file had a controlling macro.
   if (CurPPLexer) {  // Not ending a macro, ignore it.
     if (const IdentifierInfo *ControllingMacro =
@@ -442,18 +462,17 @@ bool Preprocessor::HandleEndOfFile(Token &Result, bool isEndOfMacro) {
     if (Callbacks && !isEndOfMacro && CurPPLexer)
       ExitedFID = CurPPLexer->getFileID();
 
-    bool LeavingSubmodule = CurSubmodule && CurLexer;
     if (LeavingSubmodule) {
+      // We're done with this submodule.
+      Module *M = LeaveSubmodule(/*ForPragma*/false);
+
       // Notify the parser that we've left the module.
       const char *EndPos = getCurLexerEndPos();
       Result.startToken();
       CurLexer->BufferPtr = EndPos;
       CurLexer->FormTokenWithChars(Result, EndPos, tok::annot_module_end);
       Result.setAnnotationEndLoc(Result.getLocation());
-      Result.setAnnotationValue(CurSubmodule);
-
-      // We're done with this submodule.
-      LeaveSubmodule();
+      Result.setAnnotationValue(M);
     }
 
     // We're done with the #included file.
@@ -628,11 +647,13 @@ void Preprocessor::HandleMicrosoftCommentPaste(Token &Tok) {
   assert(!FoundLexer && "Lexer should return EOD before EOF in PP mode");
 }
 
-void Preprocessor::EnterSubmodule(Module *M, SourceLocation ImportLoc) {
+void Preprocessor::EnterSubmodule(Module *M, SourceLocation ImportLoc,
+                                  bool ForPragma) {
   if (!getLangOpts().ModulesLocalVisibility) {
     // Just track that we entered this submodule.
-    BuildingSubmoduleStack.push_back(BuildingSubmoduleInfo(
-        M, ImportLoc, CurSubmoduleState, PendingModuleMacroNames.size()));
+    BuildingSubmoduleStack.push_back(
+        BuildingSubmoduleInfo(M, ImportLoc, ForPragma, CurSubmoduleState,
+                              PendingModuleMacroNames.size()));
     return;
   }
 
@@ -673,8 +694,9 @@ void Preprocessor::EnterSubmodule(Module *M, SourceLocation ImportLoc) {
   }
 
   // Track that we entered this module.
-  BuildingSubmoduleStack.push_back(BuildingSubmoduleInfo(
-      M, ImportLoc, CurSubmoduleState, PendingModuleMacroNames.size()));
+  BuildingSubmoduleStack.push_back(
+      BuildingSubmoduleInfo(M, ImportLoc, ForPragma, CurSubmoduleState,
+                            PendingModuleMacroNames.size()));
 
   // Switch to this submodule as the current submodule.
   CurSubmoduleState = &State;
@@ -697,7 +719,13 @@ bool Preprocessor::needModuleMacros() const {
   return getLangOpts().isCompilingModule();
 }
 
-void Preprocessor::LeaveSubmodule() {
+Module *Preprocessor::LeaveSubmodule(bool ForPragma) {
+  if (BuildingSubmoduleStack.empty() ||
+      BuildingSubmoduleStack.back().IsPragma != ForPragma) {
+    assert(ForPragma && "non-pragma module enter/leave mismatch");
+    return nullptr;
+  }
+
   auto &Info = BuildingSubmoduleStack.back();
 
   Module *LeavingMod = Info.M;
@@ -711,7 +739,7 @@ void Preprocessor::LeaveSubmodule() {
     // of pending names for the surrounding submodule.
     BuildingSubmoduleStack.pop_back();
     makeModuleVisible(LeavingMod, ImportLoc);
-    return;
+    return LeavingMod;
   }
 
   // Create ModuleMacros for any macros defined in this submodule.
@@ -800,4 +828,5 @@ void Preprocessor::LeaveSubmodule() {
 
   // A nested #include makes the included submodule visible.
   makeModuleVisible(LeavingMod, ImportLoc);
+  return LeavingMod;
 }
diff --git a/lib/Lex/PPMacroExpansion.cpp b/lib/Lex/PPMacroExpansion.cpp
index 196223981d..6c7663994a 100644
--- a/lib/Lex/PPMacroExpansion.cpp
+++ b/lib/Lex/PPMacroExpansion.cpp
@@ -1453,7 +1453,7 @@ static bool EvaluateHasIncludeNext(Token &Tok,
   } else if (PP.isInPrimaryFile()) {
     Lookup = nullptr;
     PP.Diag(Tok, diag::pp_include_next_in_primary);
-  } else if (PP.getCurrentSubmodule()) {
+  } else if (PP.getCurrentLexerSubmodule()) {
     // Start looking up in the directory *after* the one in which the current
     // file would be found, if any.
     assert(PP.getCurrentLexer() && "#include_next directive in macro?");
diff --git a/lib/Lex/Pragma.cpp b/lib/Lex/Pragma.cpp
index 576151a98b..51da2baac9 100644
--- a/lib/Lex/Pragma.cpp
+++ b/lib/Lex/Pragma.cpp
@@ -534,47 +534,6 @@ void Preprocessor::HandlePragmaDependency(Token &DependencyTok) {
   }
 }
 
-void Preprocessor::HandlePragmaModuleImport(Token &ImportTok) {
-  SourceLocation ImportLoc = ImportTok.getLocation();
-
-  Token Tok;
-
-  llvm::SmallVector<std::pair<IdentifierInfo *, SourceLocation>, 8> ModuleName;
-  while (true) {
-    LexUnexpandedToken(Tok);
-    if (Tok.isNot(tok::identifier)) {
-      Diag(Tok.getLocation(),
-           diag::err_pragma_module_import_expected_module_name) << 0;
-      return;
-    }
-
-    ModuleName.emplace_back(Tok.getIdentifierInfo(), Tok.getLocation());
-
-    LexUnexpandedToken(Tok);
-    assert(Tok.isNot(tok::eof));
-    if (Tok.is(tok::eod))
-      break;
-    if (Tok.isNot(tok::period)) {
-      Diag(Tok.getLocation(),
-           diag::err_pragma_module_import_expected_module_name) << 1;
-      return;
-    }
-  }
-
-  // If we have a non-empty module path, load the named module.
-  Module *Imported =
-      TheModuleLoader.loadModule(ImportLoc, ModuleName, Module::Hidden,
-                                 /*IsIncludeDirective=*/false);
-  if (!Imported)
-    return;
-
-  makeModuleVisible(Imported, ImportLoc);
-  EnterAnnotationToken(SourceRange(ImportLoc, Tok.getLocation()),
-                       tok::annot_module_include, Imported);
-  if (Callbacks)
-    Callbacks->moduleImport(ImportLoc, ModuleName, Imported);
-}
-
 /// ParsePragmaPushOrPopMacro - Handle parsing of pragma push_macro/pop_macro.
 /// Return the IdentifierInfo* associated with the macro to push or pop.
 IdentifierInfo *Preprocessor::ParsePragmaPushOrPopMacro(Token &Tok) {
@@ -1342,6 +1301,26 @@ public:
   }
 };
 
+static bool LexModuleName(
+    Preprocessor &PP, Token &Tok,
+    llvm::SmallVectorImpl<std::pair<IdentifierInfo *, SourceLocation>>
+        &ModuleName) {
+  while (true) {
+    PP.LexUnexpandedToken(Tok);
+    if (Tok.isNot(tok::identifier)) {
+      PP.Diag(Tok.getLocation(), diag::err_pp_expected_module_name)
+        << ModuleName.empty();
+      return true;
+    }
+
+    ModuleName.emplace_back(Tok.getIdentifierInfo(), Tok.getLocation());
+
+    PP.LexUnexpandedToken(Tok);
+    if (Tok.isNot(tok::period))
+      return false;
+  }
+}
+
 /// Handle the clang \#pragma module import extension. The syntax is:
 /// \code
 ///   #pragma clang module import some.module.name
@@ -1350,8 +1329,108 @@ struct PragmaModuleImportHandler : public PragmaHandler {
   PragmaModuleImportHandler() : PragmaHandler("import") {}
 
   void HandlePragma(Preprocessor &PP, PragmaIntroducerKind Introducer,
-                    Token &ImportTok) override {
-    PP.HandlePragmaModuleImport(ImportTok);
+                    Token &Tok) override {
+    SourceLocation ImportLoc = Tok.getLocation();
+
+    // Read the module name.
+    llvm::SmallVector<std::pair<IdentifierInfo *, SourceLocation>, 8>
+        ModuleName;
+    if (LexModuleName(PP, Tok, ModuleName))
+      return;
+
+    if (Tok.isNot(tok::eod))
+      PP.Diag(Tok, diag::ext_pp_extra_tokens_at_eol) << "pragma";
+
+    // If we have a non-empty module path, load the named module.
+    Module *Imported =
+        PP.getModuleLoader().loadModule(ImportLoc, ModuleName, Module::Hidden,
+                                      /*IsIncludeDirective=*/false);
+    if (!Imported)
+      return;
+
+    PP.makeModuleVisible(Imported, ImportLoc);
+    PP.EnterAnnotationToken(SourceRange(ImportLoc, ModuleName.back().second),
+                            tok::annot_module_include, Imported);
+    if (auto *CB = PP.getPPCallbacks())
+      CB->moduleImport(ImportLoc, ModuleName, Imported);
+  }
+};
+
+/// Handle the clang \#pragma module begin extension. The syntax is:
+/// \code
+///   #pragma clang module begin some.module.name
+///   ...
+///   #pragma clang module end
+/// \endcode
+struct PragmaModuleBeginHandler : public PragmaHandler {
+  PragmaModuleBeginHandler() : PragmaHandler("begin") {}
+
+  void HandlePragma(Preprocessor &PP, PragmaIntroducerKind Introducer,
+                    Token &Tok) override {
+    SourceLocation BeginLoc = Tok.getLocation();
+
+    // Read the module name.
+    llvm::SmallVector<std::pair<IdentifierInfo *, SourceLocation>, 8>
+        ModuleName;
+    if (LexModuleName(PP, Tok, ModuleName))
+      return;
+
+    if (Tok.isNot(tok::eod))
+      PP.Diag(Tok, diag::ext_pp_extra_tokens_at_eol) << "pragma";
+
+    // We can only enter submodules of the current module.
+    StringRef Current = PP.getLangOpts().CurrentModule;
+    if (ModuleName.front().first->getName() != Current) {
+      PP.Diag(ModuleName.front().second, diag::err_pp_module_begin_wrong_module)
+        << ModuleName.front().first << (ModuleName.size() > 1)
+        << Current.empty() << Current;
+      return;
+    }
+
+    // Find the module we're entering. We require that a module map for it
+    // be loaded or implicitly loadable.
+    // FIXME: We could create the submodule here. We'd need to know whether
+    // it's supposed to be explicit, but not much else.
+    Module *M = PP.getHeaderSearchInfo().getModuleMap().findModule(Current);
+    if (!M) {
+      PP.Diag(ModuleName.front().second,
+              diag::err_pp_module_begin_no_module_map) << Current;
+      return;
+    }
+    for (unsigned I = 1; I != ModuleName.size(); ++I) {
+      auto *NewM = M->findSubmodule(ModuleName[I].first->getName());
+      if (!NewM) {
+        PP.Diag(ModuleName[I].second, diag::err_pp_module_begin_no_submodule)
+          << M->getFullModuleName() << ModuleName[I].first;
+        return;
+      }
+      M = NewM;
+    }
+
+    // Enter the scope of the submodule.
+    PP.EnterSubmodule(M, BeginLoc, /*ForPragma*/true);
+    PP.EnterAnnotationToken(SourceRange(BeginLoc, ModuleName.back().second),
+                            tok::annot_module_begin, M);
+  }
+};
+
+/// Handle the clang \#pragma module end extension.
+struct PragmaModuleEndHandler : public PragmaHandler {
+  PragmaModuleEndHandler() : PragmaHandler("end") {}
+
+  void HandlePragma(Preprocessor &PP, PragmaIntroducerKind Introducer,
+                    Token &Tok) override {
+    SourceLocation Loc = Tok.getLocation();
+
+    PP.LexUnexpandedToken(Tok);
+    if (Tok.isNot(tok::eod))
+      PP.Diag(Tok, diag::ext_pp_extra_tokens_at_eol) << "pragma";
+
+    Module *M = PP.LeaveSubmodule(/*ForPragma*/true);
+    if (M)
+      PP.EnterAnnotationToken(SourceRange(Loc), tok::annot_module_end, M);
+    else
+      PP.Diag(Loc, diag::err_pp_module_end_without_module_begin);
   }
 };
 
@@ -1582,6 +1661,8 @@ void Preprocessor::RegisterBuiltinPragmas() {
   auto *ModuleHandler = new PragmaNamespace("module");
   AddPragmaHandler("clang", ModuleHandler);
   ModuleHandler->AddPragma(new PragmaModuleImportHandler());
+  ModuleHandler->AddPragma(new PragmaModuleBeginHandler());
+  ModuleHandler->AddPragma(new PragmaModuleEndHandler());
 
   AddPragmaHandler("STDC", new PragmaSTDC_FENV_ACCESSHandler());
   AddPragmaHandler("STDC", new PragmaSTDC_CX_LIMITED_RANGEHandler());
diff --git a/lib/Lex/Preprocessor.cpp b/lib/Lex/Preprocessor.cpp
index babef5dcc7..e409ab0365 100644
--- a/lib/Lex/Preprocessor.cpp
+++ b/lib/Lex/Preprocessor.cpp
@@ -85,10 +85,10 @@ Preprocessor::Preprocessor(std::shared_ptr<PreprocessorOptions> PPOpts,
       LastTokenWasAt(false), ModuleImportExpectsIdentifier(false),
       CodeCompletionReached(false), CodeCompletionII(nullptr),
       MainFileDir(nullptr), SkipMainFilePreamble(0, true), CurPPLexer(nullptr),
-      CurDirLookup(nullptr), CurLexerKind(CLK_Lexer), CurSubmodule(nullptr),
-      Callbacks(nullptr), CurSubmoduleState(&NullSubmoduleState),
-      MacroArgCache(nullptr), Record(nullptr), MIChainHead(nullptr),
-      DeserialMIChainHead(nullptr) {
+      CurDirLookup(nullptr), CurLexerKind(CLK_Lexer),
+      CurLexerSubmodule(nullptr), Callbacks(nullptr),
+      CurSubmoduleState(&NullSubmoduleState), MacroArgCache(nullptr),
+      Record(nullptr), MIChainHead(nullptr), DeserialMIChainHead(nullptr) {
   OwnsHeaderSearch = OwnsHeaders;
   
   CounterValue = 0; // __COUNTER__ starts at 0.