[include-cleaner] Record main-file macro occurences and includes

[sammccall]

Summary

The occurrences are roots for finding used headers, like walkAST.
Includes are the targets we're matching used headers against.

[sammccall]

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Record.h

123

in the prototype I reimplemented this function in clangd, but I expect we can just reuse the RecordedIncludes class instead, copying from clangd's includes list is cheap.

(That's one argument for putting this in a different header, which I can do already if you prefer)

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Types.h

22–23

(this is already fixed at head)

[hokein]

Left a few initial comments, it looks roughly good to me (for the macro-usage case, I might miss some historical context there, I think we come to an agreement that what this patch proposes is the designed behavior).

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Record.h

123

that's an interesting idea, but we need to do it carefully because of FileEntry*, the Include structure has a filed of FileEntry*, it is not feasible for clangd to propagate it (clangd's Inclusion doesn't have it, instead it uses a HeaderID which is based on the fs::UniqueID to identify a physical file).

But I think this is not something we should worry about at the moment, and the current interfaces look quite good.

127

nit: worth a comment mentioning the unsigned is the index of All.

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Types.h

122

nit: HashLoc seems clearer than Location.

And it looks like Location and Line feels somewhat redundant -- we can get the Line loc from the Location with SourceManager. But I think it is fair to hold both, Line is an important property of the include, which will probably widely used.

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

300

why using a lower case x here?

308

this is a redef error with the one defined in the header, I think it is not intended, rename it?

321

IIUC, ErrorOK is irrelevant, and should be removed.

[sammccall]

address comments

[hokein]

looks good from my side.

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Types.h

52

probably worth adding a comment: the order of the enumerators must match the the order of template arguments in Storage.

82

thinking more about it (no action required) -- this structure seems a natural good fit to be used in UsedSymbolCB in Analysis.h rather than using SourceLocation RefLoc, Symbol Target separately.

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

101

this should be hidden in anonymous namespace, right?

106

nit: virtual can be removed.

[kadircet]

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Record.h

102

s/this/the main file/

123

right, let's keep it to the minimum now. we can extend/move-around later on

129

it feels like having a set of includes corresponding to the same fileentry, rather than a single one, feels subtle enough to deserve a comment.

the only case i can think of is, two logically different files resolving to same physical file (i.e. symlinks). these will actually have completely different spellings, and probably it'll make things hard when deciding which include to "keep" or "insert". are there other cases where we can have multiple includes corresponding to the same fileentry?

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Types.h

41

what about just a stringref to name?

42

triple slashes

46

OOC: can we really have multiple macro identifiers defined at the same source location? I guess it's cheap to compare another pointer, but might be easier if we actually didn't make the identifierinfo part of the identity.

82

agreed. i'll adjust that with the refkind patch

121

i think it's worth leaving a comment here mentioning that this can be null (couldn't resolve, or maybe we shouldn't be caring, e.g. logical stl include).

even better, maybe we should actually have a Header here? that way this would be conceptual equivalent of SymbolReference, but for includes?

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

109

we can keep a include depth based on the Reason and bail out on other operations when depth is not 1 (getFileID on Loc here probably hits the SourceManager cache anyways, so might not be important, but i feel like it's better to not rely on that).

242

we might have different handles (fileentry*) to the same file (if we end up having multiple filemanagers in action, which is unlikely as things stand today, but still). so what about using uniqueids for comparison here?

AFAICT, filemanager won't have multiple fileenrtys with the same uniqueid anyway. so we wouldn't really lose any generality.

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

263

i know these look more like functions, but convention is to use PascalCase for matchers, so Spelled instead?

327

nit: ASSERT_FALSE(Recorded.MacroReferences.empty())

354

same as above, Line?

364

can you also add an include with duplicated spelling to test spelling to set of includes mapping?

[sammccall]

address comments

[sammccall]

const

[maryammo]

This commit causes build failure on https://lab.llvm.org/buildbot/#/builders/121/builds/24947 :

[43/635] Building CXX object tools/clang/tools/extra/include-cleaner/lib/CMakeFiles/obj.clangIncludeCleaner.dir/Types.cpp.o
FAILED: tools/clang/tools/extra/include-cleaner/lib/CMakeFiles/obj.clangIncludeCleaner.dir/Types.cpp.o 
CCACHE_CPP2=yes CCACHE_HASHDIR=yes /usr/bin/ccache /usr/lib64/ccache/c++  -DGTEST_HAS_RTTI=0 -D_DEBUG -D_GNU_SOURCE -D__STDC_CONSTANT_MACROS -D__STDC_FORMAT_MACROS -D__STDC_LIMIT_MACROS -Itools/clang/tools/extra/include-cleaner/lib -I/home/buildbots/ppc64le-clang-multistage-test/clang-ppc64le-multistage/llvm/clang-tools-extra/include-cleaner/lib -I/home/buildbots/ppc64le-clang-multistage-test/clang-ppc64le-multistage/llvm/clang/include -Itools/clang/include -Iinclude -I/home/buildbots/ppc64le-clang-multistage-test/clang-ppc64le-multistage/llvm/llvm/include -I/home/buildbots/ppc64le-clang-multistage-test/clang-ppc64le-multistage/llvm/clang-tools-extra/include-cleaner/include -fPIC -fno-semantic-interposition -fvisibility-inlines-hidden -Werror=date-time -Wall -Wextra -Wno-unused-parameter -Wwrite-strings -Wcast-qual -Wno-missing-field-initializers -pedantic -Wno-long-long -Wimplicit-fallthrough -Wno-maybe-uninitialized -Wno-class-memaccess -Wno-noexcept-type -Wdelete-non-virtual-dtor -Wno-comment -Wno-misleading-indentation -fdiagnostics-color -ffunction-sections -fdata-sections -fno-common -Woverloaded-virtual -fno-strict-aliasing -O3 -DNDEBUG    -fno-exceptions -fno-rtti -UNDEBUG -std=c++17 -MD -MT tools/clang/tools/extra/include-cleaner/lib/CMakeFiles/obj.clangIncludeCleaner.dir/Types.cpp.o -MF tools/clang/tools/extra/include-cleaner/lib/CMakeFiles/obj.clangIncludeCleaner.dir/Types.cpp.o.d -o tools/clang/tools/extra/include-cleaner/lib/CMakeFiles/obj.clangIncludeCleaner.dir/Types.cpp.o -c /home/buildbots/ppc64le-clang-multistage-test/clang-ppc64le-multistage/llvm/clang-tools-extra/include-cleaner/lib/Types.cpp
In file included from /home/buildbots/ppc64le-clang-multistage-test/clang-ppc64le-multistage/llvm/clang-tools-extra/include-cleaner/lib/Types.cpp:9:
/home/buildbots/ppc64le-clang-multistage-test/clang-ppc64le-multistage/llvm/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Types.h:84:10: error: declaration of ‘clang::include_cleaner::Symbol clang::include_cleaner::SymbolReference::Symbol’ [-fpermissive]
   Symbol Symbol;
          ^~~~~~
/home/buildbots/ppc64le-clang-multistage-test/clang-ppc64le-multistage/llvm/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Types.h:51:8: error: changes meaning of ‘Symbol’ from ‘struct clang::include_cleaner::Symbol’ [-fpermissive]
 struct Symbol {
        ^~~~~~
/home/buildbots/ppc64le-clang-multistage-test/clang-ppc64le-multistage/llvm/clang-tools-extra/include-cleaner/lib/Types.cpp: In function ‘llvm::raw_ostream& clang::include_cleaner::operator<<(llvm::raw_ostream&, const clang::include_cleaner::SymbolReference&)’:
/home/buildbots/ppc64le-clang-multistage-test/clang-ppc64le-multistage/llvm/clang-tools-extra/include-cleaner/lib/Types.cpp:51:51: warning: ‘*’ in boolean context, suggest ‘&&’ instead [-Wint-in-bool-context]
                                /*Width=*/CHAR_BIT *

I was able to reproduce the failure and by reverting this commit locally it passed, can you please take a look? @sammccall

[sammccall]

In D136723#3915853, @maryammo wrote:

This commit causes build failure on https://lab.llvm.org/buildbot/#/builders/121/builds/24947 :
I was able to reproduce the failure and by reverting this commit locally it passed, can you please take a look? @sammccall

It looks like this was already fixed in
https://github.com/llvm/llvm-project/commit/d19ba74dee0b9ab553bd8a6ef5b67ff349f4bf13

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Record.h

102

"this" as in the current object, as the data flow here isn't obvious. Changed to *this, clearer?

123

Good point. We could give Include a UniqueID instead I think. Let's defer this though.

123

that's an interesting idea, but we need to do it carefully because of FileEntry*, the Include structure has a filed of FileEntry*, it is not feasible for clangd to propagate it (clangd's Inclusion doesn't have it, instead it uses a HeaderID which is based on the fs::UniqueID to identify a physical file).

But I think this is not something we should worry about at the moment, and the current interfaces look quite good.

129

Sure:

#include "foo.h"
#include "bar.h"
#include "foo.h"

it isn't useful, but people write it (by mistake).
Maybe we want to clean it up, maybe we don't, but I don't see much reason to model it wrong on purpose - it doesn't really simplify anything, and may further confuse a confusing situation.

I can add a comment - what do you want it to say?
"There may be more than one #include of the same file" just seems to restate the obvious interpretation of the type here. (And this is the header, so it's not an ideal place to dump a bunch of implementation comments)

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Types.h

41

That's twice as big, and this is already the critical path for sizeof(SymbolReference) which me can accumulate a lot of.

I don't think it's critical, but I also don't see why StringRef is better. (Less indirection vs carries some semantics + lifetime information)

46

Probably not? Removed from comparison.
(I think it's *simpler* to include it in the comparison, but it doesn't matter much)

52

Done in the base patch. D136710

121

Done (comment.)

maybe we should actually have a Header here

A Header is conceptually something like a Matcher<Include>. If we put Header here then we're communicating what the reference is going to look like: #include "foo.h" is going to satisfy a symbol usage in /path/to/foo.h or one in a file with // IWYU pragma private: include "foo.h", which we don't know.
(And it could be both, which would no longer be expressible).

Technically you could store an Header that *always* has a FileEntry, but that just seems like obfuscation.

122

Agree, we could one of these but it's a bit sad: we need SourceLocation to diagnose things, but Line is much more convenient the rest of the time.

Renamed, leaving this open to see what Kadir thinks too.

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

101

Right! This used to be friended.

109

That sounds complicated, wait until it shows up in a profile?

The conceptual stack definitely isn't properly maintained in all situations. e.g. a normal parse leaks 1 entry for... builtins or something, and clangd's parse with preamble leaks 3. I'm not specifically sure there's a problem in this scenario, but it's *not* a trivial thing that obviously can't go wrong.

(& I think if we're afraid to rely on the cache for the file we're *currently lexing* then we should seriously consider fixing the cache instead)

242

This is approximately the same issue as above (> that's an interesting idea, but we need to do it carefully because of FileEntry*`).
I agree and it doesn't really add code complexity but it does add subtlety, especially if we use FileEntry in the data structure and UniqueID for matching.
So I think deferring *all* the multi-file-manager concerns makes sense (or handling them all now if you prefer)

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

263

These are functions - I don't think there's a convention here, just confusion and bugs.
OK to fix named intead?

321

Hmm, it's not irrelevant, it means I get to make lots of mistakes in my examples :-)

Removed.

327

Why? This means we get "expected empty() is false, got true" instead of "expected size() is 0, got 7", which seems strictly worse...

Added a somewhat-useful printer for SymbolReference and switched to IsEmpty() instead.

[sammccall]

(and sorry, it seems I was sitting on a pile of comments I thought I'd sent, LMK if I should follow up on them)

[kadircet]

In D136723#3915898, @sammccall wrote:

(and sorry, it seems I was sitting on a pile of comments I thought I'd sent, LMK if I should follow up on them)

no looks good. the only meaty discussion was fileentry to multiple includes/files related ones, and i guess deferring them is fine.

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Record.h

 //===----------------------------------------------------------------------===//
 
 #ifndef CLANG_INCLUDE_CLEANER_RECORD_H
 #define CLANG_INCLUDE_CLEANER_RECORD_H
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/FileSystem/UniqueID.h"
+#include "clang-include-cleaner/Types.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/StringMap.h"
 #include <memory>
 #include <vector>
 
 namespace clang {
 class ASTConsumer;
 class ASTContext;
 class CompilerInstance;
 class Decl;
 class FileEntry;
+class Preprocessor;
+class PPCallbacks;
 
 namespace include_cleaner {
 
 /// Captures #include mapping information. It analyses IWYU Pragma comments and
 /// other use-instead-like mechanisms (#pragma include_instead) on included
 /// files.
 ///
 /// This is a low-level piece being used in the "Location => Header" analysis
private:
   llvm::DenseMap<llvm::sys::fs::UniqueID, std::string /*VerbatimSpelling*/>
       IWYUPublic;
 
   // FIXME: add other IWYU supports (export etc)
   // FIXME: add support for clang use_instead pragma
   // FIXME: add selfcontained file.
 };
 
-// Contains recorded parser events relevant to include-cleaner.
+/// Recorded main-file parser events relevant to include-cleaner.
 struct RecordedAST {
-  // The consumer (when installed into clang) tracks declarations in this.
+  /// The consumer (when installed into clang) tracks declarations in `*this`.
   std::unique_ptr<ASTConsumer> record();
 
   ASTContext *Ctx = nullptr;
-  // The set of declarations written at file scope inside the main file.
-  //
-  // These are the roots of the subtrees that should be traversed to find uses.
-  // (Traversing the TranslationUnitDecl would find uses inside headers!)
+  /// The set of declarations written at file scope inside the main file.
+  ///
+  /// These are the roots of the subtrees that should be traversed to find uses.
+  /// (Traversing the TranslationUnitDecl would find uses inside headers!)
   std::vector<Decl *> Roots;
 };
 
+/// Recorded main-file preprocessor events relevant to include-cleaner.
+///
+/// This doesn't include facts that we record globally for the whole TU, even
+/// when they occur in the main file (e.g. IWYU pragmas).
+struct RecordedPP {
+  /// The callback (when installed into clang) tracks macros/includes in this.
+  std::unique_ptr<PPCallbacks> record(const Preprocessor &PP);
+
+  /// Describes where macros were used in the main file.
+  std::vector<SymbolReference> MacroReferences;
+
+  /// A container for all includes present in the main file.
+  /// Supports efficiently hit-testing Headers against Includes.
+  /// FIXME: is there a more natural header for this class?
+  class RecordedIncludes {
+  public:
+    void add(const Include &);
+
+    /// All #includes seen, in the order they appear.
+    llvm::ArrayRef<Include> all() const { return All; }
+
+    /// Determine #includes that match a header (that provides a used symbol).
+    ///
+    /// Matching is based on the type of Header specified:
+    ///  - for a physical file like /path/to/foo.h, we check Resolved
+    ///  - for a logical file like <vector>, we check Spelled
+    llvm::SmallVector<const Include *> match(Header H) const;
+
+  private:
+    std::vector<Include> All;
+    // Lookup structures for match(), values are index into All.
+    llvm::StringMap<llvm::SmallVector<unsigned>> BySpelling;
+    llvm::DenseMap<const FileEntry *, llvm::SmallVector<unsigned>> ByFile;
+  } Includes;
+};
+
 } // namespace include_cleaner
 } // namespace clang
 
 #endif

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Types.h

 // a standard library symbol was defined doesn't matter, etc.
 //
 // We define some slightly more abstract sum types to handle these cases while
 // keeping the API clean. For example, Symbol may be a Decl AST node, a macro,
 // or a recognized standard library symbol.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef CLANG_INCLUDE_CLEANER_TYPES_H
 #define CLANG_INCLUDE_CLEANER_TYPES_H
 
+#include "clang/Basic/SourceLocation.h"
 #include "clang/Tooling/Inclusions/StandardLibrary.h"
 #include <memory>
 #include <vector>
 
 namespace llvm {
 class raw_ostream;
 } // namespace llvm
 namespace clang {
 class Decl;
 class FileEntry;
+class IdentifierInfo;
 namespace include_cleaner {
 
+/// We consider a macro to be a different symbol each time it is defined.
+struct Macro {
+  IdentifierInfo *Name;
+  /// The location of the Name where the macro is defined.
+  SourceLocation Definition;
+
+  bool operator==(const Macro &S) const {
+    return Definition == S.Definition;
+  }
+};
+
 /// An entity that can be referenced in the code.
 struct Symbol {
   enum Kind {
     /// A canonical clang declaration.
     Declaration,
+    /// A preprocessor macro, as defined in a specific location.
+    Macro,
     /// A recognized symbol from the standard library, like std::string.
     Standard,
   };
 
   Symbol(const Decl &D) : Storage(&D) {}
+  Symbol(struct Macro M) : Storage(M) {}
   Symbol(tooling::stdlib::Symbol S) : Storage(S) {}
 
   Kind kind() const { return static_cast<Kind>(Storage.index()); }
   bool operator==(const Symbol &RHS) const { return Storage == RHS.Storage; }
 
+  const Decl &declaration() const { return *std::get<Declaration>(Storage); }
+  struct Macro macro() const { return std::get<Macro>(Storage); }
   tooling::stdlib::Symbol standard() const {
     return std::get<Standard>(Storage);
   }
-  const Decl &declaration() const { return *std::get<Declaration>(Storage); }
 
 private:
   // FIXME: Add support for macros.
   // Order must match Kind enum!
-  std::variant<const Decl *, tooling::stdlib::Symbol> Storage;
+  std::variant<const Decl *, struct Macro, tooling::stdlib::Symbol> Storage;
 };
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const Symbol &);
 
+/// Indicates that a piece of code refers to a symbol.
+struct SymbolReference {
+  /// The symbol referred to.
+  Symbol Symbol;
+  /// The point in the code that refers to the symbol.
+  SourceLocation RefLocation;
+};
+llvm::raw_ostream &operator<<(llvm::raw_ostream &, const SymbolReference &);
+
 /// Represents a file that provides some symbol. Might not be includeable, e.g.
 /// built-in or main-file itself.
 struct Header {
   enum Kind {
     /// A source file parsed by clang. (May also be a <built-in> buffer).
     Physical,
     /// A recognized standard library header, like <string>.
     Standard,

 
 private:
   // FIXME: Handle verbatim spellings.
   // Order must match Kind enum!
   std::variant<const FileEntry *, tooling::stdlib::Header> Storage;
 };
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const Header &);
 
+/// A single #include directive written in the main file.
+struct Include {
+  llvm::StringRef Spelled;             // e.g. vector
+  const FileEntry *Resolved = nullptr; // e.g. /path/to/c++/v1/vector
+                                       // nullptr if the header was not found
+  SourceLocation HashLocation;         // of hash in #include <vector>
+  unsigned Line = 0;                   // 1-based line number for #include
+};
+llvm::raw_ostream &operator<<(llvm::raw_ostream &, const Include &);
+
 } // namespace include_cleaner
 } // namespace clang
 
 #endif
-

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

 //===--- Record.cpp - Record compiler events ------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "clang-include-cleaner/Record.h"
 #include "clang/AST/ASTConsumer.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/DeclGroup.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Frontend/CompilerInstance.h"
+#include "clang/Lex/MacroInfo.h"
 #include "clang/Lex/PPCallbacks.h"
 #include "clang/Lex/Preprocessor.h"
 
 namespace clang::include_cleaner {
+namespace {
+
+class PPRecorder : public PPCallbacks {
+public:
+  PPRecorder(RecordedPP &Recorded, const Preprocessor &PP)
+      : Recorded(Recorded), PP(PP), SM(PP.getSourceManager()) {}
+
+  void FileChanged(SourceLocation Loc, FileChangeReason Reason,
+                   SrcMgr::CharacteristicKind FileType,
+                   FileID PrevFID) override {
+    Active = SM.isWrittenInMainFile(Loc);
+  }
+
+  void InclusionDirective(SourceLocation Hash, const Token &IncludeTok,
+                          StringRef SpelledFilename, bool IsAngled,
+                          CharSourceRange FilenameRange,
+                          llvm::Optional<FileEntryRef> File,
+                          StringRef SearchPath, StringRef RelativePath,
+                          const Module *, SrcMgr::CharacteristicKind) override {
+    if (!Active)
+      return;
+
+    Include I;
+    I.HashLocation = Hash;
+    I.Resolved = File ? &File->getFileEntry() : nullptr;
+    I.Line = SM.getSpellingLineNumber(Hash);
+    I.Spelled = SpelledFilename;
+    Recorded.Includes.add(I);
+  }
+
+  void MacroExpands(const Token &MacroName, const MacroDefinition &MD,
+                    SourceRange Range, const MacroArgs *Args) override {
+    if (!Active)
+      return;
+    recordMacroRef(MacroName, *MD.getMacroInfo());
+  }
+
+  void MacroDefined(const Token &MacroName, const MacroDirective *MD) override {
+    if (!Active)
+      return;
+
+    const auto *MI = MD->getMacroInfo();
+    // The tokens of a macro definition could refer to a macro.
+    // Formally this reference isn't resolved until this macro is expanded,
+    // but we want to treat it as a reference anyway.
+    for (const auto &Tok : MI->tokens()) {
+      auto *II = Tok.getIdentifierInfo();
+      // Could this token be a reference to a macro? (Not param to this macro).
+      if (!II || !II->hadMacroDefinition() ||
+          llvm::is_contained(MI->params(), II))
+        continue;
+      if (const MacroInfo *MI = PP.getMacroInfo(II))
+        recordMacroRef(Tok, *MI);
+    }
+  }
+
+  void MacroUndefined(const Token &MacroName, const MacroDefinition &MD,
+                      const MacroDirective *) override {
+    if (!Active)
+      return;
+    if (const auto *MI = MD.getMacroInfo())
+      recordMacroRef(MacroName, *MI);
+  }
+
+private:
+  void recordMacroRef(const Token &Tok, const MacroInfo &MI) {
+    if (MI.isBuiltinMacro())
+      return; // __FILE__ is not a reference.
+    Recorded.MacroReferences.push_back(
+        SymbolReference{Macro{Tok.getIdentifierInfo(), MI.getDefinitionLoc()},
+                        Tok.getLocation()});
+  }
+
+  bool Active = false;
+  RecordedPP &Recorded;
+  const Preprocessor &PP;
+  const SourceManager &SM;
+};
+
+} // namespace
 
 // FIXME: this is a mirror of clang::clangd::parseIWYUPragma, move to libTooling
 // to share the code?
 static llvm::Optional<StringRef> parseIWYUPragma(const char *Text) {
   assert(strncmp(Text, "//", 2) || strncmp(Text, "/*", 2));
   constexpr llvm::StringLiteral IWYUPragma = " IWYU pragma: ";
   Text += 2; // Skip the comment start, // or /*.
   if (strncmp(Text, IWYUPragma.data(), IWYUPragma.size()))
     return llvm::None;
   Text += IWYUPragma.size();
   const char *End = Text;
   while (*End != 0 && *End != '\n')
     ++End;
   return StringRef(Text, End - Text);
 }
 
 class PragmaIncludes::RecordPragma : public PPCallbacks, public CommentHandler {
 public:
bool HandleTopLevelDecl(DeclGroupRef DG) override {
       }
       return ASTConsumer::HandleTopLevelDecl(DG);
     }
   };
 
   return std::make_unique<Recorder>(this);
 }
 
+void RecordedPP::RecordedIncludes::add(const Include &I) {
+  unsigned Index = All.size();
+  All.push_back(I);
+  auto BySpellingIt = BySpelling.try_emplace(I.Spelled).first;
+  All.back().Spelled = BySpellingIt->first(); // Now we own the backing string.
+
+  BySpellingIt->second.push_back(Index);
+  if (I.Resolved)
+    ByFile[I.Resolved].push_back(Index);
+}
+
+llvm::SmallVector<const Include *>
+RecordedPP::RecordedIncludes::match(Header H) const {
+  llvm::SmallVector<const Include *> Result;
+  switch (H.kind()) {
+  case Header::Physical:
+    for (unsigned I : ByFile.lookup(H.physical()))
+      Result.push_back(&All[I]);
+    break;
+  case Header::Standard:
+    for (unsigned I : BySpelling.lookup(H.standard().name().trim("<>")))
+      Result.push_back(&All[I]);
+    break;
+  }
+  return Result;
+}
+
+std::unique_ptr<PPCallbacks> RecordedPP::record(const Preprocessor &PP) {
+  return std::make_unique<PPRecorder>(*this, PP);
+}
+
+
 } // namespace clang::include_cleaner

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

 //===--- Types.cpp --------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "clang-include-cleaner/Types.h"
 #include "clang/AST/Decl.h"
 #include "clang/Basic/FileEntry.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace clang::include_cleaner {
 
 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const Symbol &S) {
   switch (S.kind()) {
   case Symbol::Declaration:
     if (const auto *ND = llvm::dyn_cast<NamedDecl>(&S.declaration()))
       return OS << ND->getNameAsString();
     return OS << S.declaration().getDeclKindName();
+  case Symbol::Macro:
+    return OS << S.macro().Name;
   case Symbol::Standard:
     return OS << S.standard().scope() << S.standard().name();
   }
   llvm_unreachable("Unhandled Symbol kind");
 }
 
 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const Header &H) {
   switch (H.kind()) {
   case Header::Physical:
     return OS << H.physical()->getName();
   case Header::Standard:
     return OS << H.standard().name();
   }
   llvm_unreachable("Unhandled Header kind");
 }
 
+llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const Include &I) {
+  return OS << I.Line << ": " << I.Spelled << " => "
+            << (I.Resolved ? I.Resolved->getName() : "<missing>");
+}
+
+llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const SymbolReference &R) {
+  // We can't decode the Location without SourceManager. Its raw representation
+  // isn't completely useless (and distinguishes SymbolReference from Symbol).
+  return OS << R.Symbol << "@0x"
+            << llvm::utohexstr(R.RefLocation.getRawEncoding(),
+                               /*Width=*/CHAR_BIT *
+                                   sizeof(SourceLocation::UIntTy));
+}
+
 } // namespace clang::include_cleaner

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

 //===-- RecordTest.cpp ----------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "clang-include-cleaner/Record.h"
 #include "clang/Frontend/FrontendAction.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Testing/TestAST.h"
