2 files changed, 64 insertions, 3 deletions

diff --git a/clang/lib/Frontend/PrintPreprocessedOutput.cpp b/clang/lib/Frontend/PrintPreprocessedOutput.cpp
index ffedf3cac84..832eaf2926f 100644
--- a/clang/lib/Frontend/PrintPreprocessedOutput.cpp
+++ b/clang/lib/Frontend/PrintPreprocessedOutput.cpp
@@ -174,6 +174,9 @@ public:
   void MacroUndefined(const Token &MacroNameTok,
                       const MacroDefinition &MD,
                       const MacroDirective *Undef) override;
+
+  void BeginModule(const Module *M);
+  void EndModule(const Module *M);
 };
 }  // end anonymous namespace
 
@@ -372,6 +375,20 @@ void PrintPPOutputPPCallbacks::InclusionDirective(SourceLocation HashLoc,
   }
 }
 
+/// Handle entering the scope of a module during a module compilation.
+void PrintPPOutputPPCallbacks::BeginModule(const Module *M) {
+  startNewLineIfNeeded();
+  OS << "#pragma clang module begin " << M->getFullModuleName();
+  setEmittedDirectiveOnThisLine();
+}
+
+/// Handle leaving the scope of a module during a module compilation.
+void PrintPPOutputPPCallbacks::EndModule(const Module *M) {
+  startNewLineIfNeeded();
+  OS << "#pragma clang module end /*" << M->getFullModuleName() << "*/";
+  setEmittedDirectiveOnThisLine();
+}
+
 /// Ident - Handle #ident directives when read by the preprocessor.
 ///
 void PrintPPOutputPPCallbacks::Ident(SourceLocation Loc, StringRef S) {
@@ -685,13 +702,27 @@ static void PrintPreprocessedTokens(Preprocessor &PP, Token &Tok,
       // -traditional-cpp the lexer keeps /all/ whitespace, including comments.
       SourceLocation StartLoc = Tok.getLocation();
       Callbacks->MoveToLine(StartLoc.getLocWithOffset(Tok.getLength()));
-    } else if (Tok.is(tok::annot_module_include) ||
-               Tok.is(tok::annot_module_begin) ||
-               Tok.is(tok::annot_module_end)) {
+    } else if (Tok.is(tok::annot_module_include)) {
       // PrintPPOutputPPCallbacks::InclusionDirective handles producing
       // appropriate output here. Ignore this token entirely.
       PP.Lex(Tok);
       continue;
+    } else if (Tok.is(tok::annot_module_begin)) {
+      // FIXME: We retrieve this token after the FileChanged callback, and
+      // retrieve the module_end token before the FileChanged callback, so
+      // we render this within the file and render the module end outside the
+      // file, but this is backwards from the token locations: the module_begin
+      // token is at the include location (outside the file) and the module_end
+      // token is at the EOF location (within the file).
+      Callbacks->BeginModule(
+          reinterpret_cast<Module *>(Tok.getAnnotationValue()));
+      PP.Lex(Tok);
+      continue;
+    } else if (Tok.is(tok::annot_module_end)) {
+      Callbacks->EndModule(
+          reinterpret_cast<Module *>(Tok.getAnnotationValue()));
+      PP.Lex(Tok);
+      continue;
     } else if (IdentifierInfo *II = Tok.getIdentifierInfo()) {
       OS << II->getName();
     } else if (Tok.isLiteral() && !Tok.needsCleaning() &&
diff --git a/clang/lib/Frontend/Rewrite/InclusionRewriter.cpp b/clang/lib/Frontend/Rewrite/InclusionRewriter.cpp
index ee61f76d029..d45cbc01df8 100644
--- a/clang/lib/Frontend/Rewrite/InclusionRewriter.cpp
+++ b/clang/lib/Frontend/Rewrite/InclusionRewriter.cpp
@@ -46,6 +46,8 @@ class InclusionRewriter : public PPCallbacks {
   std::map<unsigned, IncludedFile> FileIncludes;
   /// Tracks where inclusions that import modules are found.
   std::map<unsigned, const Module *> ModuleIncludes;
+  /// Tracks where inclusions that enter modules (in a module build) are found.
+  std::map<unsigned, const Module *> ModuleEntryIncludes;
   /// Used transitively for building up the FileIncludes mapping over the
   /// various \c PPCallbacks callbacks.
   SourceLocation LastInclusionLocation;
@@ -57,6 +59,11 @@ public:
     PredefinesBuffer = Buf;
   }
   void detectMainFileEOL();
+  void handleModuleBegin(Token &Tok) {
+    assert(Tok.getKind() == tok::annot_module_begin);
+    ModuleEntryIncludes.insert({Tok.getLocation().getRawEncoding(),
+                                (Module *)Tok.getAnnotationValue()});
+  }
 private:
   void FileChanged(SourceLocation Loc, FileChangeReason Reason,
                    SrcMgr::CharacteristicKind FileType,
@@ -84,6 +91,7 @@ private:
                         bool &FileExists);
   const IncludedFile *FindIncludeAtLocation(SourceLocation Loc) const;
   const Module *FindModuleAtLocation(SourceLocation Loc) const;
+  const Module *FindEnteredModule(SourceLocation Loc) const;
   StringRef NextIdentifierName(Lexer &RawLex, Token &RawToken);
 };
 
@@ -211,6 +219,16 @@ InclusionRewriter::FindModuleAtLocation(SourceLocation Loc) const {
   return nullptr;
 }
 
+/// Simple lookup for a SourceLocation (specifically one denoting the hash in
+/// an inclusion directive) in the map of module entry information.
+const Module *
+InclusionRewriter::FindEnteredModule(SourceLocation Loc) const {
+  const auto I = ModuleEntryIncludes.find(Loc.getRawEncoding());
+  if (I != ModuleEntryIncludes.end())
+    return I->second;
+  return nullptr;
+}
+
 /// Detect the likely line ending style of \p FromFile by examining the first
 /// newline found within it.
 static StringRef DetectEOL(const MemoryBuffer &FromFile) {
@@ -452,8 +470,18 @@ void InclusionRewriter::Process(FileID FileId,
             if (const Module *Mod = FindModuleAtLocation(Loc))
               WriteImplicitModuleImport(Mod);
             else if (const IncludedFile *Inc = FindIncludeAtLocation(Loc)) {
+              const Module *Mod = FindEnteredModule(Loc);
+              if (Mod)
+                OS << "#pragma clang module begin " << Mod->getFullModuleName()
+                   << "\n";
+
               // Include and recursively process the file.
               Process(Inc->Id, Inc->FileType);
+
+              if (Mod)
+                OS << "#pragma clang module end /*" << Mod->getFullModuleName()
+                   << "*/\n";
+
               // Add line marker to indicate we're returning from an included
               // file.
               LineInfoExtra = " 2";
@@ -590,6 +618,8 @@ void clang::RewriteIncludesInInput(Preprocessor &PP, raw_ostream *OS,
   PP.SetMacroExpansionOnlyInDirectives();
   do {
     PP.Lex(Tok);
+    if (Tok.is(tok::annot_module_begin))
+      Rewrite->handleModuleBegin(Tok);
   } while (Tok.isNot(tok::eof));
   Rewrite->setPredefinesBuffer(SM.getBuffer(PP.getPredefinesFileID()));
   Rewrite->Process(PP.getPredefinesFileID(), SrcMgr::C_User);