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

[nridge]

Summary

Fixes https://github.com/clangd/clangd/issues/234

[merge_guards_bot]

Unit tests: pass. 61162 tests passed, 0 failed and 728 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[sammccall]

I like the goal here, a few high-level comments:

• it seems unfortunate to solve this mostly in xrefs rather than mostly in SelectionTree. All the other clients potentially want this behavior.
• the allSelectedNodes API on SelectionTree feels non-orthogonal: it's only meaningful here because the input happened to be a point and thus touches a single token and therefore node. If the input was a range, then this API doesn't help, because it loses the sense of "topmost" node.
• maybe the function we want walks down the tree as if finding the common ancestor, but may descend to return the roots of multiple subtrees if they are located in distinct expansions of the same macro arg
• or maybe we just want to treat non-first expansions as not-selected? That would about this issue I think

(BTW, Haojian is out the first few weeks of Jan. I'm still technically OOO, back later this week)

[nridge]

In D72041#1800189, @sammccall wrote:

• the allSelectedNodes API on SelectionTree feels non-orthogonal: it's only meaningful here because the input happened to be a point and thus touches a single token and therefore node. If the input was a range, then this API doesn't help, because it loses the sense of "topmost" node.

Good point!

• maybe the function we want walks down the tree as if finding the common ancestor, but may descend to return the roots of multiple subtrees if they are located in distinct expansions of the same macro arg
• or maybe we just want to treat non-first expansions as not-selected? That would about this issue I think

I'm thinking of going with the second suggestion, because it doesn't require changes to the SelectionTree API, and therefore addresses this concern:

• it seems unfortunate to solve this mostly in xrefs rather than mostly in SelectionTree. All the other clients potentially want this behavior.

without touching all the clients of SelectionTree.

[nridge]

Revise the approach to treat non-first expansions as not-selected

[merge_guards_bot]

Unit tests: pass. 61848 tests passed, 0 failed and 781 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[sammccall]

Thanks for fixing this. I think it could maybe use a further tweak, LMK if I'm missing something as my intuition around macros is pretty bad.

clang-tools-extra/clangd/Selection.cpp

259

Given the following program:

#define SQUARE(x) x * x;
int four = [[SQUARE(2)]];

We're going to now report the binary operator and one of the operands as selected and not the other, which doesn't seem desirable.

I think we want to accept macro-selected || arg-selected, so probably doing the current "non-argument macro expansion" first unconditionally or factoring it out into a function.

This will change the behavior of int four = [[SQUARE]](2) to consider the literal children selected too, I think this is fine.

259

I don't think it's a good idea to add hidden state and side-effects to testChunk() - it breaks a lot of assumptions that help reason about the code, and using mutable hides the violation of them.
(And a possible source of bugs - this is first in traversal order rather than first in source order - these are mostly but IIRC not always the same).

Instead I think you can do this statelessly: from the top-level spelling location, walk down with SM.getMacroArgExpandedLocation until you hit the target FileID (this is the first-expansion of first-expansion of first-expansion...) or the FileID stops changing (you've reached the innermost macro invocation, and your target location was on a different branch).

clang-tools-extra/clangd/Selection.h

58

Worth mentioning the new behavior here, IMO. e.g.

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).

[nridge]

clang-tools-extra/clangd/Selection.cpp

259

I agree that adding state is not great. I thought it was icky as I was writing it, I just couldn't think of an alternative. Thank you for suggesting one!

I implemented what you suggested, and it seems to work. I did want to ask a clarifying question to make sure I understand correctly: when an argument occurs multiple times in a macro exapnsion, the occurrences will have distinct FileIDs (as opposed just different offsets in the same macro FileID)?

[nridge]

Address review comments

[sammccall]

Awesome, thanks for fixing this!

clang-tools-extra/clangd/Selection.cpp

160

nit: return false directly?

259

That's right. My mental model is:

• a TU is a sequence of expanded (i.e. post-PP) tokens
• a FileID always corresponds to a single subrange of these expanded tokens...
• ...with holes in it for child FileIDs. Parents/children nest properly.

One of the things I really like about the syntax::TokenBuffer API is that it exposes this structure.

259

when false is the fallthrough to handling as if part of the macro body deliberate?

Thinking about it I suppose either that or returning NoTokens works, but please document it, e.g. } else { /* fall through and treat as part of the macro body */}

[nridge]

clang-tools-extra/clangd/Selection.cpp

259

It is deliberate. In fact, it is intended to address this previous comment. Please let me know if I've misunderstood and the solution doesn't match the request.

I'll add a comment as you suggested.

[sammccall]

clang-tools-extra/clangd/Selection.cpp

259

Yeah I think the behavior is great, as is the implementation :-) just a little surprising.