+#include "clang/Tooling/Inclusions/StandardLibrary.h"
+#include "llvm/Support/VirtualFileSystem.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Testing/Support/Annotations.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace clang::include_cleaner {
 namespace {
+using testing::ElementsAre;
+using testing::ElementsAreArray;
+using testing::IsEmpty;
 
 // Matches a Decl* if it is a NamedDecl with the given name.
-MATCHER_P(Named, N, "") {
+MATCHER_P(named, N, "") {
   if (const NamedDecl *ND = llvm::dyn_cast<NamedDecl>(arg)) {
     if (N == ND->getNameAsString())
       return true;
   }
   std::string S;
   llvm::raw_string_ostream OS(S);
   arg->dump(OS);
   *result_listener << S;
Inputs.Code =
       namespace ns {
         int x;
         namespace {
           int y;
         }
       }
     )cpp";
   auto AST = build();
-  EXPECT_THAT(Recorded.Roots, testing::ElementsAre(Named("ns")));
+  EXPECT_THAT(Recorded.Roots, testing::ElementsAre(named("ns")));
 }
 
 // Decl in included file is not a root.
 TEST_F(RecordASTTest, Inclusion) {
   Inputs.ExtraFiles["header.h"] = "void headerFunc();";
   Inputs.Code = R"cpp(
     #include "header.h"
     void mainFunc();
   )cpp";
   auto AST = build();
-  EXPECT_THAT(Recorded.Roots, testing::ElementsAre(Named("mainFunc")));
+  EXPECT_THAT(Recorded.Roots, testing::ElementsAre(named("mainFunc")));
 }
 
 // Decl from macro expanded into the main file is a root.
 TEST_F(RecordASTTest, Macros) {
   Inputs.ExtraFiles["header.h"] = "#define X void x();";
   Inputs.Code = R"cpp(
     #include "header.h"
     X
   )cpp";
   auto AST = build();
