[clangd][wip] Get rid of the include-cleaner hacky macro-reference implementation

[hokein]

Summary

By extending and using the existing MainFileMacros structure.

• record more information (InConditionalDirective) in MacroOccurrence
• collect macro references inside macro body

When emiting include-cleaner diagnostics, we convert the result to the
RecordedPP structure just-in-time to interop with include-cleaner

Put all together into a patch (might need to split)
TODO: the newly-added test doesn't pass, because there are some other
issues in missing-include diagnostics.

clang-tools-extra/clangd/CollectMacros.h

 namespace clang {
 namespace clangd {
 
 struct MacroOccurrence {
   // Instead of storing SourceLocation, we have to store the token range because
   // SourceManager from preamble is not available when we build the AST.
   Range Rng;
   bool IsDefinition;
+  // True if the occurence is used in a conditional directive, e.g. #ifdef MACRO
+  bool InConditionalDirective;
 };
 
 struct MainFileMacros {
   llvm::StringSet<> Names;
   llvm::DenseMap<SymbolID, std::vector<MacroOccurrence>> MacroRefs;
   // Somtimes it is not possible to compute the SymbolID for the Macro, e.g. a
   // reference to an undefined macro. Store them separately, e.g. for semantic
   // highlighting.
   std::vector<MacroOccurrence> UnknownMacros;
   // Ranges skipped by the preprocessor due to being inactive.
   std::vector<Range> SkippedRanges;
 };
 
 /// Collects macro references (e.g. definitions, expansions) in the main file.
 /// It is used to:
 ///  - collect macros in the preamble section of the main file (in Preamble.cpp)
 ///  - collect macros after the preamble of the main file (in ParsedAST.cpp)
 class CollectMainFileMacros : public PPCallbacks {
 public:
-  explicit CollectMainFileMacros(const SourceManager &SM, MainFileMacros &Out)
-      : SM(SM), Out(Out) {}
+  explicit CollectMainFileMacros(const SourceManager &SM,
+                                 const Preprocessor &PP, MainFileMacros &Out)
+      : SM(SM), PP(PP), Out(Out) {}
 
   void FileChanged(SourceLocation Loc, FileChangeReason,
                    SrcMgr::CharacteristicKind, FileID) override {
     InMainFile = isInsideMainFile(Loc, SM);
   }
 
-  void MacroDefined(const Token &MacroName, const MacroDirective *MD) override {
-    add(MacroName, MD->getMacroInfo(), /*IsDefinition=*/true);
-  }
+  void MacroDefined(const Token &MacroName, const MacroDirective *MD) override;
 
   void MacroExpands(const Token &MacroName, const MacroDefinition &MD,
                     SourceRange Range, const MacroArgs *Args) override {
     add(MacroName, MD.getMacroInfo());
   }
 
   void MacroUndefined(const clang::Token &MacroName,
                       const clang::MacroDefinition &MD,
                       const clang::MacroDirective *Undef) override {
     add(MacroName, MD.getMacroInfo());
   }
 
   void Ifdef(SourceLocation Loc, const Token &MacroName,
              const MacroDefinition &MD) override {
-    add(MacroName, MD.getMacroInfo());
+    add(MacroName, MD.getMacroInfo(), /*IsDefinition=*/false,
+        /*InConditionalDirective=*/true);
   }
 
   void Ifndef(SourceLocation Loc, const Token &MacroName,
               const MacroDefinition &MD) override {
-    add(MacroName, MD.getMacroInfo());
+    add(MacroName, MD.getMacroInfo(), /*IsDefinition=*/false,
+        /*InConditionalDirective=*/true);
   }
 
   void Defined(const Token &MacroName, const MacroDefinition &MD,
                SourceRange Range) override {
-    add(MacroName, MD.getMacroInfo());
+    add(MacroName, MD.getMacroInfo(), /*IsDefinition=*/false,
+        /*InConditionalDirective=*/true);
   }
 
   void SourceRangeSkipped(SourceRange R, SourceLocation EndifLoc) override {
     if (!InMainFile)
       return;
     Position Begin = sourceLocToPosition(SM, R.getBegin());
     Position End = sourceLocToPosition(SM, R.getEnd());
     Out.SkippedRanges.push_back(Range{Begin, End});
   }
 
 private:
   void add(const Token &MacroNameTok, const MacroInfo *MI,
-           bool IsDefinition = false);
+           bool IsDefinition = false, bool InConditionalDirective = false);
   const SourceManager &SM;
+  const Preprocessor &PP;
   bool InMainFile = true;
   MainFileMacros &Out;
 };
 
 /// Represents a `#pragma mark` in the main file.
 ///
 /// There can be at most one pragma mark per line.
 struct PragmaMark {

clang-tools-extra/clangd/CollectMacros.cpp

 //===--- CollectMacros.cpp ---------------------------------------*- C++-*-===//
 //
 // 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 "CollectMacros.h"
 #include "AST.h"
 #include "clang/Basic/SourceLocation.h"
+#include "llvm/ADT/STLExtras.h"
 
 namespace clang {
 namespace clangd {
 
 void CollectMainFileMacros::add(const Token &MacroNameTok, const MacroInfo *MI,
-                                bool IsDefinition) {
+                                bool IsDefinition, bool InIfCondition) {
   if (!InMainFile)
     return;
   auto Loc = MacroNameTok.getLocation();
   if (Loc.isInvalid() || Loc.isMacroID())
     return;
 
   auto Name = MacroNameTok.getIdentifierInfo()->getName();
   Out.Names.insert(Name);
   auto Range = halfOpenToRange(
       SM, CharSourceRange::getCharRange(Loc, MacroNameTok.getEndLoc()));
   if (auto SID = getSymbolID(Name, MI, SM))
-    Out.MacroRefs[SID].push_back({Range, IsDefinition});
+    Out.MacroRefs[SID].push_back({Range, IsDefinition, InIfCondition});
   else
-    Out.UnknownMacros.push_back({Range, IsDefinition});
+    Out.UnknownMacros.push_back({Range, IsDefinition, InIfCondition});
 }
 
 class CollectPragmaMarks : public PPCallbacks {
 public:
   explicit CollectPragmaMarks(const SourceManager &SM,
                               std::vector<clangd::PragmaMark> &Out)
       : SM(SM), Out(Out) {}
 

 };
 
 std::unique_ptr<PPCallbacks>
 collectPragmaMarksCallback(const SourceManager &SM,
                            std::vector<PragmaMark> &Out) {
   return std::make_unique<CollectPragmaMarks>(SM, Out);
 }
 
+void CollectMainFileMacros::MacroDefined(const Token &MacroName,
+                                         const MacroDirective *MD) {
+
+  if (!InMainFile)
+    return;
+  const auto *MI = MD->getMacroInfo();
+  add(MacroName, MD->getMacroInfo(), true);
+  if (MI)
+    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))
+        add(Tok, MI);
+    }
+}
+
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/IncludeCleaner.cpp

if (auto CanonicalPath =
         return SpelledHeader;
     }
   }
   return include_cleaner::spellHeader(
       Provider, AST.getPreprocessor().getHeaderSearchInfo(), MainFile);
 }
 
 std::vector<include_cleaner::SymbolReference>
