Allow IWYU pragma: keep on forward declarations as an escape hatch for cases where IWYU incorrectly concludes that a forward declaration is not necessary and recommends its removal.

6 files changed, 37 insertions, 22 deletions

diff --git a/docs/IWYUPragmas.md b/docs/IWYUPragmas.md
index b1568b4..fba0ff6 100644
--- a/docs/IWYUPragmas.md
+++ b/docs/IWYUPragmas.md
@@ -7,12 +7,14 @@ All pragmas start with `// IWYU pragma: ` or `/* IWYU pragma: `. They are case-s
 
 ### IWYU pragma: keep ###
 
-This pragma applies to a single `#include` directive. It forces IWYU to keep an inclusion even if it is deemed unnecessary.
+This pragma applies to a single `#include` directive or forward declaration. It forces IWYU to keep an inclusion even if it is deemed unnecessary.
 
     main.cc:
       #include <vector> // IWYU pragma: keep
+     
+      class ForwardDeclaration; // IWYU pragma: keep
 
-In this case, `std::vector` isn't used, so `<vector>` would normally be discarded, but the pragma instructs IWYU to leave it.
+In this case, `std::vector` isn't used, so `<vector>` would normally be discarded, but the pragma instructs IWYU to leave it. Similarly the class `ForwardDeclaration` isn't used but is kept because of the pragma on it.
 
 
 ### IWYU pragma: export ###
diff --git a/iwyu.cc b/iwyu.cc
index 8650546..f9b1982 100644
--- a/iwyu.cc
+++ b/iwyu.cc
@@ -105,6 +105,7 @@
 #include "iwyu_ast_util.h"
 #include "iwyu_cache.h"
 #include "iwyu_globals.h"
+#include "iwyu_lexer_utils.h"
 #include "iwyu_location_util.h"
 #include "iwyu_output.h"
 #include "iwyu_path_util.h"
@@ -3661,6 +3662,12 @@ class IwyuAstConsumer
               definitely_keep_fwd_decl = true;
           }
         }
