Make -frewrite-includes put an endif at the end of the included text

clang/docs/ReleaseNotes.rst

@@ -150,6 +150,9 @@ Modified Compiler Flags
 
 * ``-Woverriding-t-option`` is renamed to ``-Woverriding-option``.
 * ``-Winterrupt-service-routine`` is renamed to ``-Wexcessive-regsave`` as a generalization
+* ``-frewrite-includes`` now guards the original #include directives with
+  ``__CLANG_REWRITTEN_INCLUDES``, and ``__CLANG_REWRITTEN_SYSTEM_INCLUDES`` as
+  appropriate.
 
 Removed Compiler Flags
 -------------------------

clang/lib/Frontend/Rewrite/InclusionRewriter.cpp

@@ -90,8 +90,10 @@ class InclusionRewriter : public PPCallbacks {
                          bool EnsureNewline);
   void CommentOutDirective(Lexer &DirectivesLex, const Token &StartToken,
                            const MemoryBufferRef &FromFile, StringRef EOL,
-                           unsigned &NextToWrite, int &Lines);
+                           unsigned &NextToWrite, int &Lines,
+                           const IncludedFile *Inc = nullptr);
   const IncludedFile *FindIncludeAtLocation(SourceLocation Loc) const;
+  StringRef getIncludedFileName(const IncludedFile *Inc) const;
   const Module *FindModuleAtLocation(SourceLocation Loc) const;
   const Module *FindEnteredModule(SourceLocation Loc) const;
   bool IsIfAtLocationTrue(SourceLocation Loc) const;
@@ -311,6 +313,17 @@ void InclusionRewriter::OutputContentUpTo(const MemoryBufferRef &FromFile,
   WriteFrom = WriteTo;
 }
 
+StringRef
+InclusionRewriter::getIncludedFileName(const IncludedFile *Inc) const {
+  if (Inc) {
+    auto B = SM.getBufferOrNone(Inc->Id);
+    assert(B && "Attempting to process invalid inclusion");
+    if (B)
+      return llvm::sys::path::filename(B->getBufferIdentifier());
+  }
+  return StringRef();
+}
+
 /// Print characters from \p FromFile starting at \p NextToWrite up until the
 /// inclusion directive at \p StartToken, then print out the inclusion
 /// inclusion directive disabled by a #if directive, updating \p NextToWrite
@@ -320,7 +333,8 @@ void InclusionRewriter::CommentOutDirective(Lexer &DirectiveLex,
                                             const Token &StartToken,
                                             const MemoryBufferRef &FromFile,
                                             StringRef LocalEOL,
-                                            unsigned &NextToWrite, int &Line) {
+                                            unsigned &NextToWrite, int &Line,
+                                            const IncludedFile *Inc) {
   OutputContentUpTo(FromFile, NextToWrite,
                     SM.getFileOffset(StartToken.getLocation()), LocalEOL, Line,
                     false);
@@ -332,12 +346,21 @@ void InclusionRewriter::CommentOutDirective(Lexer &DirectiveLex,
     // OutputContentUpTo() would not output anything anyway.
     return;
   }
-  OS << "#if 0 /* expanded by -frewrite-includes */" << MainEOL;
+  if (Inc) {
+    OS << "#if defined(__CLANG_REWRITTEN_INCLUDES) ";
+    if (isSystem(Inc->FileType))
+      OS << "|| defined(__CLANG_REWRITTEN_SYSTEM_INCLUDES) ";
+    OS << "/* " << getIncludedFileName(Inc);
+  } else {
+    OS << "#if 0 /*";
+  }
+  OS << " expanded by -frewrite-includes */" << MainEOL;
   OutputContentUpTo(FromFile, NextToWrite,
                     SM.getFileOffset(DirectiveToken.getLocation()) +
                         DirectiveToken.getLength(),
                     LocalEOL, Line, true);
-  OS << "#endif /* expanded by -frewrite-includes */" << MainEOL;
+  OS << (Inc ? "#else /* " : "#endif /*") << getIncludedFileName(Inc)
+     << " expanded by -frewrite-includes */" << MainEOL;
 }
 
 /// Find the next identifier in the pragma directive specified by \p RawToken.
@@ -400,15 +423,16 @@ void InclusionRewriter::Process(FileID FileId,
           case tok::pp_include:
           case tok::pp_include_next:
           case tok::pp_import: {
-            CommentOutDirective(RawLex, HashToken, FromFile, LocalEOL, NextToWrite,
-              Line);
+            SourceLocation Loc = HashToken.getLocation();
+            const IncludedFile *Inc = FindIncludeAtLocation(Loc);
+            CommentOutDirective(RawLex, HashToken, FromFile, LocalEOL,
+                                NextToWrite, Line, Inc);
             if (FileId != PP.getPredefinesFileID())
               WriteLineInfo(FileName, Line - 1, FileType, "");
             StringRef LineInfoExtra;
-            SourceLocation Loc = HashToken.getLocation();
             if (const Module *Mod = FindModuleAtLocation(Loc))
               WriteImplicitModuleImport(Mod);
-            else if (const IncludedFile *Inc = FindIncludeAtLocation(Loc)) {
+            else if (Inc) {
               const Module *Mod = FindEnteredModule(Loc);
               if (Mod)
                 OS << "#pragma clang module begin "
@@ -420,6 +444,8 @@ void InclusionRewriter::Process(FileID FileId,
               if (Mod)
                 OS << "#pragma clang module end /*"
                    << Mod->getFullModuleName(true) << "*/\n";
+              OS << "#endif /* " << getIncludedFileName(Inc)
+                 << " expanded by -frewrite-includes */" << LocalEOL;
 
               // Add line marker to indicate we're returning from an included
               // file.

clang/test/Frontend/rewrite-includes-cli-include.c

@@ -3,6 +3,7 @@ main_file_line
 // CHECK: {{^}}# 1 "<built-in>"{{$}}
 // CHECK-NEXT: {{^}}# 1 "{{.*[/\\]Inputs(/|\\\\)}}rewrite-includes2.h" 1{{$}}
 // CHECK-NEXT: {{^}}int included_line2;{{$}}
+// CHECK-NEXT: {{^}}#endif /* rewrite-includes2.h expanded by -frewrite-includes */{{$}}
 // CHECK-NEXT: {{^}}# 1 "<built-in>" 2{{$}}
 // CHECK-NEXT: {{^}}# 1 "{{.*}}rewrite-includes-cli-include.c"{{$}}
 // CHECK-NEXT: FileCheck