[clangd] Extend CollectMainFileMacros.

[hokein]

Summary

Extend the existing MainFileMacros structure:

• record more information (InConditionalDirective) in MacroOccurrence
• collect macro references inside macro body (fix a long-time FIXME)

So that the MainFileMacros preseve enough information, which allows a
just-in-time convertion to interop with include-cleaner::Macro for
include-cleaer features.

See the context in https://reviews.llvm.org/D146017.

[kadircet]

thanks!

clang-tools-extra/clangd/CollectMacros.h

48

nit: you can get SM out of PP, no need to pass both.

70

nit: can you actually move macro bodies out-of-line ? it's unfortunate that we need to build half of the project everytime we touch this header.

clang-tools-extra/clangd/ParsedAST.cpp

615

nit: we seem to be calling Clang->getPreprocessor() in many places, maybe extract it into a variable instead?

[hokein]

address comments

clang-tools-extra/clangd/CollectMacros.h

48

oh, right!

70

Done (thanks for the define-out-line code action!)

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);
   }
+  // FIXME: handle C++23 #elifdef, #elifndef
 
   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/ParsedAST.cpp

ParsedAST::build(llvm::StringRef Filename, const ParseInputs &Inputs,
   // with non-preamble macros below.
   MainFileMacros Macros;
   std::vector<PragmaMark> Marks;
   if (Preamble) {
     Macros = Patch->mainFileMacros();
     Marks = Patch->marks();
   }
   Clang->getPreprocessor().addPPCallbacks(
-      std::make_unique<CollectMainFileMacros>(Clang->getSourceManager(),
-                                              Macros));
+      std::make_unique<CollectMainFileMacros>(
+          Clang->getSourceManager(), Clang->getPreprocessor(), Macros));
 
   Clang->getPreprocessor().addPPCallbacks(
       collectPragmaMarksCallback(Clang->getSourceManager(), Marks));
 
   // Copy over the includes from the preamble, then combine with the
   // non-preamble includes below.
   CanonicalIncludes CanonIncludes;
   if (Preamble)

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 {
scanPreamble(llvm::StringRef Contents, const tooling::CompileCommand &Cmd) {
   const auto &SM = PP.getSourceManager();
   IncludeStructure Includes;
   Includes.collect(*Clang);
   ScannedPreamble SP;
   SP.Bounds = Bounds;
   PP.addPPCallbacks(
       std::make_unique<DirectiveCollector>(PP, SP.TextualDirectives));
   PP.addPPCallbacks(collectPragmaMarksCallback(SM, SP.Marks));
-  PP.addPPCallbacks(std::make_unique<CollectMainFileMacros>(SM, SP.Macros));
+  PP.addPPCallbacks(std::make_unique<CollectMainFileMacros>(SM, PP, SP.Macros));
   if (llvm::Error Err = Action.Execute())
     return std::move(Err);
   Action.EndSourceFile();
   SP.Includes = std::move(Includes.MainFileIncludes);
   llvm::append_range(SP.Lines, llvm::split(Contents, "\n"));
   return SP;
 }
 

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

 //===-- CollectMacrosTests.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 "AST.h"
 #include "Annotations.h"
 #include "CollectMacros.h"
+#include "Matchers.h"
 #include "SourceCode.h"
 #include "TestTU.h"
 #include "clang/Basic/SourceLocation.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include <vector>
 
 namespace clang {
 namespace clangd {
 namespace {
 
 using testing::UnorderedElementsAreArray;
 
+MATCHER_P(rangeIs, R, "") { return arg.Rng == R; }
+MATCHER(isDef, "") { return arg.IsDefinition; }
+MATCHER(inConditionalDirective, "") { return arg.InConditionalDirective; }
+
 TEST(CollectMainFileMacros, SelectedMacros) {
   // References of the same symbol must have the ranges with the same
   // name(integer). If there are N different symbols then they must be named
   // from 1 to N. Macros for which SymbolID cannot be computed must be named
-  // "Unknown".
+  // "Unknown". The payload of the annotation describes the extra bit
+  // information of the MacroOccurrence (e.g. $1(def) => IsDefinition).
   const char *Tests[] = {
       R"cpp(// Macros: Cursor on definition.
-        #define $1[[FOO]](x,y) (x + y)
+        #define $1(def)[[FOO]](x,y) (x + y)
         int main() { int x = $1[[FOO]]($1[[FOO]](3, 4), $1[[FOO]](5, 6)); }
       )cpp",
       R"cpp(
-        #define $1[[M]](X) X;
-        #define $2[[abc]] 123
+        #define $1(def)[[M]](X) X;
+        #define $2(def)[[abc]] 123
         int s = $1[[M]]($2[[abc]]);
       )cpp",
       // FIXME: Locating macro in duplicate definitions doesn't work. Enable
       // this once LocateMacro is fixed.
       // R"cpp(// Multiple definitions.
       //   #define $1[[abc]] 1
       //   int func1() { int a = $1[[abc]];}
       //   #undef $1[[abc]]
 
       //   #define $2[[abc]] 2
       //   int func2() { int a = $2[[abc]];}
       //   #undef $2[[abc]]
       // )cpp",
       R"cpp(
-        #ifdef $Unknown[[UNDEFINED]]
+        #ifdef $Unknown(condit)[[UNDEFINED]]
+        #endif
+
+        #ifndef $Unknown(condit)[[UNDEFINED]]
+        #endif
+
+        #if defined($Unknown(condit)[[UNDEFINED]])
         #endif
       )cpp",
       R"cpp(
-        #ifndef $Unknown[[abc]]
-        #define $1[[abc]]
-        #ifdef $1[[abc]]
+        #ifndef $Unknown(condit)[[abc]]
+        #define $1(def)[[abc]]
+        #ifdef $1(condit)[[abc]]
         #endif
         #endif
       )cpp",
       R"cpp(
         // Macros from token concatenations not included.
-        #define $1[[CONCAT]](X) X##A()
-        #define $2[[PREPEND]](X) MACRO##X()
-        #define $3[[MACROA]]() 123
+        #define $1(def)[[CONCAT]](X) X##A()
+        #define $2(def)[[PREPEND]](X) MACRO##X()
+        #define $3(def)[[MACROA]]() 123
         int B = $1[[CONCAT]](MACRO);
         int D = $2[[PREPEND]](A);
       )cpp",
       R"cpp(
-        // FIXME: Macro names in a definition are not detected.
-        #define $1[[MACRO_ARGS2]](X, Y) X Y
-        #define $2[[FOO]] BAR
-        #define $3[[BAR]] 1
+        #define $1(def)[[MACRO_ARGS2]](X, Y) X Y
+        #define $3(def)[[BAR]] 1
+        #define $2(def)[[FOO]] $3[[BAR]]
         int A = $2[[FOO]];
       )cpp"};
