[clangd] Handle go-to-definition in macro invocations where the target appears in the expansion multiple times

Fixes clangd/clangd#234

Reviewers: sammccall

Subscribers: ilya-biryukov, MaskRay, jkorous, arphaman, kadircet, usaxena95, cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D72041

Author: HighCommander4

clang-tools-extra/clangd/Selection.cpp

@@ -52,8 +52,7 @@ using ast_type_traits::DynTypedNode;
 // On traversing an AST node, its token range is erased from the unclaimed set.
 // The tokens actually removed are associated with that node, and hit-tested
 // against the selection to determine whether the node is selected.
-template <typename T>
-class IntervalSet {
+template <typename T> class IntervalSet {
 public:
   IntervalSet(llvm::ArrayRef<T> Range) { UnclaimedRanges.insert(Range); }
 
@@ -78,7 +77,7 @@ class IntervalSet {
       --Overlap.first;
       // ...unless B isn't selected at all.
       if (Overlap.first->end() <= Claim.begin())
-          ++Overlap.first;
+        ++Overlap.first;
     }
     if (Overlap.first == Overlap.second)
       return Out;
@@ -118,8 +117,7 @@ class IntervalSet {
   };
 
   // Disjoint sorted unclaimed ranges of expanded tokens.
-  std::set<llvm::ArrayRef<T>, RangeLess>
-      UnclaimedRanges;
+  std::set<llvm::ArrayRef<T>, RangeLess> UnclaimedRanges;
 };
 
 // Sentinel value for the selectedness of a node where we've seen no tokens yet.
@@ -148,11 +146,37 @@ bool shouldIgnore(const syntax::Token &Tok) {
   return Tok.kind() == tok::comment || Tok.kind() == tok::semi;
 }
 
+// Determine whether 'Target' is the first expansion of the macro
+// argument whose top-level spelling location is 'SpellingLoc'.
+bool isFirstExpansion(FileID Target, SourceLocation SpellingLoc,
+                      const SourceManager &SM) {
+  SourceLocation Prev = SpellingLoc;
+  while (true) {
+    // If the arg is expanded multiple times, getMacroArgExpandedLocation()
+    // returns the first expansion.
+    SourceLocation Next = SM.getMacroArgExpandedLocation(Prev);
+    // So if we reach the target, target is the first-expansion of the
+    // first-expansion ...
+    if (SM.getFileID(Next) == Target)
+      return true;
+
+    // Otherwise, if the FileID stops changing, we've reached the innermost
+    // macro expansion, and Target was on a different branch.
+    if (SM.getFileID(Next) == SM.getFileID(Prev))
+      return false;
+
+    Prev = Next;
+  }
+  return false;
+}
+
 // SelectionTester can determine whether a range of tokens from the PP-expanded
 // stream (corresponding to an AST node) is considered selected.
 //
 // When the tokens result from macro expansions, the appropriate tokens in the
 // main file are examined (macro invocation or args). Similarly for #includes.
+// However, only the first expansion of a given spelled token is considered
+// selected.
 //
 // It tests each token in the range (not just the endpoints) as contiguous
 // expanded tokens may not have contiguous spellings (with macros).
@@ -260,9 +284,14 @@ class SelectionTester {
     // Handle tokens that were passed as a macro argument.
     SourceLocation ArgStart = SM.getTopMacroCallerLoc(StartLoc);
     if (SM.getFileID(ArgStart) == SelFile) {
-      SourceLocation ArgEnd = SM.getTopMacroCallerLoc(Batch.back().location());
-      return testTokenRange(SM.getFileOffset(ArgStart),
-                            SM.getFileOffset(ArgEnd));
+      if (isFirstExpansion(FID, ArgStart, SM)) {
+        SourceLocation ArgEnd =
+            SM.getTopMacroCallerLoc(Batch.back().location());
+        return testTokenRange(SM.getFileOffset(ArgStart),
+                              SM.getFileOffset(ArgEnd));
+      } else {
+        /* fall through and treat as part of the macro body */
+      }
     }
 
     // Handle tokens produced by non-argument macro expansion.
@@ -346,7 +375,7 @@ std::string printNodeToString(const DynTypedNode &N, const PrintingPolicy &PP) {
 }
 #endif
 
-bool isImplicit(const Stmt* S) {
+bool isImplicit(const Stmt *S) {
   // Some Stmts are implicit and shouldn't be traversed, but there's no
   // "implicit" attribute on Stmt/Expr.
   // Unwrap implicit casts first if present (other nodes too?).
@@ -611,7 +640,7 @@ class SelectionVisitor : public RecursiveASTVisitor<SelectionVisitor> {
       // int (*[[s]])();
       else if (auto *VD = llvm::dyn_cast<VarDecl>(D))
         return VD->getLocation();
-    } else if (const auto* CCI = N.get<CXXCtorInitializer>()) {
+    } else if (const auto *CCI = N.get<CXXCtorInitializer>()) {
       // : [[b_]](42)
       return CCI->getMemberLocation();
     }
@@ -747,10 +776,10 @@ const Node *SelectionTree::commonAncestor() const {
   return Ancestor != Root ? Ancestor : nullptr;
 }
 
-const DeclContext& SelectionTree::Node::getDeclContext() const {
-  for (const Node* CurrentNode = this; CurrentNode != nullptr;
+const DeclContext &SelectionTree::Node::getDeclContext() const {
+  for (const Node *CurrentNode = this; CurrentNode != nullptr;
        CurrentNode = CurrentNode->Parent) {
-    if (const Decl* Current = CurrentNode->ASTNode.get<Decl>()) {
+    if (const Decl *Current = CurrentNode->ASTNode.get<Decl>()) {
       if (CurrentNode != this)
         if (auto *DC = dyn_cast<DeclContext>(Current))
           return *DC;

clang-tools-extra/clangd/Selection.h

@@ -64,6 +64,9 @@ namespace clangd {
 //
 // SelectionTree tries to behave sensibly in the presence of macros, but does
 // not model any preprocessor concepts: the output is a subset of the AST.
+// When a macro argument is specifically selected, only its first expansion is
+// selected in the AST. (Returning a selection forest is unreasonably difficult
+// for callers to handle correctly.)
 //
 // Comments, directives and whitespace are completely ignored.
 // Semicolons are also ignored, as the AST generally does not model them well.
@@ -127,15 +130,15 @@ class SelectionTree {
     Selection Selected;
     // Walk up the AST to get the DeclContext of this Node,
     // which is not the node itself.
-    const DeclContext& getDeclContext() const;
+    const DeclContext &getDeclContext() const;
     // Printable node kind, like "CXXRecordDecl" or "AutoTypeLoc".
     std::string kind() const;
     // If this node is a wrapper with no syntax (e.g. implicit cast), return
     // its contents. (If multiple wrappers are present, unwraps all of them).
-    const Node& ignoreImplicit() const;
+    const Node &ignoreImplicit() const;
     // If this node is inside a wrapper with no syntax (e.g. implicit cast),
     // return that wrapper. (If multiple are present, unwraps all of them).
-    const Node& outerImplicit() const;
+    const Node &outerImplicit() const;
   };
   // The most specific common ancestor of all the selected nodes.
   // Returns nullptr if the common ancestor is the root.

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

@@ -415,7 +415,7 @@ TEST(SelectionTest, CommonAncestor) {
 
 // Regression test: this used to match the injected X, not the outer X.
 TEST(SelectionTest, InjectedClassName) {
-  const char* Code = "struct ^X { int x; };";
+  const char *Code = "struct ^X { int x; };";
   auto AST = TestTU::withCode(Annotations(Code).code()).build();
   auto T = makeSelectionTree(Code, AST);
   ASSERT_EQ("CXXRecordDecl", nodeKind(T.commonAncestor())) << T;
@@ -508,7 +508,8 @@ TEST(SelectionTest, IncludedFile) {
 }
 
 TEST(SelectionTest, MacroArgExpansion) {
-  // If a macro arg is expanded several times, we consider them all selected.
+  // If a macro arg is expanded several times, we only consider the first one
+  // selected.
   const char *Case = R"cpp(
     int mul(int, int);
     #define SQUARE(X) mul(X, X);
@@ -517,15 +518,8 @@ TEST(SelectionTest, MacroArgExpansion) {
   Annotations Test(Case);
   auto AST = TestTU::withCode(Test.code()).build();
   auto T = makeSelectionTree(Case, AST);
-  // Unfortunately, this makes the common ancestor the CallExpr...
-  // FIXME: hack around this by picking one?
-  EXPECT_EQ("CallExpr", T.commonAncestor()->kind());
-  EXPECT_FALSE(T.commonAncestor()->Selected);
-  EXPECT_EQ(2u, T.commonAncestor()->Children.size());
-  for (const auto* N : T.commonAncestor()->Children) {
-    EXPECT_EQ("IntegerLiteral", N->kind());
-    EXPECT_TRUE(N->Selected);
-  }
+  EXPECT_EQ("IntegerLiteral", T.commonAncestor()->kind());
+  EXPECT_TRUE(T.commonAncestor()->Selected);
 
   // Verify that the common assert() macro doesn't suffer from this.
   // (This is because we don't associate the stringified token with the arg).
@@ -542,7 +536,7 @@ TEST(SelectionTest, MacroArgExpansion) {
 }
 
 TEST(SelectionTest, Implicit) {
-  const char* Test = R"cpp(
+  const char *Test = R"cpp(
     struct S { S(const char*); };
     int f(S);
     int x = f("^");

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

@@ -338,7 +338,18 @@ TEST(LocateSymbol, All) {
        #define ADDRESSOF(X) &X;
        int *j = ADDRESSOF(^i);
       )cpp",
-
+      R"cpp(// Macro argument appearing multiple times in expansion
+        #define VALIDATE_TYPE(x) (void)x;
+        #define ASSERT(expr)       \
+          do {                     \
+            VALIDATE_TYPE(expr);   \
+            if (!expr);            \
+          } while (false)
+        bool [[waldo]]() { return true; }
+        void foo() {
+          ASSERT(wa^ldo());
+        }
+      )cpp",
       R"cpp(// Symbol concatenated inside macro (not supported)
        int *pi;
        #define POINTER(X) p ## X;