[include-cleaner] Respect language when tracking IWYU pragmas on stdlib headers (llvm#152142)

In absence we were associating all the IWYU pragmas with C++ versions of stdlib headers, which would result in false negatives when querying them from a C source file.

Author: kadircet

clang-tools-extra/include-cleaner/lib/Record.cpp

@@ -14,6 +14,7 @@
 #include "clang/Basic/FileEntry.h"
 #include "clang/Basic/FileManager.h"
 #include "clang/Basic/LLVM.h"
+#include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/Specifiers.h"
@@ -179,8 +180,10 @@ class PragmaIncludes::RecordPragma : public PPCallbacks, public CommentHandler {
   RecordPragma(const CompilerInstance &CI, PragmaIncludes *Out)
       : RecordPragma(CI.getPreprocessor(), Out) {}
   RecordPragma(const Preprocessor &P, PragmaIncludes *Out)
-      : SM(P.getSourceManager()), HeaderInfo(P.getHeaderSearchInfo()), Out(Out),
-        Arena(std::make_shared<llvm::BumpPtrAllocator>()),
+      : SM(P.getSourceManager()), HeaderInfo(P.getHeaderSearchInfo()),
+        L(P.getLangOpts().CPlusPlus ? tooling::stdlib::Lang::CXX
+                                    : tooling::stdlib::Lang::C),
+        Out(Out), Arena(std::make_shared<llvm::BumpPtrAllocator>()),
         UniqueStrings(*Arena),
         MainFileStem(llvm::sys::path::stem(
             SM.getNonBuiltinFilenameForID(SM.getMainFileID()).value_or(""))) {}
@@ -224,7 +227,7 @@ class PragmaIncludes::RecordPragma : public PPCallbacks, public CommentHandler {
     std::optional<Header> IncludedHeader;
     if (IsAngled)
       if (auto StandardHeader =
-              tooling::stdlib::Header::named("<" + FileName.str() + ">")) {
+              tooling::stdlib::Header::named("<" + FileName.str() + ">", L)) {
         IncludedHeader = *StandardHeader;
       }
     if (!IncludedHeader && File)
@@ -363,6 +366,7 @@ class PragmaIncludes::RecordPragma : public PPCallbacks, public CommentHandler {
   bool InMainFile = false;
   const SourceManager &SM;
   const HeaderSearch &HeaderInfo;
+  const tooling::stdlib::Lang L;
   PragmaIncludes *Out;
   std::shared_ptr<llvm::BumpPtrAllocator> Arena;
   /// Intern table for strings. Contents are on the arena.

clang-tools-extra/include-cleaner/unittests/RecordTest.cpp

@@ -18,6 +18,7 @@
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Frontend/FrontendOptions.h"
 #include "clang/Serialization/PCHContainerOperations.h"
+#include "clang/Testing/CommandLineArgs.h"
 #include "clang/Testing/TestAST.h"
 #include "clang/Tooling/Inclusions/StandardLibrary.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -514,6 +515,26 @@ TEST_F(PragmaIncludeTest, IWYUExportForStandardHeaders) {
               testing::UnorderedElementsAre(FileNamed("export.h")));
 }
 
+TEST_F(PragmaIncludeTest, IWYUExportForStandardHeadersRespectsLang) {
+  Inputs.Code = R"cpp(
+    #include "export.h"
+  )cpp";
+  Inputs.Language = TestLanguage::Lang_C99;
+  Inputs.ExtraFiles["export.h"] = R"cpp(
+    #include <stdlib.h> // IWYU pragma: export
+  )cpp";
+  Inputs.ExtraFiles["stdlib.h"] = "";
+  Inputs.ExtraArgs = {"-isystem."};
+  TestAST Processed = build();
+  auto &FM = Processed.fileManager();
+  EXPECT_THAT(PI.getExporters(*tooling::stdlib::Header::named(
+                                  "<stdlib.h>", tooling::stdlib::Lang::C),
+                              FM),
+              testing::UnorderedElementsAre(FileNamed("export.h")));
+  EXPECT_THAT(PI.getExporters(llvm::cantFail(FM.getFileRef("stdlib.h")), FM),
+              testing::UnorderedElementsAre(FileNamed("export.h")));
+}
+
 TEST_F(PragmaIncludeTest, IWYUExportBlock) {
   Inputs.Code = R"cpp(// Line 1
    #include "normal.h"