clang-tools-extra/clangd/Selection.h

 // Usually commonAncestor() is the place to start:
 //  - it's the simplest answer to "what node is under the cursor"
 //  - the selected Expr (for example) can be found by walking up the parent
 //    chain and checking Node->ASTNode.
 //  - if you want to traverse the selected nodes, they are all under
 //    commonAncestor() in the tree.
 //
 // 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.
 //
 // Comments, directives and whitespace are completely ignored.
 // Semicolons are also ignored, as the AST generally does not model them well.
 //
 // The SelectionTree owns the Node structures, but the ASTNode attributes
 // point back into the AST it was constructed with.
 class SelectionTree {
 public:
struct Node {
     // Direct children within the selection tree.
     llvm::SmallVector<const Node *, 8> Children;
     // The corresponding node from the full AST.
     ast_type_traits::DynTypedNode ASTNode;
     // The extent to which this node is covered by the selection.
     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.
   // (This is to avoid accidentally traversing the TUDecl and thus preamble).
   const Node *commonAncestor() const;
   // The selection node corresponding to TranslationUnitDecl.
   const Node &root() const { return *Root; }
+  // Get all individual selected nodes, in case the caller wants something
+  // other than their common ancestor.
+  std::vector<const Node *> allSelectedNodes() const;
 
 private:
   std::deque<Node> Nodes; // Stable-pointer storage.
   const Node *Root;
   clang::PrintingPolicy PrintPolicy;
 
   void print(llvm::raw_ostream &OS, const Node &N, int Indent) const;
   friend llvm::raw_ostream &operator<<(llvm::raw_ostream &OS,

clang-tools-extra/clangd/Selection.cpp

 //  Claim:          [---------------------------------------]
 //  Out:            [---------]                      [------]
 //  After:   [-----]                                         [-----------]
 //
 // It is used to track (expanded) tokens not yet associated with an AST node.
 // 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); }
 
   // Removes the elements of Claim from the set, modifying or removing ranges
   // that overlap it.
   // Returns the continuous subranges of Claim that were actually removed.
   llvm::SmallVector<llvm::ArrayRef<T>, 4> erase(llvm::ArrayRef<T> Claim) {
     llvm::SmallVector<llvm::ArrayRef<T>, 4> Out;
     if (Claim.empty())
       return Out;
 
     // General case:
     // Claim:                   [-----------------]
     // UnclaimedRanges: [-A-] [-B-] [-C-] [-D-] [-E-] [-F-] [-G-]
     // Overlap:               ^first                  ^second
     // Ranges C and D are fully included. Ranges B and E must be trimmed.
     auto Overlap = std::make_pair(
         UnclaimedRanges.lower_bound({Claim.begin(), Claim.begin()}), // C
         UnclaimedRanges.lower_bound({Claim.end(), Claim.end()}));    // F
     // Rewind to cover B.
     if (Overlap.first != UnclaimedRanges.begin()) {
       --Overlap.first;
       // ...unless B isn't selected at all.
       if (Overlap.first->end() <= Claim.begin())
-          ++Overlap.first;
     }
     if (Overlap.first == Overlap.second)
       return Out;
 
     // First, copy all overlapping ranges into the output.
     auto OutFirst = Out.insert(Out.end(), Overlap.first, Overlap.second);
     // If any of the overlapping ranges were sliced by the claim, split them:
     //  - restrict the returned range to the claimed part
private:
   using TokenRange = llvm::ArrayRef<T>;
   struct RangeLess {
     bool operator()(llvm::ArrayRef<T> L, llvm::ArrayRef<T> R) const {
       return L.begin() < R.begin();
     }
   };
 
   // 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.
 // This resolves to Unselected if no tokens are ever seen.
 // But Unselected + Complete -> Partial, while NoTokens + Complete --> Complete.
 // This value is never exposed publicly.
 constexpr SelectionTree::Selection NoTokens =
     static_cast<SelectionTree::Selection>(
         static_cast<unsigned char>(SelectionTree::Complete + 1));
 
 // Nodes start with NoTokens, and then use this function to aggregate the
 // selectedness as more tokens are found.
 void update(SelectionTree::Selection &Result, SelectionTree::Selection New) {
   if (New == NoTokens)
     return;
   if (Result == NoTokens)
     Result = New;
   else if (Result != New)
     // Can only be completely selected (or unselected) if all tokens are.
     Result = SelectionTree::Partial;
 }
 
-
 // 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.
 //
 // It tests each token in the range (not just the endpoints) as contiguous
 // expanded tokens may not have contiguous spellings (with macros).
 //
 // Non-token text, and tokens not modeled in the AST (comments, semicolons)
 // are ignored when determining selectedness.
 class SelectionTester {
 public:
   // The selection is offsets [SelBegin, SelEnd) in SelFile.
   SelectionTester(const syntax::TokenBuffer &Buf, FileID SelFile,
                   unsigned SelBegin, unsigned SelEnd, const SourceManager &SM)
       : SelFile(SelFile), SM(SM) {
     // Find all tokens (partially) selected in the file.
     auto AllSpelledTokens = Buf.spelledTokens(SelFile);
     const syntax::Token *SelFirst =
         llvm::partition_point(AllSpelledTokens, [&](const syntax::Token &Tok) {
           return SM.getFileOffset(Tok.endLocation()) <= SelBegin;
         });
     const syntax::Token *SelLimit = std::partition_point(
         SelFirst, AllSpelledTokens.end(), [&](const syntax::Token &Tok) {
           return SM.getFileOffset(Tok.location()) < SelEnd;
testChunk(FileID FID, llvm::ArrayRef<syntax::Token> Batch) const {
     }
 
     assert(StartLoc.isMacroID());
     // 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));
     }
 
     // Handle tokens produced by non-argument macro expansion.
     // Check if the macro name is selected, don't claim it exclusively.
     auto Expansion = SM.getDecomposedExpansionLoc(StartLoc);
     if (Expansion.first == SelFile)
       // FIXME: also check ( and ) for function-like macros?
       return testToken(Expansion.second);
std::string printNodeToString(const DynTypedNode &N, const PrintingPolicy &PP) {
   std::string S;
   llvm::raw_string_ostream OS(S);
   printNodeKind(OS, N);
   OS << " ";
   return std::move(OS.str());
 }
 #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?).
   if (auto *ICE = llvm::dyn_cast<ImplicitCastExpr>(S))
     S = ICE->getSubExprAsWritten();
   // Implicit this in a MemberExpr is not filtered out by RecursiveASTVisitor.
   // It would be nice if RAV handled this (!shouldTraverseImplicitCode()).
   if (auto *CTI = llvm::dyn_cast<CXXThisExpr>(S))
#endif
   SourceRange earlySourceRange(const DynTypedNode &N) {
     if (const Decl *D = N.get<Decl>()) {
       // void [[foo]]();
       if (auto *FD = llvm::dyn_cast<FunctionDecl>(D))
         return FD->getNameInfo().getSourceRange();
       // 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();
     }
     return SourceRange();
   }
 
   // Perform hit-testing of a complete Node against the selection.
   // This runs for every node in the AST, and must be fast in common cases.
const Node *SelectionTree::commonAncestor() const {
   while (Ancestor->Children.size() == 1 && !Ancestor->Selected)
     Ancestor = Ancestor->Children.front();
   // Returning nullptr here is a bit unprincipled, but it makes the API safer:
   // the TranslationUnitDecl contains all of the preamble, so traversing it is a
   // performance cliff. Callers can check for null and use root() if they want.
   return Ancestor != Root ? Ancestor : nullptr;
 }
 
-const DeclContext& SelectionTree::Node::getDeclContext() const {
-  for (const Node* CurrentNode = this; CurrentNode != nullptr;
+namespace {
+void gatherSelectedNodes(const Node *Root, std::vector<const Node *> &Result) {
+  if (Root->Selected) {
+    Result.push_back(Root);
+    return;
+  }
+  for (const Node *Child : Root->Children) {
+    gatherSelectedNodes(Child, Result);
+  }
+}
+
+} // namespace
+
+std::vector<const Node *> SelectionTree::allSelectedNodes() const {
+  std::vector<const Node *> Result;
+  gatherSelectedNodes(Root, Result);
+  return Result;
+}
+
+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;
       return *Current->getDeclContext();
     }
   }
   llvm_unreachable("A tree must always be rooted at TranslationUnitDecl.");
 }

clang-tools-extra/clangd/XRefs.cpp

std::vector<const Decl *> getDeclAtPosition(ParsedAST &AST, SourceLocation Pos,
   unsigned Offset;
   std::tie(FID, Offset) = AST.getSourceManager().getDecomposedSpellingLoc(Pos);
   SelectionTree Selection(AST.getASTContext(), AST.getTokens(), Offset);
   std::vector<const Decl *> Result;
   if (const SelectionTree::Node *N = Selection.commonAncestor()) {
     auto Decls = targetDecl(N->ASTNode, Relations);
     Result.assign(Decls.begin(), Decls.end());
   }
+  // If the common ancestor didn't have a target, the individual selected
+  // nodes still might. This can happen if the selection is in a macro
+  // invocation and is over a token that appears in multiple places in
+  // the macro expansion. As long as all occurrences map to the same target,
+  // return that target.
+  if (Result.empty()) {
+    auto SelectedNodes = Selection.allSelectedNodes();
+    if (SelectedNodes.size() <= 1)
+      return {};
+    for (const SelectionTree::Node *N : SelectedNodes) {
+      auto Decls = targetDecl(N->ASTNode, Relations);
+      if (Decls.size() != 1) {
+        return Result;
+      }
+      auto Candidate = Decls[0];
+      if (Result.empty()) {
+        Result.push_back(Decls[0]);
+      } else if (Result[0] != Candidate) {
+        return {};
+      }
+    }
+  }
   return Result;
 }
 
 llvm::Optional<Location> makeLocation(ASTContext &AST, SourceLocation TokLoc,
                                       llvm::StringRef TUPath) {
   const SourceManager &SourceMgr = AST.getSourceManager();
   const FileEntry *F = SourceMgr.getFileEntryForID(SourceMgr.getFileID(TokLoc));
   if (!F)

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

const char *Tests[] = {
        }
       )cpp",
 
       R"cpp(// Macro argument
        int [[i]];
        #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;
        int i = *POINTER(^i);
       )cpp",
 
       R"cpp(// Forward class declaration
         class Foo;