-  EXPECT_THAT(Recorded.Roots, testing::ElementsAre(Named("x")));
+  EXPECT_THAT(Recorded.Roots, testing::ElementsAre(named("x")));
+}
+
+class RecordPPTest : public ::testing::Test {
+protected:
+  TestInputs Inputs;
+  RecordedPP Recorded;
+
+  RecordPPTest() {
+    struct RecordAction : public PreprocessOnlyAction {
+      RecordedPP &Out;
+      RecordAction(RecordedPP &Out) : Out(Out) {}
+
+      void ExecuteAction() override {
+        auto &PP = getCompilerInstance().getPreprocessor();
+        PP.addPPCallbacks(Out.record(PP));
+        PreprocessOnlyAction::ExecuteAction();
+      }
+    };
+    Inputs.MakeAction = [this] {
+      return std::make_unique<RecordAction>(Recorded);
+    };
+  }
+
+  TestAST build() { return TestAST(Inputs); }
+};
+
+// Matches an Include with a particular spelling.
+MATCHER_P(spelled, S, "") { return arg.Spelled == S; }
+
+TEST_F(RecordPPTest, CapturesIncludes) {
+  llvm::Annotations MainFile(R"cpp(
+    $H^#include "./header.h"
+    $M^#include "missing.h"
+  )cpp");
+  Inputs.Code = MainFile.code();
+  Inputs.ExtraFiles["header.h"] = "";
+  Inputs.ErrorOK = true; // missing header
+  auto AST = build();
+
+  ASSERT_THAT(
+      Recorded.Includes.all(),
+      testing::ElementsAre(spelled("./header.h"), spelled("missing.h")));
+
+  auto &H = Recorded.Includes.all().front();
+  EXPECT_EQ(H.Line, 2u);
+  EXPECT_EQ(H.HashLocation,
+            AST.sourceManager().getComposedLoc(
+                AST.sourceManager().getMainFileID(), MainFile.point("H")));
+  EXPECT_EQ(H.Resolved, AST.fileManager().getFile("header.h").get());
+
+  auto &M = Recorded.Includes.all().back();
+  EXPECT_EQ(M.Line, 3u);
+  EXPECT_EQ(M.HashLocation,
+            AST.sourceManager().getComposedLoc(
+                AST.sourceManager().getMainFileID(), MainFile.point("M")));
+  EXPECT_EQ(M.Resolved, nullptr);
+}
+
+TEST_F(RecordPPTest, CapturesMacroRefs) {
+  llvm::Annotations Header(R"cpp(
+    #define $def^X 1
+
+    // Refs, but not in main file.
+    #define Y X
+    int one = X;
+  )cpp");
+  llvm::Annotations MainFile(R"cpp(
+    #define EARLY X // not a ref, no definition
+    #include "header.h"
+    #define LATE ^X
+    #define LATE2 ^X // a ref even if not expanded
+
+    int uno = ^X;
+    int jeden = $exp^LATE; // a ref in LATE's expansion
+
+    #define IDENT(X) X // not a ref, shadowed
+    int eins = IDENT(^X);
+
+    #undef ^X
+    // Not refs, rather a new macro with the same name.
+    #define X 2
+    int two = X;
+  )cpp");
+  Inputs.Code = MainFile.code();
+  Inputs.ExtraFiles["header.h"] = Header.code();
+  auto AST = build();
+  const auto &SM = AST.sourceManager();
+
+  SourceLocation Def = SM.getComposedLoc(
+      SM.translateFile(AST.fileManager().getFile("header.h").get()),
+      Header.point("def"));
+  ASSERT_THAT(Recorded.MacroReferences, Not(IsEmpty()));
+  Symbol OrigX = Recorded.MacroReferences.front().Symbol;
+  EXPECT_EQ("X", OrigX.macro().Name->getName());
+  EXPECT_EQ(Def, OrigX.macro().Definition);
+
+  std::vector<unsigned> RefOffsets;
+  std::vector<unsigned> ExpOffsets; // Expansion locs of refs in macro locs.
+  std::vector<SourceLocation> RefMacroLocs;
+  for (const auto &Ref : Recorded.MacroReferences) {
+    if (Ref.Symbol == OrigX) {
+      auto [FID, Off] = SM.getDecomposedLoc(Ref.RefLocation);
+      if (FID == SM.getMainFileID()) {
+        RefOffsets.push_back(Off);
+      } else if (Ref.RefLocation.isMacroID() &&
+                 SM.isWrittenInMainFile(SM.getExpansionLoc(Ref.RefLocation))) {
+        ExpOffsets.push_back(
+            SM.getDecomposedExpansionLoc(Ref.RefLocation).second);
+      } else {
+        ADD_FAILURE() << Ref.RefLocation.printToString(SM);
+      }
+    }
+  }
+  EXPECT_THAT(RefOffsets, ElementsAreArray(MainFile.points()));
+  EXPECT_THAT(ExpOffsets, ElementsAreArray(MainFile.points("exp")));
+}
+
+// Matches an Include* on the specified line;
+MATCHER_P(line, N, "") { return arg->Line == (unsigned)N; }
+
+TEST(RecordedIncludesTest, Match) {
+  // We're using synthetic data, but need a FileManager to obtain FileEntry*s.
+  // Ensure it doesn't do any actual IO.
+  auto FS = llvm::makeIntrusiveRefCnt<llvm::vfs::InMemoryFileSystem>();
+  FileManager FM(FileSystemOptions{});
+  const FileEntry *A = FM.getVirtualFile("/path/a", /*Size=*/0, time_t{});
+  const FileEntry *B = FM.getVirtualFile("/path/b", /*Size=*/0, time_t{});
+
+  RecordedPP::RecordedIncludes Includes;
+  Includes.add(Include{"a", A, SourceLocation(), 1});
+  Includes.add(Include{"a2", A, SourceLocation(), 2});
+  Includes.add(Include{"b", B, SourceLocation(), 3});
+  Includes.add(Include{"vector", B, SourceLocation(), 4});
+  Includes.add(Include{"vector", B, SourceLocation(), 5});
+  Includes.add(Include{"missing", nullptr, SourceLocation(), 6});
+
+  EXPECT_THAT(Includes.match(A), ElementsAre(line(1), line(2)));
+  EXPECT_THAT(Includes.match(B), ElementsAre(line(3), line(4), line(5)));
+  EXPECT_THAT(Includes.match(*tooling::stdlib::Header::named("<vector>")),
+              ElementsAre(line(4), line(5)));
 }
 
 class PragmaIncludeTest : public ::testing::Test {
 protected:
   // We don't build an AST, we just run a preprocessor action!
   TestInputs Inputs;
   PragmaIncludes PI;
 
Inputs.MakeAction = [this] {
       return std::make_unique<Hook>(&PI);
     };
   }
 
   TestAST build() { return TestAST(Inputs); }
 };
 
 TEST_F(PragmaIncludeTest, IWYUKeep) {
   Inputs.Code = R"cpp(// Line 1
     #include "keep1.h" // IWYU pragma: keep
     #include "keep2.h" /* IWYU pragma: keep */
     #include "normal.h"
   )cpp";
   Inputs.ExtraFiles["keep1.h"] = Inputs.ExtraFiles["keep2.h"] =
       Inputs.ExtraFiles["normal.h"] = "";
 
   TestAST Processed = build();
TEST_F(PragmaIncludeTest, IWYUPrivate) {
   assert(PrivateFE);
   EXPECT_EQ(PI.getPublic(PrivateFE.get()), "\"public2.h\"");
   auto PublicFE = Processed.fileManager().getFile("public.h");
   assert(PublicFE);
   EXPECT_EQ(PI.getPublic(PublicFE.get()), ""); // no mapping.
 }
 
 } // namespace
 } // namespace clang::include_cleaner