-collectMacroReferences(ParsedAST &AST) {
+convertMacroReferences(ParsedAST &AST) {
   const auto &SM = AST.getSourceManager();
-  //  FIXME: !!this is a hacky way to collect macro references.
-  std::vector<include_cleaner::SymbolReference> Macros;
+  const auto &MacroRefs = AST.getMacros();
   auto &PP = AST.getPreprocessor();
-  for (const syntax::Token &Tok :
-       AST.getTokens().spelledTokens(SM.getMainFileID())) {
-    auto Macro = locateMacroAt(Tok, PP);
+  std::vector<include_cleaner::SymbolReference> Macros;
+  for (const auto &MAndRefs : MacroRefs.MacroRefs) {
+    for (const auto &Ref : MAndRefs.second) {
+      auto L = sourceLocationInMainFile(SM, Ref.Rng.start);
+      if (!L) {
+        llvm::consumeError(L.takeError());
+        continue;
+      }
+      if (const auto *Tok =
+              syntax::spelledIdentifierTouching(*L, AST.getTokens())) {
+        auto Macro = locateMacroAt(*Tok, PP);
         if (!Macro)
           continue;
+
         if (auto DefLoc = Macro->Info->getDefinitionLoc(); DefLoc.isValid())
           Macros.push_back(
-          {Tok.location(),
-           include_cleaner::Macro{/*Name=*/PP.getIdentifierInfo(Tok.text(SM)),
-                                  DefLoc},
-           include_cleaner::RefType::Explicit});
+              {Tok->location(),
+               include_cleaner::Macro{
+                   /*Name=*/PP.getIdentifierInfo(Tok->text(SM)), DefLoc},
+               Ref.InConditionalDirective
+                   ? include_cleaner::RefType::Implicit
+                   : include_cleaner::RefType::Explicit});
+      }
+    }
   }
   return Macros;
 }
 
 llvm::StringRef getResolvedPath(const include_cleaner::Header &SymProvider) {
   switch (SymProvider.kind()) {
   case include_cleaner::Header::Physical:
     return SymProvider.physical()->tryGetRealPathName();

 IncludeCleanerFindings computeIncludeCleanerFindings(ParsedAST &AST) {
   const auto &SM = AST.getSourceManager();
   const auto &Includes = AST.getIncludeStructure();
   include_cleaner::Includes ConvertedIncludes =
       convertIncludes(SM, Includes.MainFileIncludes);
   const FileEntry *MainFile = SM.getFileEntryForID(SM.getMainFileID());
 
   std::vector<include_cleaner::SymbolReference> Macros =
-      collectMacroReferences(AST);
+      convertMacroReferences(AST);
   std::vector<MissingIncludeDiagInfo> MissingIncludes;
   llvm::DenseSet<IncludeStructure::HeaderID> Used;
   trace::Span Tracer("include_cleaner::walkUsed");
   include_cleaner::walkUsed(
       AST.getLocalTopLevelDecls(), /*MacroRefs=*/Macros,
       AST.getPragmaIncludes(), SM,
       [&](const include_cleaner::SymbolReference &Ref,
           llvm::ArrayRef<include_cleaner::Header> Providers) {

clang-tools-extra/clangd/ParsedAST.cpp

ParsedAST::build(llvm::StringRef Filename, const ParseInputs &Inputs,
   // (We can't *just* use the replayed includes, they don't have Resolved path).
   Includes.collect(*Clang);
   // Copy over the macros in the preamble region of the main file, and combine
   // with non-preamble macros below.
   MainFileMacros Macros;
   if (Preamble)
     Macros = Preamble->Macros;
   Clang->getPreprocessor().addPPCallbacks(
-      std::make_unique<CollectMainFileMacros>(Clang->getSourceManager(),
-                                              Macros));
+      std::make_unique<CollectMainFileMacros>(
+          Clang->getSourceManager(), Clang->getPreprocessor(), Macros));
 
   std::vector<PragmaMark> Marks;
   // FIXME: We need to patch the marks for stale preambles.
   if (Preamble)
     Marks = Preamble->Marks;
   Clang->getPreprocessor().addPPCallbacks(
       collectPragmaMarksCallback(Clang->getSourceManager(), Marks));
 

clang-tools-extra/clangd/Preamble.cpp

if (Stats) {
       Stats->BuildSize += PP.getHeaderSearchInfo().getTotalMemory();
     }
   }
 
   void BeforeExecute(CompilerInstance &CI) override {
     CanonIncludes.addSystemHeadersMapping(CI.getLangOpts());
     LangOpts = &CI.getLangOpts();
     SourceMgr = &CI.getSourceManager();
+    PP = &CI.getPreprocessor();
     Includes.collect(CI);
     if (Config::current().Diagnostics.UnusedIncludes ==
                 Config::IncludesPolicy::Strict ||
             Config::current().Diagnostics.MissingIncludes ==
             Config::IncludesPolicy::Strict)
       Pragmas.record(CI);
     if (BeforeExecuteCallback)
       BeforeExecuteCallback(CI);
   }
 
   std::unique_ptr<PPCallbacks> createPPCallbacks() override {
-    assert(SourceMgr && LangOpts &&
-           "SourceMgr and LangOpts must be set at this point");
+    assert(SourceMgr && LangOpts && PP &&
+           "SourceMgr, LangOpts and PP must be set at this point");
 
     return std::make_unique<PPChainedCallbacks>(
-        std::make_unique<CollectMainFileMacros>(*SourceMgr, Macros),
+        std::make_unique<CollectMainFileMacros>(*SourceMgr, *PP, Macros),
         collectPragmaMarksCallback(*SourceMgr, Marks));
   }
 
   CommentHandler *getCommentHandler() override {
     IWYUHandler = collectIWYUHeaderMaps(&CanonIncludes);
     return IWYUHandler.get();
   }
 
private:
   CanonicalIncludes CanonIncludes;
   include_cleaner::PragmaIncludes Pragmas;
   MainFileMacros Macros;
   std::vector<PragmaMark> Marks;
   bool IsMainFileIncludeGuarded = false;
   std::unique_ptr<CommentHandler> IWYUHandler = nullptr;
   const clang::LangOptions *LangOpts = nullptr;
   const SourceManager *SourceMgr = nullptr;
+  const Preprocessor *PP = nullptr;
   PreambleBuildStats *Stats;
   bool ParseForwardingFunctions;
   std::function<void(CompilerInstance &)> BeforeExecuteCallback;
 };
 
 // Represents directives other than includes, where basic textual information is
 // enough.
 struct TextualPPDirective {

clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp

 #include "clang/AST/DeclBase.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Tooling/Syntax/Tokens.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ScopedPrinter.h"
+#include "llvm/Testing/Annotations/Annotations.h"
 #include "llvm/Testing/Support/SupportHelpers.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <cstddef>
 #include <string>
 #include <vector>
 
 namespace clang {

 Matcher<const Diag &> withFix(::testing::Matcher<Fix> FixMatcher) {
   return Field(&Diag::Fixes, ElementsAre(FixMatcher));
 }
 
 MATCHER_P2(Diag, Range, Message,
            "Diag at " + llvm::to_string(Range) + " = [" + Message + "]") {
   return arg.Range == Range && arg.Message == Message;
 }
-
+MATCHER_P(Range, R, "") {
+  return arg.SymRefRange == R;
+}
+MATCHER_P(SingleHeader, Header, "") {
+  return arg.Providers.size() == 1 && arg.Providers[0] == Header;
+}
 MATCHER_P3(Fix, Range, Replacement, Message,
            "Fix " + llvm::to_string(Range) + " => " +
                ::testing::PrintToString(Replacement) + " = [" + Message + "]") {
   return arg.Message == Message && arg.Edits.size() == 1 &&
          arg.Edits[0].range == Range && arg.Edits[0].newText == Replacement;
 }
 
 std::string guard(llvm::StringRef Code) {
TEST(IncludeCleaner, ComputeMissingHeaders) {
   size_t End = llvm::cantFail(positionToOffset(MainFile.code(), Range.end));
   syntax::FileRange BRange{SM.getMainFileID(), static_cast<unsigned int>(Start),
                            static_cast<unsigned int>(End)};
   include_cleaner::Header Header{*SM.getFileManager().getFile("b.h")};
   MissingIncludeDiagInfo BInfo{B, BRange, {Header}};
   EXPECT_THAT(Findings.MissingIncludes, ElementsAre(BInfo));
 }
 
+TEST(IncludeCleaner, MissingIncludesForMacroRefs) {
+  Config Cfg;
+  Cfg.Diagnostics.MissingIncludes = Config::IncludesPolicy::Strict;
+  WithContextValue Ctx(Config::Key, std::move(Cfg));
+  llvm::StringRef TestCases[] = {
+      R"cpp(
+        #include "all.h"
+        int a = [[FOO]];
+      )cpp",
+
+      R"cpp(
+        #include "all.h"
+        #define BAR [[FOO]];
+      )cpp",
+
+      // No missing include insertion for ambiguous macro refs.
+      R"cpp(
+       #include "all.h"
+       #if defined(FOO)
+       #endif
+
+       #ifdef FOO
+       #endif
+      )cpp",
+  };
+
+  for (const auto& T : TestCases) {
+    TestTU TU;
+    llvm::Annotations MainFile = T;
+    TU.Code = MainFile.code();
+    TU.AdditionalFiles["all.h"] = guard("#include \"foo.h\"");
+    TU.AdditionalFiles["foo.h"] = guard("#define FOO 1");
+    ParsedAST AST = TU.build();
+    const auto &SM = AST.getSourceManager();
+    include_cleaner::Header FooH = *SM.getFileManager().getFile("foo.h");
+
+    std::vector<testing::Matcher<MissingIncludeDiagInfo>> Matcher;
+    for (const auto &R : MainFile.ranges())
+      Matcher.push_back(
+          AllOf(Range(syntax::FileRange(AST.getSourceManager().getMainFileID(),
+                                        R.Begin, R.End)),
+                SingleHeader(FooH)));
+    EXPECT_THAT(computeIncludeCleanerFindings(AST).MissingIncludes,
+                testing::UnorderedElementsAreArray(Matcher));
+  }
+}
+
 TEST(IncludeCleaner, GenerateMissingHeaderDiags) {
   Config Cfg;
   Cfg.Diagnostics.MissingIncludes = Config::IncludesPolicy::Strict;
   Cfg.Diagnostics.Includes.IgnoreHeader = {
       [](llvm::StringRef Header) { return Header.ends_with("buzz.h"); }};
   WithContextValue Ctx(Config::Key, std::move(Cfg));
   Annotations MainFile(R"cpp(
 #include "a.h"

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

)cpp",
       #define $Macro_decl[[DEF_CLASS]](T) class T {};
       // Preamble ends.
       $Macro[[DEF_MULTIPLE]](XYZ);
       $Macro[[DEF_MULTIPLE]](XYZW);
       $Macro[[DEF_CLASS]]($Class_def[[A]])
       #define $Macro_decl[[MACRO_CONCAT]](X, V, T) T foo##X = V
       #define $Macro_decl[[DEF_VAR]](X, V) int X = V
       #define $Macro_decl[[DEF_VAR_T]](T, X, V) T X = V
-      #define $Macro_decl[[DEF_VAR_REV]](V, X) DEF_VAR(X, V)
+      #define $Macro_decl[[DEF_VAR_REV]](V, X) $Macro[[DEF_VAR]](X, V)
       #define $Macro_decl[[CPY]](X) X
       #define $Macro_decl[[DEF_VAR_TYPE]](X, Y) X Y
       #define $Macro_decl[[SOME_NAME]] variable
       #define $Macro_decl[[SOME_NAME_SET]] variable2 = 123
       #define $Macro_decl[[INC_VAR]](X) X += 2
       void $Function_def[[foo]]() {
         $Macro[[DEF_VAR]]($LocalVariable_def[[X]],  123);
         $Macro[[DEF_VAR_REV]](908, $LocalVariable_def[[XY]]);
)cpp",
       #define $Macro_decl[[CALL_FN]](F) F();
       #define $Macro_decl[[DEF_FN]](F) void F ()
       $Macro[[DEF_FN]]($Function_def[[g]]) {
         $Macro[[CALL_FN]]($Function[[foo]]);
       }
     )cpp",
       R"cpp(
       #define $Macro_decl[[fail]](expr) expr
-      #define $Macro_decl[[assert]](COND) if (!(COND)) { fail("assertion failed" #COND); }
+      #define $Macro_decl[[assert]](COND) if (!(COND)) { $Macro[[fail]]("assertion failed" #COND); }
       // Preamble ends.
       int $Variable_def[[x]];
       int $Variable_def[[y]];
       int $Function_decl[[f]]();
       void $Function_def[[foo]]() {
         $Macro[[assert]]($Variable[[x]] $Operator[[!=]] $Variable[[y]]);
         $Macro[[assert]]($Variable[[x]] $Operator[[!=]] $Function[[f]]());
       }