+      } else {
+        SourceLocation decl_end_location = decl->getSourceRange().getEnd();
+        if (LineHasText(decl_end_location, "// IWYU pragma: keep") ||
+            LineHasText(decl_end_location, "/* IWYU pragma: keep")) {
+          definitely_keep_fwd_decl = true;
+        }
       }
 
       preprocessor_info().FileInfoFor(CurrentFileEntry())->AddForwardDeclare(
diff --git a/iwyu_lexer_utils.cc b/iwyu_lexer_utils.cc
index c0bf67c..9a3545f 100644
--- a/iwyu_lexer_utils.cc
+++ b/iwyu_lexer_utils.cc
@@ -8,6 +8,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "iwyu_lexer_utils.h"
+#include "iwyu_globals.h"
 
 #include <cstring>
 #include <string>
@@ -28,11 +29,18 @@ using clang::SourceLocation;
 using clang::SourceManager;
 using clang::SourceRange;
 using clang::Token;
+using llvm::StringRef;
 using std::string;
 using std::vector;
 
 namespace include_what_you_use {
 
+bool LineHasText(SourceLocation source_location, StringRef text) {
+  const StringRef data =
+      GetSourceTextUntilEndOfLine(source_location, DefaultDataGetter());
+  return data.find(text) != StringRef::npos;
+}
+
 // SourceManagerCharacterDataGetter method implementations.
 SourceManagerCharacterDataGetter::SourceManagerCharacterDataGetter(
     const SourceManager& source_manager)
@@ -48,14 +56,13 @@ const char* SourceManagerCharacterDataGetter::GetCharacterData(
   return data;
 }
 
-string GetSourceTextUntilEndOfLine(
-    SourceLocation start_loc,
-    const CharacterDataGetterInterface& data_getter) {
+StringRef GetSourceTextUntilEndOfLine(
+    SourceLocation start_loc, const CharacterDataGetterInterface& data_getter) {
   const char* data = data_getter.GetCharacterData(start_loc);
   const char* line_end = strchr(data, '\n');
   if (!line_end)
     return data;
-  return string(data, line_end - data);
+  return StringRef(data, line_end - data);
 }
 
 SourceLocation GetLocationAfter(
diff --git a/iwyu_lexer_utils.h b/iwyu_lexer_utils.h
index b0719d2..872b5e5 100644
--- a/iwyu_lexer_utils.h
+++ b/iwyu_lexer_utils.h
@@ -25,6 +25,11 @@ namespace include_what_you_use {
 using std::string;
 using std::vector;
 
+// For a particular source line that source_location points to,
+// returns true if the given text occurs on the line.
+// (Case sensitive.)
+bool LineHasText(clang::SourceLocation source_location, llvm::StringRef text);
+
 // Interface to get character data from a SourceLocation. This allows
 // tests to avoid constructing a SourceManager yet still allow iwyu to
 // get the character data from SourceLocations.
@@ -46,7 +51,7 @@ class SourceManagerCharacterDataGetter : public CharacterDataGetterInterface {
 };
 
 // Returns the source-code line from the current location until \n.
-string GetSourceTextUntilEndOfLine(
+llvm::StringRef GetSourceTextUntilEndOfLine(
     clang::SourceLocation start_loc,
     const CharacterDataGetterInterface& data_getter);
 
diff --git a/iwyu_preprocessor.cc b/iwyu_preprocessor.cc
index 7882aa3..00bcb32 100644
--- a/iwyu_preprocessor.cc
+++ b/iwyu_preprocessor.cc
@@ -93,16 +93,6 @@ static string GetIncludeNameAsWritten(SourceLocation include_loc) {
   return GetIncludeNameAsWritten(include_loc, DefaultDataGetter());
 }
 
-// For a particular #include line that include_loc points to,
-// returns true if the given text occurs on the line.
-// (Case sensitive.)
-static bool IncludeLineHasText(SourceLocation include_loc,
-                               const string& text) {
-  const string data = GetSourceTextUntilEndOfLine(include_loc,
-                                                  DefaultDataGetter());
-  return data.find(text) != string::npos;
-}
-
 //------------------------------------------------------------
 // Utilities on macros.
 
@@ -405,13 +395,13 @@ void IwyuPreprocessorInfo::MaybeProtectInclude(
   // TODO(dsturtevant): As written "// // IWYU pragma: keep" is incorrectly
   // interpreted as a pragma. Maybe do "keep" and "export" pragma handling
   // in HandleComment?
-  if (IncludeLineHasText(includer_loc, "// IWYU pragma: keep") ||
-      IncludeLineHasText(includer_loc, "/* IWYU pragma: keep")) {
+  if (LineHasText(includer_loc, "// IWYU pragma: keep") ||
+      LineHasText(includer_loc, "/* IWYU pragma: keep")) {
     protect_reason = "pragma_keep";
     FileInfoFor(includer)->ReportKnownDesiredFile(includee);
 
-  } else if (IncludeLineHasText(includer_loc, "// IWYU pragma: export") ||
-             IncludeLineHasText(includer_loc, "/* IWYU pragma: export") ||
+  } else if (LineHasText(includer_loc, "// IWYU pragma: export") ||
+             LineHasText(includer_loc, "/* IWYU pragma: export") ||
              HasOpenBeginExports(includer)) {
     protect_reason = "pragma_export";
     const string quoted_includer =
diff --git a/tests/cxx/comment_pragmas.cc b/tests/cxx/comment_pragmas.cc
index 9e9fbce..98c8bfa 100644
--- a/tests/cxx/comment_pragmas.cc
+++ b/tests/cxx/comment_pragmas.cc
@@ -106,10 +106,12 @@
 #include "tests/cxx/comment_pragmas-d22.h" // IWYU pragma: keep
 #include "tests/cxx/comment_pragmas-d22.h"
 
-
 class CommentPragmasD19;  // Needed, but removed due to no_forward_declare.
 class CommentPragmasTest21a;  // Needed but removed due to no_forward_declare.
 
+class ForwardDeclaredUnnecessary1;  // IWYU pragma: keep
+class ForwardDeclaredUnnecessary2;  /* IWYU pragma: keep */
+
 // The following classes are all defined in public files exported by i2.h.
 // IWYU: CommentPragmasI2 is...*comment_pragmas-i1.h
 CommentPragmasI2 cpi2;
@@ -243,5 +245,7 @@ The full include-list for tests/cxx/comment_pragmas.cc:
 #include "tests/cxx/comment_pragmas-i8.h"  // for CommentPragmasI8
 #include "tests/cxx/indirect.h"  // for IndirectClass
 #include "tests/cxx/no_such_file.h"  // for CommentPragmasD2
+class ForwardDeclaredUnnecessary1;  // lines XX-XX
+class ForwardDeclaredUnnecessary2;  // lines XX-XX
 
 ***** IWYU_SUMMARY */