[llvm-branch-commits] [clang-tools-extra] [include-cleaner][NFC] share main-file location classification (PR #196762)

[Daniil Dudkin via llvm-branch-commits]

https://github.com/unterumarmung created https://github.com/llvm/llvm-project/pull/196762

None

>From 3987d7ec5bfbfecf8e1e1673ca7b57425002e3d6 Mon Sep 17 00:00:00 2001
From: Daniil Dudkin <unterumarmung at yandex.ru>
Date: Sat, 9 May 2026 16:53:29 +0300
Subject: [PATCH] [include-cleaner][NFC] share main-file location
 classification

---
 .../include-cleaner/lib/Record.cpp            |  9 ++---
 .../include-cleaner/lib/Types.cpp             | 15 +++++++++
 .../include-cleaner/lib/TypesInternal.h       | 10 ++++++
 .../include-cleaner/unittests/TypesTest.cpp   | 33 +++++++++++++++++++
 4 files changed, 60 insertions(+), 7 deletions(-)

diff --git a/clang-tools-extra/include-cleaner/lib/Record.cpp b/clang-tools-extra/include-cleaner/lib/Record.cpp
index a19fcf525919d..2632f58ba8c9d 100644
--- a/clang-tools-extra/include-cleaner/lib/Record.cpp
+++ b/clang-tools-extra/include-cleaner/lib/Record.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang-include-cleaner/Record.h"
+#include "TypesInternal.h"
 #include "clang-include-cleaner/Types.h"
 #include "clang/AST/ASTConsumer.h"
 #include "clang/AST/ASTContext.h"
@@ -159,13 +160,7 @@ class PPRecorder : public PPCallbacks {
 
 private:
   bool shouldRecordMacroRef(SourceLocation Loc) const {
-    const SourceLocation ExpandedLoc = SM.getExpansionLoc(Loc);
-    const FileID FID = SM.getFileID(ExpandedLoc);
-    if (FID == SM.getMainFileID())
-      return true;
-    const SourceLocation IncludeLoc = SM.getIncludeLoc(FID);
-    return IncludeLoc.isValid() &&
-           SM.getFileID(IncludeLoc) == SM.getMainFileID();
+    return locateInMainFile(Loc, SM) != MainFileLocation::Other;
   }
 
   void recordMacroRef(const Token &Tok, const MacroInfo &MI,
diff --git a/clang-tools-extra/include-cleaner/lib/Types.cpp b/clang-tools-extra/include-cleaner/lib/Types.cpp
index d7635c258e522..975bd2ceb164e 100644
--- a/clang-tools-extra/include-cleaner/lib/Types.cpp
+++ b/clang-tools-extra/include-cleaner/lib/Types.cpp
@@ -10,6 +10,7 @@
 #include "TypesInternal.h"
 #include "clang/AST/Decl.h"
 #include "clang/Basic/FileEntry.h"
+#include "clang/Basic/SourceManager.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
@@ -109,6 +110,20 @@ llvm::SmallString<128> normalizePath(llvm::StringRef Path) {
   return P;
 }
 
+MainFileLocation locateInMainFile(SourceLocation Loc, const SourceManager &SM) {
+  const SourceLocation ExpandedLoc = SM.getExpansionLoc(Loc);
+  const FileID FID = SM.getFileID(ExpandedLoc);
+  if (FID == SM.getMainFileID())
+    return MainFileLocation::MainFile;
+  const SourceLocation IncludeLoc = SM.getIncludeLoc(FID);
+  if (!IncludeLoc.isValid())
+    return MainFileLocation::Other;
+  const FileID IncludeFID = SM.getFileID(SM.getExpansionLoc(IncludeLoc));
+  if (IncludeFID == SM.getMainFileID())
+    return MainFileLocation::DirectInclude;
+  return MainFileLocation::Other;
+}
+
 void Includes::addSearchDirectory(llvm::StringRef Path) {
   SearchPath.try_emplace(normalizePath(Path));
 }
diff --git a/clang-tools-extra/include-cleaner/lib/TypesInternal.h b/clang-tools-extra/include-cleaner/lib/TypesInternal.h
index 60c28956c13c9..217c03c5bd7b5 100644
--- a/clang-tools-extra/include-cleaner/lib/TypesInternal.h
+++ b/clang-tools-extra/include-cleaner/lib/TypesInternal.h
@@ -19,6 +19,9 @@
 namespace llvm {
 class raw_ostream;
 }
+namespace clang {
+class SourceManager;
+} // namespace clang
 namespace clang::include_cleaner {
 /// A place where a symbol can be provided.
 /// It is either a physical file of the TU (SourceLocation) or a logical
@@ -98,6 +101,13 @@ template <typename T> struct Hinted : public T {
 
 llvm::SmallString<128> normalizePath(llvm::StringRef Path);
 
+enum class MainFileLocation {
+  MainFile,
+  DirectInclude,
+  Other,
+};
+MainFileLocation locateInMainFile(SourceLocation Loc, const SourceManager &SM);
+
 } // namespace clang::include_cleaner
 
 #endif
diff --git a/clang-tools-extra/include-cleaner/unittests/TypesTest.cpp b/clang-tools-extra/include-cleaner/unittests/TypesTest.cpp
index 6f7491928fb05..9ce7e9b972dd1 100644
--- a/clang-tools-extra/include-cleaner/unittests/TypesTest.cpp
+++ b/clang-tools-extra/include-cleaner/unittests/TypesTest.cpp
@@ -10,6 +10,7 @@
 #include "TypesInternal.h"
 #include "clang/Basic/FileManager.h"
 #include "clang/Basic/FileSystemOptions.h"
+#include "clang/Testing/TestAST.h"
 #include "clang/Tooling/Inclusions/StandardLibrary.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/Support/VirtualFileSystem.h"
@@ -114,5 +115,37 @@ TEST(NormalizePathTest, PreservesRootPath) {
   EXPECT_EQ(normalizePath("/").str(), "/");
 }
 
+TEST(MainFileLocationTest, ClassifiesSourceLocations) {
+  TestInputs Inputs;
+  Inputs.Code = R"cpp(
+    #include "direct.h"
+    int main_file;
+  )cpp";
+  Inputs.ExtraFiles["direct.h"] = R"cpp(
+    int direct_include;
+    #include "transitive.h"
+  )cpp";
+  Inputs.ExtraFiles["transitive.h"] = "int transitive_include;\n";
+  TestAST AST(Inputs);
+  const SourceManager &SM = AST.sourceManager();
+
+  EXPECT_EQ(locateInMainFile(SM.translateLineCol(SM.getMainFileID(), 3, 9), SM),
+            MainFileLocation::MainFile);
+  EXPECT_EQ(
+      locateInMainFile(
+          SM.translateFileLineCol(
+              &AST.fileManager().getOptionalFileRef("direct.h")->getFileEntry(),
+              2, 9),
+          SM),
+      MainFileLocation::DirectInclude);
+  EXPECT_EQ(locateInMainFile(
+                SM.translateFileLineCol(&AST.fileManager()
+                                             .getOptionalFileRef("transitive.h")
+                                             ->getFileEntry(),
+                                        1, 9),
+                SM),
+            MainFileLocation::Other);
+}
+
 } // namespace
 } // namespace clang::include_cleaner