1 //===--- Selection.cpp ----------------------------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "Selection.h"
10 #include "AST.h"
11 #include "SourceCode.h"
12 #include "support/Logger.h"
13 #include "support/Trace.h"
14 #include "clang/AST/ASTTypeTraits.h"
15 #include "clang/AST/Decl.h"
16 #include "clang/AST/DeclCXX.h"
17 #include "clang/AST/Expr.h"
18 #include "clang/AST/ExprCXX.h"
19 #include "clang/AST/PrettyPrinter.h"
20 #include "clang/AST/RecursiveASTVisitor.h"
21 #include "clang/AST/TypeLoc.h"
22 #include "clang/Basic/OperatorKinds.h"
23 #include "clang/Basic/SourceLocation.h"
24 #include "clang/Basic/SourceManager.h"
25 #include "clang/Basic/TokenKinds.h"
26 #include "clang/Lex/Lexer.h"
27 #include "clang/Tooling/Syntax/Tokens.h"
28 #include "llvm/ADT/STLExtras.h"
29 #include "llvm/ADT/StringExtras.h"
30 #include "llvm/Support/Casting.h"
31 #include "llvm/Support/raw_ostream.h"
32 #include <algorithm>
33 #include <string>
34 
35 namespace clang {
36 namespace clangd {
37 namespace {
38 using Node = SelectionTree::Node;
39 
40 // Measure the fraction of selections that were enabled by recovery AST.
41 void recordMetrics(const SelectionTree &S, const LangOptions &Lang) {
42   if (!trace::enabled())
43     return;
44   const char *LanguageLabel = Lang.CPlusPlus ? "C++" : Lang.ObjC ? "ObjC" : "C";
45   static constexpr trace::Metric SelectionUsedRecovery(
46       "selection_recovery", trace::Metric::Distribution, "language");
47   static constexpr trace::Metric RecoveryType(
48       "selection_recovery_type", trace::Metric::Distribution, "language");
49   const auto *Common = S.commonAncestor();
50   for (const auto *N = Common; N; N = N->Parent) {
51     if (const auto *RE = N->ASTNode.get<RecoveryExpr>()) {
52       SelectionUsedRecovery.record(1, LanguageLabel); // used recovery ast.
53       RecoveryType.record(RE->isTypeDependent() ? 0 : 1, LanguageLabel);
54       return;
55     }
56   }
57   if (Common)
58     SelectionUsedRecovery.record(0, LanguageLabel); // unused.
59 }
60 
61 SourceRange getSourceRange(const DynTypedNode &N) {
62   // MemberExprs to implicitly access anonymous fields should not claim any
63   // tokens for themselves. Given:
64   //   struct A { struct { int b; }; };
65   // The clang AST reports the following nodes for an access to b:
66   //   A().b;
67   //   [----] MemberExpr, base = A().<anonymous>, member = b
68   //   [----] MemberExpr: base = A(), member = <anonymous>
69   //   [-]    CXXConstructExpr
70   // For our purposes, we don't want the second MemberExpr to own any tokens,
71   // so we reduce its range to match the CXXConstructExpr.
72   // (It's not clear that changing the clang AST would be correct in general).
73   if (const auto *ME = N.get<MemberExpr>()) {
74     if (!ME->getMemberDecl()->getDeclName())
75       return ME->getBase()
76                  ? getSourceRange(DynTypedNode::create(*ME->getBase()))
77                  : SourceRange();
78   }
79   return N.getSourceRange();
80 }
81 
82 // An IntervalSet maintains a set of disjoint subranges of an array.
83 //
84 // Initially, it contains the entire array.
85 //           [-----------------------------------------------------------]
86 //
87 // When a range is erased(), it will typically split the array in two.
88 //  Claim:                     [--------------------]
89 //  after:   [----------------]                      [-------------------]
90 //
91 // erase() returns the segments actually erased. Given the state above:
92 //  Claim:          [---------------------------------------]
93 //  Out:            [---------]                      [------]
94 //  After:   [-----]                                         [-----------]
95 //
96 // It is used to track (expanded) tokens not yet associated with an AST node.
97 // On traversing an AST node, its token range is erased from the unclaimed set.
98 // The tokens actually removed are associated with that node, and hit-tested
99 // against the selection to determine whether the node is selected.
100 template <typename T> class IntervalSet {
101 public:
102   IntervalSet(llvm::ArrayRef<T> Range) { UnclaimedRanges.insert(Range); }
103 
104   // Removes the elements of Claim from the set, modifying or removing ranges
105   // that overlap it.
106   // Returns the continuous subranges of Claim that were actually removed.
107   llvm::SmallVector<llvm::ArrayRef<T>> erase(llvm::ArrayRef<T> Claim) {
108     llvm::SmallVector<llvm::ArrayRef<T>> Out;
109     if (Claim.empty())
110       return Out;
111 
112     // General case:
113     // Claim:                   [-----------------]
114     // UnclaimedRanges: [-A-] [-B-] [-C-] [-D-] [-E-] [-F-] [-G-]
115     // Overlap:               ^first                  ^second
116     // Ranges C and D are fully included. Ranges B and E must be trimmed.
117     auto Overlap = std::make_pair(
118         UnclaimedRanges.lower_bound({Claim.begin(), Claim.begin()}), // C
119         UnclaimedRanges.lower_bound({Claim.end(), Claim.end()}));    // F
120     // Rewind to cover B.
121     if (Overlap.first != UnclaimedRanges.begin()) {
122       --Overlap.first;
123       // ...unless B isn't selected at all.
124       if (Overlap.first->end() <= Claim.begin())
125         ++Overlap.first;
126     }
127     if (Overlap.first == Overlap.second)
128       return Out;
129 
130     // First, copy all overlapping ranges into the output.
131     auto OutFirst = Out.insert(Out.end(), Overlap.first, Overlap.second);
132     // If any of the overlapping ranges were sliced by the claim, split them:
133     //  - restrict the returned range to the claimed part
134     //  - save the unclaimed part so it can be reinserted
135     llvm::ArrayRef<T> RemainingHead, RemainingTail;
136     if (Claim.begin() > OutFirst->begin()) {
137       RemainingHead = {OutFirst->begin(), Claim.begin()};
138       *OutFirst = {Claim.begin(), OutFirst->end()};
139     }
140     if (Claim.end() < Out.back().end()) {
141       RemainingTail = {Claim.end(), Out.back().end()};
142       Out.back() = {Out.back().begin(), Claim.end()};
143     }
144 
145     // Erase all the overlapping ranges (invalidating all iterators).
146     UnclaimedRanges.erase(Overlap.first, Overlap.second);
147     // Reinsert ranges that were merely trimmed.
148     if (!RemainingHead.empty())
149       UnclaimedRanges.insert(RemainingHead);
150     if (!RemainingTail.empty())
151       UnclaimedRanges.insert(RemainingTail);
152 
153     return Out;
154   }
155 
156 private:
157   using TokenRange = llvm::ArrayRef<T>;
158   struct RangeLess {
159     bool operator()(llvm::ArrayRef<T> L, llvm::ArrayRef<T> R) const {
160       return L.begin() < R.begin();
161     }
162   };
163 
164   // Disjoint sorted unclaimed ranges of expanded tokens.
165   std::set<llvm::ArrayRef<T>, RangeLess> UnclaimedRanges;
166 };
167 
168 // Sentinel value for the selectedness of a node where we've seen no tokens yet.
169 // This resolves to Unselected if no tokens are ever seen.
170 // But Unselected + Complete -> Partial, while NoTokens + Complete --> Complete.
171 // This value is never exposed publicly.
172 constexpr SelectionTree::Selection NoTokens =
173     static_cast<SelectionTree::Selection>(
174         static_cast<unsigned char>(SelectionTree::Complete + 1));
175 
176 // Nodes start with NoTokens, and then use this function to aggregate the
177 // selectedness as more tokens are found.
178 void update(SelectionTree::Selection &Result, SelectionTree::Selection New) {
179   if (New == NoTokens)
180     return;
181   if (Result == NoTokens)
182     Result = New;
183   else if (Result != New)
184     // Can only be completely selected (or unselected) if all tokens are.
185     Result = SelectionTree::Partial;
186 }
187 
188 // As well as comments, don't count semicolons as real tokens.
189 // They're not properly claimed as expr-statement is missing from the AST.
190 bool shouldIgnore(const syntax::Token &Tok) {
191   return Tok.kind() == tok::comment || Tok.kind() == tok::semi;
192 }
193 
194 // Determine whether 'Target' is the first expansion of the macro
195 // argument whose top-level spelling location is 'SpellingLoc'.
196 bool isFirstExpansion(FileID Target, SourceLocation SpellingLoc,
197                       const SourceManager &SM) {
198   SourceLocation Prev = SpellingLoc;
199   while (true) {
200     // If the arg is expanded multiple times, getMacroArgExpandedLocation()
201     // returns the first expansion.
202     SourceLocation Next = SM.getMacroArgExpandedLocation(Prev);
203     // So if we reach the target, target is the first-expansion of the
204     // first-expansion ...
205     if (SM.getFileID(Next) == Target)
206       return true;
207 
208     // Otherwise, if the FileID stops changing, we've reached the innermost
209     // macro expansion, and Target was on a different branch.
210     if (SM.getFileID(Next) == SM.getFileID(Prev))
211       return false;
212 
213     Prev = Next;
214   }
215   return false;
216 }
217 
218 // SelectionTester can determine whether a range of tokens from the PP-expanded
219 // stream (corresponding to an AST node) is considered selected.
220 //
221 // When the tokens result from macro expansions, the appropriate tokens in the
222 // main file are examined (macro invocation or args). Similarly for #includes.
223 // However, only the first expansion of a given spelled token is considered
224 // selected.
225 //
226 // It tests each token in the range (not just the endpoints) as contiguous
227 // expanded tokens may not have contiguous spellings (with macros).
228 //
229 // Non-token text, and tokens not modeled in the AST (comments, semicolons)
230 // are ignored when determining selectedness.
231 class SelectionTester {
232 public:
233   // The selection is offsets [SelBegin, SelEnd) in SelFile.
234   SelectionTester(const syntax::TokenBuffer &Buf, FileID SelFile,
235                   unsigned SelBegin, unsigned SelEnd, const SourceManager &SM)
236       : SelFile(SelFile), SM(SM) {
237     // Find all tokens (partially) selected in the file.
238     auto AllSpelledTokens = Buf.spelledTokens(SelFile);
239     const syntax::Token *SelFirst =
240         llvm::partition_point(AllSpelledTokens, [&](const syntax::Token &Tok) {
241           return SM.getFileOffset(Tok.endLocation()) <= SelBegin;
242         });
243     const syntax::Token *SelLimit = std::partition_point(
244         SelFirst, AllSpelledTokens.end(), [&](const syntax::Token &Tok) {
245           return SM.getFileOffset(Tok.location()) < SelEnd;
246         });
247     auto Sel = llvm::makeArrayRef(SelFirst, SelLimit);
248     // Find which of these are preprocessed to nothing and should be ignored.
249     std::vector<bool> PPIgnored(Sel.size(), false);
250     for (const syntax::TokenBuffer::Expansion &X :
251          Buf.expansionsOverlapping(Sel)) {
252       if (X.Expanded.empty()) {
253         for (const syntax::Token &Tok : X.Spelled) {
254           if (&Tok >= SelFirst && &Tok < SelLimit)
255             PPIgnored[&Tok - SelFirst] = true;
256         }
257       }
258     }
259     // Precompute selectedness and offset for selected spelled tokens.
260     for (unsigned I = 0; I < Sel.size(); ++I) {
261       if (shouldIgnore(Sel[I]) || PPIgnored[I])
262         continue;
263       SpelledTokens.emplace_back();
264       Tok &S = SpelledTokens.back();
265       S.Offset = SM.getFileOffset(Sel[I].location());
266       if (S.Offset >= SelBegin && S.Offset + Sel[I].length() <= SelEnd)
267         S.Selected = SelectionTree::Complete;
268       else
269         S.Selected = SelectionTree::Partial;
270     }
271   }
272 
273   // Test whether a consecutive range of tokens is selected.
274   // The tokens are taken from the expanded token stream.
275   SelectionTree::Selection
276   test(llvm::ArrayRef<syntax::Token> ExpandedTokens) const {
277     if (SpelledTokens.empty())
278       return NoTokens;
279     SelectionTree::Selection Result = NoTokens;
280     while (!ExpandedTokens.empty()) {
281       // Take consecutive tokens from the same context together for efficiency.
282       FileID FID = SM.getFileID(ExpandedTokens.front().location());
283       auto Batch = ExpandedTokens.take_while([&](const syntax::Token &T) {
284         return SM.getFileID(T.location()) == FID;
285       });
286       assert(!Batch.empty());
287       ExpandedTokens = ExpandedTokens.drop_front(Batch.size());
288 
289       update(Result, testChunk(FID, Batch));
290     }
291     return Result;
292   }
293 
294   // Cheap check whether any of the tokens in R might be selected.
295   // If it returns false, test() will return NoTokens or Unselected.
296   // If it returns true, test() may return any value.
297   bool mayHit(SourceRange R) const {
298     if (SpelledTokens.empty())
299       return false;
300     auto B = SM.getDecomposedLoc(R.getBegin());
301     auto E = SM.getDecomposedLoc(R.getEnd());
302     if (B.first == SelFile && E.first == SelFile)
303       if (E.second < SpelledTokens.front().Offset ||
304           B.second > SpelledTokens.back().Offset)
305         return false;
306     return true;
307   }
308 
309 private:
310   // Hit-test a consecutive range of tokens from a single file ID.
311   SelectionTree::Selection
312   testChunk(FileID FID, llvm::ArrayRef<syntax::Token> Batch) const {
313     assert(!Batch.empty());
314     SourceLocation StartLoc = Batch.front().location();
315     // There are several possible categories of FileID depending on how the
316     // preprocessor was used to generate these tokens:
317     //   main file, #included file, macro args, macro bodies.
318     // We need to identify the main-file tokens that represent Batch, and
319     // determine whether we want to exclusively claim them. Regular tokens
320     // represent one AST construct, but a macro invocation can represent many.
321 
322     // Handle tokens written directly in the main file.
323     if (FID == SelFile) {
324       return testTokenRange(SM.getFileOffset(Batch.front().location()),
325                             SM.getFileOffset(Batch.back().location()));
326     }
327 
328     // Handle tokens in another file #included into the main file.
329     // Check if the #include is selected, but don't claim it exclusively.
330     if (StartLoc.isFileID()) {
331       for (SourceLocation Loc = Batch.front().location(); Loc.isValid();
332            Loc = SM.getIncludeLoc(SM.getFileID(Loc))) {
333         if (SM.getFileID(Loc) == SelFile)
334           // FIXME: use whole #include directive, not just the filename string.
335           return testToken(SM.getFileOffset(Loc));
336       }
337       return NoTokens;
338     }
339 
340     assert(StartLoc.isMacroID());
341     // Handle tokens that were passed as a macro argument.
342     SourceLocation ArgStart = SM.getTopMacroCallerLoc(StartLoc);
343     if (SM.getFileID(ArgStart) == SelFile) {
344       if (isFirstExpansion(FID, ArgStart, SM)) {
345         SourceLocation ArgEnd =
346             SM.getTopMacroCallerLoc(Batch.back().location());
347         return testTokenRange(SM.getFileOffset(ArgStart),
348                               SM.getFileOffset(ArgEnd));
349       } else { // NOLINT(llvm-else-after-return)
350         /* fall through and treat as part of the macro body */
351       }
352     }
353 
354     // Handle tokens produced by non-argument macro expansion.
355     // Check if the macro name is selected, don't claim it exclusively.
356     auto Expansion = SM.getDecomposedExpansionLoc(StartLoc);
357     if (Expansion.first == SelFile)
358       // FIXME: also check ( and ) for function-like macros?
359       return testToken(Expansion.second);
360     return NoTokens;
361   }
362 
363   // Is the closed token range [Begin, End] selected?
364   SelectionTree::Selection testTokenRange(unsigned Begin, unsigned End) const {
365     assert(Begin <= End);
366     // Outside the selection entirely?
367     if (End < SpelledTokens.front().Offset ||
368         Begin > SpelledTokens.back().Offset)
369       return SelectionTree::Unselected;
370 
371     // Compute range of tokens.
372     auto B = llvm::partition_point(
373         SpelledTokens, [&](const Tok &T) { return T.Offset < Begin; });
374     auto E = std::partition_point(
375         B, SpelledTokens.end(), [&](const Tok &T) { return T.Offset <= End; });
376 
377     // Aggregate selectedness of tokens in range.
378     bool ExtendsOutsideSelection = Begin < SpelledTokens.front().Offset ||
379                                    End > SpelledTokens.back().Offset;
380     SelectionTree::Selection Result =
381         ExtendsOutsideSelection ? SelectionTree::Unselected : NoTokens;
382     for (auto It = B; It != E; ++It)
383       update(Result, It->Selected);
384     return Result;
385   }
386 
387   // Is the token at `Offset` selected?
388   SelectionTree::Selection testToken(unsigned Offset) const {
389     // Outside the selection entirely?
390     if (Offset < SpelledTokens.front().Offset ||
391         Offset > SpelledTokens.back().Offset)
392       return SelectionTree::Unselected;
393     // Find the token, if it exists.
394     auto It = llvm::partition_point(
395         SpelledTokens, [&](const Tok &T) { return T.Offset < Offset; });
396     if (It != SpelledTokens.end() && It->Offset == Offset)
397       return It->Selected;
398     return NoTokens;
399   }
400 
401   struct Tok {
402     unsigned Offset;
403     SelectionTree::Selection Selected;
404   };
405   std::vector<Tok> SpelledTokens;
406   FileID SelFile;
407   const SourceManager &SM;
408 };
409 
410 // Show the type of a node for debugging.
411 void printNodeKind(llvm::raw_ostream &OS, const DynTypedNode &N) {
412   if (const TypeLoc *TL = N.get<TypeLoc>()) {
413     // TypeLoc is a hierarchy, but has only a single ASTNodeKind.
414     // Synthesize the name from the Type subclass (except for QualifiedTypeLoc).
415     if (TL->getTypeLocClass() == TypeLoc::Qualified)
416       OS << "QualifiedTypeLoc";
417     else
418       OS << TL->getType()->getTypeClassName() << "TypeLoc";
419   } else {
420     OS << N.getNodeKind().asStringRef();
421   }
422 }
423 
424 #ifndef NDEBUG
425 std::string printNodeToString(const DynTypedNode &N, const PrintingPolicy &PP) {
426   std::string S;
427   llvm::raw_string_ostream OS(S);
428   printNodeKind(OS, N);
429   OS << " ";
430   return std::move(OS.str());
431 }
432 #endif
433 
434 bool isImplicit(const Stmt *S) {
435   // Some Stmts are implicit and shouldn't be traversed, but there's no
436   // "implicit" attribute on Stmt/Expr.
437   // Unwrap implicit casts first if present (other nodes too?).
438   if (auto *ICE = llvm::dyn_cast<ImplicitCastExpr>(S))
439     S = ICE->getSubExprAsWritten();
440   // Implicit this in a MemberExpr is not filtered out by RecursiveASTVisitor.
441   // It would be nice if RAV handled this (!shouldTraverseImplicitCode()).
442   if (auto *CTI = llvm::dyn_cast<CXXThisExpr>(S))
443     if (CTI->isImplicit())
444       return true;
445   // Make sure implicit access of anonymous structs don't end up owning tokens.
446   if (auto *ME = llvm::dyn_cast<MemberExpr>(S)) {
447     if (auto *FD = llvm::dyn_cast<FieldDecl>(ME->getMemberDecl()))
448       if (FD->isAnonymousStructOrUnion())
449         // If Base is an implicit CXXThis, then the whole MemberExpr has no
450         // tokens. If it's a normal e.g. DeclRef, we treat the MemberExpr like
451         // an implicit cast.
452         return isImplicit(ME->getBase());
453   }
454   // Refs to operator() and [] are (almost?) always implicit as part of calls.
455   if (auto *DRE = llvm::dyn_cast<DeclRefExpr>(S)) {
456     if (auto *FD = llvm::dyn_cast<FunctionDecl>(DRE->getDecl())) {
457       switch (FD->getOverloadedOperator()) {
458       case OO_Call:
459       case OO_Subscript:
460         return true;
461       default:
462         break;
463       }
464     }
465   }
466   return false;
467 }
468 
469 // We find the selection by visiting written nodes in the AST, looking for nodes
470 // that intersect with the selected character range.
471 //
472 // While traversing, we maintain a parent stack. As nodes pop off the stack,
473 // we decide whether to keep them or not. To be kept, they must either be
474 // selected or contain some nodes that are.
475 //
476 // For simple cases (not inside macros) we prune subtrees that don't intersect.
477 class SelectionVisitor : public RecursiveASTVisitor<SelectionVisitor> {
478 public:
479   // Runs the visitor to gather selected nodes and their ancestors.
480   // If there is any selection, the root (TUDecl) is the first node.
481   static std::deque<Node> collect(ASTContext &AST,
482                                   const syntax::TokenBuffer &Tokens,
483                                   const PrintingPolicy &PP, unsigned Begin,
484                                   unsigned End, FileID File) {
485     SelectionVisitor V(AST, Tokens, PP, Begin, End, File);
486     V.TraverseAST(AST);
487     assert(V.Stack.size() == 1 && "Unpaired push/pop?");
488     assert(V.Stack.top() == &V.Nodes.front());
489     return std::move(V.Nodes);
490   }
491 
492   // We traverse all "well-behaved" nodes the same way:
493   //  - push the node onto the stack
494   //  - traverse its children recursively
495   //  - pop it from the stack
496   //  - hit testing: is intersection(node, selection) - union(children) empty?
497   //  - attach it to the tree if it or any children hit the selection
498   //
499   // Two categories of nodes are not "well-behaved":
500   //  - those without source range information, we don't record those
501   //  - those that can't be stored in DynTypedNode.
502   bool TraverseDecl(Decl *X) {
503     if (llvm::isa_and_nonnull<TranslationUnitDecl>(X))
504       return Base::TraverseDecl(X); // Already pushed by constructor.
505     // Base::TraverseDecl will suppress children, but not this node itself.
506     if (X && X->isImplicit())
507       return true;
508     return traverseNode(X, [&] { return Base::TraverseDecl(X); });
509   }
510   bool TraverseTypeLoc(TypeLoc X) {
511     return traverseNode(&X, [&] { return Base::TraverseTypeLoc(X); });
512   }
513   bool TraverseTemplateArgumentLoc(const TemplateArgumentLoc &X) {
514     return traverseNode(&X,
515                         [&] { return Base::TraverseTemplateArgumentLoc(X); });
516   }
517   bool TraverseNestedNameSpecifierLoc(NestedNameSpecifierLoc X) {
518     return traverseNode(
519         &X, [&] { return Base::TraverseNestedNameSpecifierLoc(X); });
520   }
521   bool TraverseConstructorInitializer(CXXCtorInitializer *X) {
522     return traverseNode(
523         X, [&] { return Base::TraverseConstructorInitializer(X); });
524   }
525   bool TraverseCXXBaseSpecifier(const CXXBaseSpecifier &X) {
526     return traverseNode(&X, [&] { return Base::TraverseCXXBaseSpecifier(X); });
527   }
528   bool TraverseAttr(Attr *X) {
529     return traverseNode(X, [&] { return Base::TraverseAttr(X); });
530   }
531   // Stmt is the same, but this form allows the data recursion optimization.
532   bool dataTraverseStmtPre(Stmt *X) {
533     if (!X || isImplicit(X))
534       return false;
535     auto N = DynTypedNode::create(*X);
536     if (canSafelySkipNode(N))
537       return false;
538     push(std::move(N));
539     if (shouldSkipChildren(X)) {
540       pop();
541       return false;
542     }
543     return true;
544   }
545   bool dataTraverseStmtPost(Stmt *X) {
546     pop();
547     return true;
548   }
549   // QualifiedTypeLoc is handled strangely in RecursiveASTVisitor: the derived
550   // TraverseTypeLoc is not called for the inner UnqualTypeLoc.
551   // This means we'd never see 'int' in 'const int'! Work around that here.
552   // (The reason for the behavior is to avoid traversing the nested Type twice,
553   // but we ignore TraverseType anyway).
554   bool TraverseQualifiedTypeLoc(QualifiedTypeLoc QX) {
555     return traverseNode<TypeLoc>(
556         &QX, [&] { return TraverseTypeLoc(QX.getUnqualifiedLoc()); });
557   }
558   // Uninteresting parts of the AST that don't have locations within them.
559   bool TraverseNestedNameSpecifier(NestedNameSpecifier *) { return true; }
560   bool TraverseType(QualType) { return true; }
561 
562   // The DeclStmt for the loop variable claims to cover the whole range
563   // inside the parens, this causes the range-init expression to not be hit.
564   // Traverse the loop VarDecl instead, which has the right source range.
565   bool TraverseCXXForRangeStmt(CXXForRangeStmt *S) {
566     return traverseNode(S, [&] {
567       return TraverseStmt(S->getInit()) && TraverseDecl(S->getLoopVariable()) &&
568              TraverseStmt(S->getRangeInit()) && TraverseStmt(S->getBody());
569     });
570   }
571   // OpaqueValueExpr blocks traversal, we must explicitly traverse it.
572   bool TraverseOpaqueValueExpr(OpaqueValueExpr *E) {
573     return traverseNode(E, [&] { return TraverseStmt(E->getSourceExpr()); });
574   }
575   // We only want to traverse the *syntactic form* to understand the selection.
576   bool TraversePseudoObjectExpr(PseudoObjectExpr *E) {
577     return traverseNode(E, [&] { return TraverseStmt(E->getSyntacticForm()); });
578   }
579 
580 private:
581   using Base = RecursiveASTVisitor<SelectionVisitor>;
582 
583   SelectionVisitor(ASTContext &AST, const syntax::TokenBuffer &Tokens,
584                    const PrintingPolicy &PP, unsigned SelBegin, unsigned SelEnd,
585                    FileID SelFile)
586       : SM(AST.getSourceManager()), LangOpts(AST.getLangOpts()),
587 #ifndef NDEBUG
588         PrintPolicy(PP),
589 #endif
590         TokenBuf(Tokens), SelChecker(Tokens, SelFile, SelBegin, SelEnd, SM),
591         UnclaimedExpandedTokens(Tokens.expandedTokens()) {
592     // Ensure we have a node for the TU decl, regardless of traversal scope.
593     Nodes.emplace_back();
594     Nodes.back().ASTNode = DynTypedNode::create(*AST.getTranslationUnitDecl());
595     Nodes.back().Parent = nullptr;
596     Nodes.back().Selected = SelectionTree::Unselected;
597     Stack.push(&Nodes.back());
598   }
599 
600   // Generic case of TraverseFoo. Func should be the call to Base::TraverseFoo.
601   // Node is always a pointer so the generic code can handle any null checks.
602   template <typename T, typename Func>
603   bool traverseNode(T *Node, const Func &Body) {
604     if (Node == nullptr)
605       return true;
606     auto N = DynTypedNode::create(*Node);
607     if (canSafelySkipNode(N))
608       return true;
609     push(DynTypedNode::create(*Node));
610     bool Ret = Body();
611     pop();
612     return Ret;
613   }
614 
615   // HIT TESTING
616   //
617   // We do rough hit testing on the way down the tree to avoid traversing
618   // subtrees that don't touch the selection (canSafelySkipNode), but
619   // fine-grained hit-testing is mostly done on the way back up (in pop()).
620   // This means children get to claim parts of the selection first, and parents
621   // are only selected if they own tokens that no child owned.
622   //
623   // Nodes *usually* nest nicely: a child's getSourceRange() lies within the
624   // parent's, and a node (transitively) owns all tokens in its range.
625   //
626   // Exception 1: child range claims tokens that should be owned by the parent.
627   //              e.g. in `void foo(int);`, the FunctionTypeLoc should own
628   //              `void (int)` but the parent FunctionDecl should own `foo`.
629   // To handle this case, certain nodes claim small token ranges *before*
630   // their children are traversed. (see earlySourceRange).
631   //
632   // Exception 2: siblings both claim the same node.
633   //              e.g. `int x, y;` produces two sibling VarDecls.
634   //                    ~~~~~ x
635   //                    ~~~~~~~~ y
636   // Here the first ("leftmost") sibling claims the tokens it wants, and the
637   // other sibling gets what's left. So selecting "int" only includes the left
638   // VarDecl in the selection tree.
639 
640   // An optimization for a common case: nodes outside macro expansions that
641   // don't intersect the selection may be recursively skipped.
642   bool canSafelySkipNode(const DynTypedNode &N) {
643     SourceRange S = getSourceRange(N);
644     if (auto *TL = N.get<TypeLoc>()) {
645       // FIXME: TypeLoc::getBeginLoc()/getEndLoc() are pretty fragile
646       // heuristics. We should consider only pruning critical TypeLoc nodes, to
647       // be more robust.
648 
649       // DeclTypeTypeLoc::getSourceRange() is incomplete, which would lead to
650       // failing
651       // to descend into the child expression.
652       // decltype(2+2);
653       // ~~~~~~~~~~~~~ <-- correct range
654       // ~~~~~~~~      <-- range reported by getSourceRange()
655       // ~~~~~~~~~~~~  <-- range with this hack(i.e, missing closing paren)
656       // FIXME: Alter DecltypeTypeLoc to contain parentheses locations and get
657       // rid of this patch.
658       if (auto DT = TL->getAs<DecltypeTypeLoc>())
659         S.setEnd(DT.getUnderlyingExpr()->getEndLoc());
660       // AttributedTypeLoc may point to the attribute's range, NOT the modified
661       // type's range.
662       if (auto AT = TL->getAs<AttributedTypeLoc>())
663         S = AT.getModifiedLoc().getSourceRange();
664     }
665     // SourceRange often doesn't manage to accurately cover attributes.
666     // Fortunately, attributes are rare.
667     if (llvm::any_of(getAttributes(N),
668                      [](const Attr *A) { return !A->isImplicit(); }))
669       return false;
670     if (!SelChecker.mayHit(S)) {
671       dlog("{1}skip: {0}", printNodeToString(N, PrintPolicy), indent());
672       dlog("{1}skipped range = {0}", S.printToString(SM), indent(1));
673       return true;
674     }
675     return false;
676   }
677 
678   // There are certain nodes we want to treat as leaves in the SelectionTree,
679   // although they do have children.
680   bool shouldSkipChildren(const Stmt *X) const {
681     // UserDefinedLiteral (e.g. 12_i) has two children (12 and _i).
682     // Unfortunately TokenBuffer sees 12_i as one token and can't split it.
683     // So we treat UserDefinedLiteral as a leaf node, owning the token.
684     return llvm::isa<UserDefinedLiteral>(X);
685   }
686 
687   // Pushes a node onto the ancestor stack. Pairs with pop().
688   // Performs early hit detection for some nodes (on the earlySourceRange).
689   void push(DynTypedNode Node) {
690     SourceRange Early = earlySourceRange(Node);
691     dlog("{1}push: {0}", printNodeToString(Node, PrintPolicy), indent());
692     Nodes.emplace_back();
693     Nodes.back().ASTNode = std::move(Node);
694     Nodes.back().Parent = Stack.top();
695     Nodes.back().Selected = NoTokens;
696     Stack.push(&Nodes.back());
697     claimRange(Early, Nodes.back().Selected);
698   }
699 
700   // Pops a node off the ancestor stack, and finalizes it. Pairs with push().
701   // Performs primary hit detection.
702   void pop() {
703     Node &N = *Stack.top();
704     dlog("{1}pop: {0}", printNodeToString(N.ASTNode, PrintPolicy), indent(-1));
705     claimRange(getSourceRange(N.ASTNode), N.Selected);
706     if (N.Selected == NoTokens)
707       N.Selected = SelectionTree::Unselected;
708     if (N.Selected || !N.Children.empty()) {
709       // Attach to the tree.
710       N.Parent->Children.push_back(&N);
711     } else {
712       // Neither N any children are selected, it doesn't belong in the tree.
713       assert(&N == &Nodes.back());
714       Nodes.pop_back();
715     }
716     Stack.pop();
717   }
718 
719   // Returns the range of tokens that this node will claim directly, and
720   // is not available to the node's children.
721   // Usually empty, but sometimes children cover tokens but shouldn't own them.
722   SourceRange earlySourceRange(const DynTypedNode &N) {
723     if (const Decl *D = N.get<Decl>()) {
724       // We want constructor name to be claimed by TypeLoc not the constructor
725       // itself. Similar for deduction guides, we rather want to select the
726       // underlying TypeLoc.
727       // FIXME: Unfortunately this doesn't work, even though RecursiveASTVisitor
728       // traverses the underlying TypeLoc inside DeclarationName, it is null for
729       // constructors.
730       if (isa<CXXConstructorDecl>(D) || isa<CXXDeductionGuideDecl>(D))
731         return SourceRange();
732       // This will capture Field, Function, MSProperty, NonTypeTemplateParm and
733       // VarDecls. We want the name in the declarator to be claimed by the decl
734       // and not by any children. For example:
735       // void [[foo]]();
736       // int (*[[s]])();
737       // struct X { int [[hash]] [32]; [[operator]] int();}
738       if (const auto *DD = llvm::dyn_cast<DeclaratorDecl>(D))
739         return DD->getLocation();
740     } else if (const auto *CCI = N.get<CXXCtorInitializer>()) {
741       // : [[b_]](42)
742       return CCI->getMemberLocation();
743     }
744     return SourceRange();
745   }
746 
747   // Perform hit-testing of a complete Node against the selection.
748   // This runs for every node in the AST, and must be fast in common cases.
749   // This is usually called from pop(), so we can take children into account.
750   // The existing state of Result is relevant (early/late claims can interact).
751   void claimRange(SourceRange S, SelectionTree::Selection &Result) {
752     for (const auto &ClaimedRange :
753          UnclaimedExpandedTokens.erase(TokenBuf.expandedTokens(S)))
754       update(Result, SelChecker.test(ClaimedRange));
755 
756     if (Result && Result != NoTokens)
757       dlog("{1}hit selection: {0}", S.printToString(SM), indent());
758   }
759 
760   std::string indent(int Offset = 0) {
761     // Cast for signed arithmetic.
762     int Amount = int(Stack.size()) + Offset;
763     assert(Amount >= 0);
764     return std::string(Amount, ' ');
765   }
766 
767   SourceManager &SM;
768   const LangOptions &LangOpts;
769 #ifndef NDEBUG
770   const PrintingPolicy &PrintPolicy;
771 #endif
772   const syntax::TokenBuffer &TokenBuf;
773   std::stack<Node *> Stack;
774   SelectionTester SelChecker;
775   IntervalSet<syntax::Token> UnclaimedExpandedTokens;
776   std::deque<Node> Nodes; // Stable pointers as we add more nodes.
777 };
778 
779 } // namespace
780 
781 llvm::SmallString<256> abbreviatedString(DynTypedNode N,
782                                          const PrintingPolicy &PP) {
783   llvm::SmallString<256> Result;
784   {
785     llvm::raw_svector_ostream OS(Result);
786     N.print(OS, PP);
787   }
788   auto Pos = Result.find('\n');
789   if (Pos != llvm::StringRef::npos) {
790     bool MoreText = !llvm::all_of(Result.str().drop_front(Pos), llvm::isSpace);
791     Result.resize(Pos);
792     if (MoreText)
793       Result.append(" …");
794   }
795   return Result;
796 }
797 
798 void SelectionTree::print(llvm::raw_ostream &OS, const SelectionTree::Node &N,
799                           int Indent) const {
800   if (N.Selected)
801     OS.indent(Indent - 1) << (N.Selected == SelectionTree::Complete ? '*'
802                                                                     : '.');
803   else
804     OS.indent(Indent);
805   printNodeKind(OS, N.ASTNode);
806   OS << ' ' << abbreviatedString(N.ASTNode, PrintPolicy) << "\n";
807   for (const Node *Child : N.Children)
808     print(OS, *Child, Indent + 2);
809 }
810 
811 std::string SelectionTree::Node::kind() const {
812   std::string S;
813   llvm::raw_string_ostream OS(S);
814   printNodeKind(OS, ASTNode);
815   return std::move(OS.str());
816 }
817 
818 // Decide which selections emulate a "point" query in between characters.
819 // If it's ambiguous (the neighboring characters are selectable tokens), returns
820 // both possibilities in preference order.
821 // Always returns at least one range - if no tokens touched, and empty range.
822 static llvm::SmallVector<std::pair<unsigned, unsigned>, 2>
823 pointBounds(unsigned Offset, const syntax::TokenBuffer &Tokens) {
824   const auto &SM = Tokens.sourceManager();
825   SourceLocation Loc = SM.getComposedLoc(SM.getMainFileID(), Offset);
826   llvm::SmallVector<std::pair<unsigned, unsigned>, 2> Result;
827   // Prefer right token over left.
828   for (const syntax::Token &Tok :
829        llvm::reverse(spelledTokensTouching(Loc, Tokens))) {
830     if (shouldIgnore(Tok))
831       continue;
832     unsigned Offset = Tokens.sourceManager().getFileOffset(Tok.location());
833     Result.emplace_back(Offset, Offset + Tok.length());
834   }
835   if (Result.empty())
836     Result.emplace_back(Offset, Offset);
837   return Result;
838 }
839 
840 bool SelectionTree::createEach(ASTContext &AST,
841                                const syntax::TokenBuffer &Tokens,
842                                unsigned Begin, unsigned End,
843                                llvm::function_ref<bool(SelectionTree)> Func) {
844   if (Begin != End)
845     return Func(SelectionTree(AST, Tokens, Begin, End));
846   for (std::pair<unsigned, unsigned> Bounds : pointBounds(Begin, Tokens))
847     if (Func(SelectionTree(AST, Tokens, Bounds.first, Bounds.second)))
848       return true;
849   return false;
850 }
851 
852 SelectionTree SelectionTree::createRight(ASTContext &AST,
853                                          const syntax::TokenBuffer &Tokens,
854                                          unsigned int Begin, unsigned int End) {
855   llvm::Optional<SelectionTree> Result;
856   createEach(AST, Tokens, Begin, End, [&](SelectionTree T) {
857     Result = std::move(T);
858     return true;
859   });
860   return std::move(*Result);
861 }
862 
863 SelectionTree::SelectionTree(ASTContext &AST, const syntax::TokenBuffer &Tokens,
864                              unsigned Begin, unsigned End)
865     : PrintPolicy(AST.getLangOpts()) {
866   // No fundamental reason the selection needs to be in the main file,
867   // but that's all clangd has needed so far.
868   const SourceManager &SM = AST.getSourceManager();
869   FileID FID = SM.getMainFileID();
870   PrintPolicy.TerseOutput = true;
871   PrintPolicy.IncludeNewlines = false;
872 
873   dlog("Computing selection for {0}",
874        SourceRange(SM.getComposedLoc(FID, Begin), SM.getComposedLoc(FID, End))
875            .printToString(SM));
876   Nodes = SelectionVisitor::collect(AST, Tokens, PrintPolicy, Begin, End, FID);
877   Root = Nodes.empty() ? nullptr : &Nodes.front();
878   recordMetrics(*this, AST.getLangOpts());
879   dlog("Built selection tree\n{0}", *this);
880 }
881 
882 const Node *SelectionTree::commonAncestor() const {
883   const Node *Ancestor = Root;
884   while (Ancestor->Children.size() == 1 && !Ancestor->Selected)
885     Ancestor = Ancestor->Children.front();
886   // Returning nullptr here is a bit unprincipled, but it makes the API safer:
887   // the TranslationUnitDecl contains all of the preamble, so traversing it is a
888   // performance cliff. Callers can check for null and use root() if they want.
889   return Ancestor != Root ? Ancestor : nullptr;
890 }
891 
892 const DeclContext &SelectionTree::Node::getDeclContext() const {
893   for (const Node *CurrentNode = this; CurrentNode != nullptr;
894        CurrentNode = CurrentNode->Parent) {
895     if (const Decl *Current = CurrentNode->ASTNode.get<Decl>()) {
896       if (CurrentNode != this)
897         if (auto *DC = dyn_cast<DeclContext>(Current))
898           return *DC;
899       return *Current->getLexicalDeclContext();
900     }
901   }
902   llvm_unreachable("A tree must always be rooted at TranslationUnitDecl.");
903 }
904 
905 const SelectionTree::Node &SelectionTree::Node::ignoreImplicit() const {
906   if (Children.size() == 1 &&
907       getSourceRange(Children.front()->ASTNode) == getSourceRange(ASTNode))
908     return Children.front()->ignoreImplicit();
909   return *this;
910 }
911 
912 const SelectionTree::Node &SelectionTree::Node::outerImplicit() const {
913   if (Parent && getSourceRange(Parent->ASTNode) == getSourceRange(ASTNode))
914     return Parent->outerImplicit();
915   return *this;
916 }
917 
918 } // namespace clangd
919 } // namespace clang
920