+  auto ExpectedResults = [](const Annotations &T, StringRef Name) {
+    std::vector<Matcher<MacroOccurrence>> ExpectedLocations;
+    for (const auto &[R, Bits] : T.rangesWithPayload(Name)) {
+      if (Bits == "def")
+        ExpectedLocations.push_back(testing::AllOf(rangeIs(R), isDef()));
+      else if (Bits == "condit")
+        ExpectedLocations.push_back(
+            testing::AllOf(rangeIs(R), inConditionalDirective()));
+      else
+        ExpectedLocations.push_back(testing::AllOf(rangeIs(R)));
+    }
+    return ExpectedLocations;
+  };
+
   for (const char *Test : Tests) {
     Annotations T(Test);
     auto AST = TestTU::withCode(T.code()).build();
     auto ActualMacroRefs = AST.getMacros();
     auto &SM = AST.getSourceManager();
     auto &PP = AST.getPreprocessor();
 
-    // Known macros.
-    for (int I = 1;; I++) {
-      const auto ExpectedRefs = T.ranges(llvm::to_string(I));
-      if (ExpectedRefs.empty())
-        break;
+    for (const auto &[Name, Ranges] : T.all_ranges()) {
+      if (Name == "Unknown") {
+        EXPECT_THAT(ActualMacroRefs.UnknownMacros,
+                    UnorderedElementsAreArray(ExpectedResults(T, "Unknown")))
+            << "Unknown macros doesn't match in " << Test;
+        continue;
+      }
 
-      auto Loc = sourceLocationInMainFile(SM, ExpectedRefs.begin()->start);
+      auto Loc = sourceLocationInMainFile(
+          SM, offsetToPosition(T.code(), Ranges.front().Begin));
       ASSERT_TRUE(bool(Loc));
       const auto *Id = syntax::spelledIdentifierTouching(*Loc, AST.getTokens());
       ASSERT_TRUE(Id);
       auto Macro = locateMacroAt(*Id, PP);
       assert(Macro);
       auto SID = getSymbolID(Macro->Name, Macro->Info, SM);
 
-      std::vector<Range> Ranges;
-      for (const auto &Ref : ActualMacroRefs.MacroRefs[SID])
-        Ranges.push_back(Ref.Rng);
-      EXPECT_THAT(ExpectedRefs, UnorderedElementsAreArray(Ranges))
-          << "Annotation=" << I << ", MacroName=" << Macro->Name
+      EXPECT_THAT(ActualMacroRefs.MacroRefs[SID],
+                  UnorderedElementsAreArray(ExpectedResults(T, Name)))
+          << "Annotation=" << Name << ", MacroName=" << Macro->Name
           << ", Test = " << Test;
     }
-    // Unknown macros.
-    std::vector<Range> Ranges;
-    for (const auto &Ref : AST.getMacros().UnknownMacros)
-      Ranges.push_back(Ref.Rng);
-    EXPECT_THAT(Ranges, UnorderedElementsAreArray(T.ranges("Unknown")))
-        << "Unknown macros doesn't match in " << Test;
   }
 }
 } // namespace
 } // namespace clangd
 } // namespace clang

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]]());
       }