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 "support/Logger.h"
12 #include "support/Trace.h"
13 #include "clang/AST/ASTConcept.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/BitVector.h"
29 #include "llvm/ADT/STLExtras.h"
30 #include "llvm/ADT/StringExtras.h"
31 #include "llvm/Support/Casting.h"
32 #include "llvm/Support/raw_ostream.h"
33 #include <algorithm>
34 #include <set>
35 #include <string>
36 
37 namespace clang {
38 namespace clangd {
39 namespace {
40 using Node = SelectionTree::Node;
41 
42 // Measure the fraction of selections that were enabled by recovery AST.
recordMetrics(const SelectionTree & S,const LangOptions & Lang)43 void recordMetrics(const SelectionTree &S, const LangOptions &Lang) {
44   if (!trace::enabled())
45     return;
46   const char *LanguageLabel = Lang.CPlusPlus ? "C++" : Lang.ObjC ? "ObjC" : "C";
47   static constexpr trace::Metric SelectionUsedRecovery(
48       "selection_recovery", trace::Metric::Distribution, "language");
49   static constexpr trace::Metric RecoveryType(
50       "selection_recovery_type", trace::Metric::Distribution, "language");
51   const auto *Common = S.commonAncestor();
52   for (const auto *N = Common; N; N = N->Parent) {
53     if (const auto *RE = N->ASTNode.get<RecoveryExpr>()) {
54       SelectionUsedRecovery.record(1, LanguageLabel); // used recovery ast.
55       RecoveryType.record(RE->isTypeDependent() ? 0 : 1, LanguageLabel);
56       return;
57     }
58   }
59   if (Common)
60     SelectionUsedRecovery.record(0, LanguageLabel); // unused.
61 }
62 
63 // Return the range covering a node and all its children.
getSourceRange(const DynTypedNode & N)64 SourceRange getSourceRange(const DynTypedNode &N) {
65   // MemberExprs to implicitly access anonymous fields should not claim any
66   // tokens for themselves. Given:
67   //   struct A { struct { int b; }; };
68   // The clang AST reports the following nodes for an access to b:
69   //   A().b;
70   //   [----] MemberExpr, base = A().<anonymous>, member = b
71   //   [----] MemberExpr: base = A(), member = <anonymous>
72   //   [-]    CXXConstructExpr
73   // For our purposes, we don't want the second MemberExpr to own any tokens,
74   // so we reduce its range to match the CXXConstructExpr.
75   // (It's not clear that changing the clang AST would be correct in general).
76   if (const auto *ME = N.get<MemberExpr>()) {
77     if (!ME->getMemberDecl()->getDeclName())
78       return ME->getBase()
79                  ? getSourceRange(DynTypedNode::create(*ME->getBase()))
80                  : SourceRange();
81   }
82   return N.getSourceRange();
83 }
84 
85 // An IntervalSet maintains a set of disjoint subranges of an array.
86 //
87 // Initially, it contains the entire array.
88 //           [-----------------------------------------------------------]
89 //
90 // When a range is erased(), it will typically split the array in two.
91 //  Claim:                     [--------------------]
92 //  after:   [----------------]                      [-------------------]
93 //
94 // erase() returns the segments actually erased. Given the state above:
95 //  Claim:          [---------------------------------------]
96 //  Out:            [---------]                      [------]
97 //  After:   [-----]                                         [-----------]
98 //
99 // It is used to track (expanded) tokens not yet associated with an AST node.
100 // On traversing an AST node, its token range is erased from the unclaimed set.
101 // The tokens actually removed are associated with that node, and hit-tested
102 // against the selection to determine whether the node is selected.
103 template <typename T> class IntervalSet {
104 public:
IntervalSet(llvm::ArrayRef<T> Range)105   IntervalSet(llvm::ArrayRef<T> Range) { UnclaimedRanges.insert(Range); }
106 
107   // Removes the elements of Claim from the set, modifying or removing ranges
108   // that overlap it.
109   // Returns the continuous subranges of Claim that were actually removed.
erase(llvm::ArrayRef<T> Claim)110   llvm::SmallVector<llvm::ArrayRef<T>> erase(llvm::ArrayRef<T> Claim) {
111     llvm::SmallVector<llvm::ArrayRef<T>> Out;
112     if (Claim.empty())
113       return Out;
114 
115     // General case:
116     // Claim:                   [-----------------]
117     // UnclaimedRanges: [-A-] [-B-] [-C-] [-D-] [-E-] [-F-] [-G-]
118     // Overlap:               ^first                  ^second
119     // Ranges C and D are fully included. Ranges B and E must be trimmed.
120     auto Overlap = std::make_pair(
121         UnclaimedRanges.lower_bound({Claim.begin(), Claim.begin()}), // C
122         UnclaimedRanges.lower_bound({Claim.end(), Claim.end()}));    // F
123     // Rewind to cover B.
124     if (Overlap.first != UnclaimedRanges.begin()) {
125       --Overlap.first;
126       // ...unless B isn't selected at all.
127       if (Overlap.first->end() <= Claim.begin())
128         ++Overlap.first;
129     }
130     if (Overlap.first == Overlap.second)
131       return Out;
132 
133     // First, copy all overlapping ranges into the output.
134     auto OutFirst = Out.insert(Out.end(), Overlap.first, Overlap.second);
135     // If any of the overlapping ranges were sliced by the claim, split them:
136     //  - restrict the returned range to the claimed part
137     //  - save the unclaimed part so it can be reinserted
138     llvm::ArrayRef<T> RemainingHead, RemainingTail;
139     if (Claim.begin() > OutFirst->begin()) {
140       RemainingHead = {OutFirst->begin(), Claim.begin()};
141       *OutFirst = {Claim.begin(), OutFirst->end()};
142     }
143     if (Claim.end() < Out.back().end()) {
144       RemainingTail = {Claim.end(), Out.back().end()};
145       Out.back() = {Out.back().begin(), Claim.end()};
146     }
147 
148     // Erase all the overlapping ranges (invalidating all iterators).
149     UnclaimedRanges.erase(Overlap.first, Overlap.second);
150     // Reinsert ranges that were merely trimmed.
151     if (!RemainingHead.empty())
152       UnclaimedRanges.insert(RemainingHead);
153     if (!RemainingTail.empty())
154       UnclaimedRanges.insert(RemainingTail);
155 
156     return Out;
157   }
158 
159 private:
160   using TokenRange = llvm::ArrayRef<T>;
161   struct RangeLess {
operator ()clang::clangd::__anon2b7dfc710111::IntervalSet::RangeLess162     bool operator()(llvm::ArrayRef<T> L, llvm::ArrayRef<T> R) const {
163       return L.begin() < R.begin();
164     }
165   };
166 
167   // Disjoint sorted unclaimed ranges of expanded tokens.
168   std::set<llvm::ArrayRef<T>, RangeLess> UnclaimedRanges;
169 };
170 
171 // Sentinel value for the selectedness of a node where we've seen no tokens yet.
172 // This resolves to Unselected if no tokens are ever seen.
173 // But Unselected + Complete -> Partial, while NoTokens + Complete --> Complete.
174 // This value is never exposed publicly.
175 constexpr SelectionTree::Selection NoTokens =
176     static_cast<SelectionTree::Selection>(
177         static_cast<unsigned char>(SelectionTree::Complete + 1));
178 
179 // Nodes start with NoTokens, and then use this function to aggregate the
180 // selectedness as more tokens are found.
update(SelectionTree::Selection & Result,SelectionTree::Selection New)181 void update(SelectionTree::Selection &Result, SelectionTree::Selection New) {
182   if (New == NoTokens)
183     return;
184   if (Result == NoTokens)
185     Result = New;
186   else if (Result != New)
187     // Can only be completely selected (or unselected) if all tokens are.
188     Result = SelectionTree::Partial;
189 }
190 
191 // As well as comments, don't count semicolons as real tokens.
192 // They're not properly claimed as expr-statement is missing from the AST.
shouldIgnore(const syntax::Token & Tok)193 bool shouldIgnore(const syntax::Token &Tok) {
194   switch (Tok.kind()) {
195     // Even "attached" comments are not considered part of a node's range.
196     case tok::comment:
197     // The AST doesn't directly store locations for terminating semicolons.
198     case tok::semi:
199     // We don't have locations for cvr-qualifiers: see QualifiedTypeLoc.
200     case tok::kw_const:
201     case tok::kw_volatile:
202     case tok::kw_restrict:
203       return true;
204     default:
205       return false;
206   }
207 }
208 
209 // Determine whether 'Target' is the first expansion of the macro
210 // argument whose top-level spelling location is 'SpellingLoc'.
isFirstExpansion(FileID Target,SourceLocation SpellingLoc,const SourceManager & SM)211 bool isFirstExpansion(FileID Target, SourceLocation SpellingLoc,
212                       const SourceManager &SM) {
213   SourceLocation Prev = SpellingLoc;
214   while (true) {
215     // If the arg is expanded multiple times, getMacroArgExpandedLocation()
216     // returns the first expansion.
217     SourceLocation Next = SM.getMacroArgExpandedLocation(Prev);
218     // So if we reach the target, target is the first-expansion of the
219     // first-expansion ...
220     if (SM.getFileID(Next) == Target)
221       return true;
222 
223     // Otherwise, if the FileID stops changing, we've reached the innermost
224     // macro expansion, and Target was on a different branch.
225     if (SM.getFileID(Next) == SM.getFileID(Prev))
226       return false;
227 
228     Prev = Next;
229   }
230   return false;
231 }
232 
233 // SelectionTester can determine whether a range of tokens from the PP-expanded
234 // stream (corresponding to an AST node) is considered selected.
235 //
236 // When the tokens result from macro expansions, the appropriate tokens in the
237 // main file are examined (macro invocation or args). Similarly for #includes.
238 // However, only the first expansion of a given spelled token is considered
239 // selected.
240 //
241 // It tests each token in the range (not just the endpoints) as contiguous
242 // expanded tokens may not have contiguous spellings (with macros).
243 //
244 // Non-token text, and tokens not modeled in the AST (comments, semicolons)
245 // are ignored when determining selectedness.
246 class SelectionTester {
247 public:
248   // The selection is offsets [SelBegin, SelEnd) in SelFile.
SelectionTester(const syntax::TokenBuffer & Buf,FileID SelFile,unsigned SelBegin,unsigned SelEnd,const SourceManager & SM)249   SelectionTester(const syntax::TokenBuffer &Buf, FileID SelFile,
250                   unsigned SelBegin, unsigned SelEnd, const SourceManager &SM)
251       : SelFile(SelFile), SelFileBounds(SM.getLocForStartOfFile(SelFile),
252                                         SM.getLocForEndOfFile(SelFile)),
253         SM(SM) {
254     // Find all tokens (partially) selected in the file.
255     auto AllSpelledTokens = Buf.spelledTokens(SelFile);
256     const syntax::Token *SelFirst =
257         llvm::partition_point(AllSpelledTokens, [&](const syntax::Token &Tok) {
258           return SM.getFileOffset(Tok.endLocation()) <= SelBegin;
259         });
260     const syntax::Token *SelLimit = std::partition_point(
261         SelFirst, AllSpelledTokens.end(), [&](const syntax::Token &Tok) {
262           return SM.getFileOffset(Tok.location()) < SelEnd;
263         });
264     auto Sel = llvm::makeArrayRef(SelFirst, SelLimit);
265     // Find which of these are preprocessed to nothing and should be ignored.
266     llvm::BitVector PPIgnored(Sel.size(), false);
267     for (const syntax::TokenBuffer::Expansion &X :
268          Buf.expansionsOverlapping(Sel)) {
269       if (X.Expanded.empty()) {
270         for (const syntax::Token &Tok : X.Spelled) {
271           if (&Tok >= SelFirst && &Tok < SelLimit)
272             PPIgnored[&Tok - SelFirst] = true;
273         }
274       }
275     }
276     // Precompute selectedness and offset for selected spelled tokens.
277     for (unsigned I = 0; I < Sel.size(); ++I) {
278       if (shouldIgnore(Sel[I]) || PPIgnored[I])
279         continue;
280       SelectedSpelled.emplace_back();
281       Tok &S = SelectedSpelled.back();
282       S.Offset = SM.getFileOffset(Sel[I].location());
283       if (S.Offset >= SelBegin && S.Offset + Sel[I].length() <= SelEnd)
284         S.Selected = SelectionTree::Complete;
285       else
286         S.Selected = SelectionTree::Partial;
287     }
288     MaybeSelectedExpanded = computeMaybeSelectedExpandedTokens(Buf);
289   }
290 
291   // Test whether a consecutive range of tokens is selected.
292   // The tokens are taken from the expanded token stream.
293   SelectionTree::Selection
test(llvm::ArrayRef<syntax::Token> ExpandedTokens) const294   test(llvm::ArrayRef<syntax::Token> ExpandedTokens) const {
295     if (ExpandedTokens.empty())
296       return NoTokens;
297     if (SelectedSpelled.empty())
298       return SelectionTree::Unselected;
299     // Cheap (pointer) check whether any of the tokens could touch selection.
300     // In most cases, the node's overall source range touches ExpandedTokens,
301     // or we would have failed mayHit(). However now we're only considering
302     // the *unclaimed* spans of expanded tokens.
303     // This is a significant performance improvement when a lot of nodes
304     // surround the selection, including when generated by macros.
305     if (MaybeSelectedExpanded.empty() ||
306         &ExpandedTokens.front() > &MaybeSelectedExpanded.back() ||
307         &ExpandedTokens.back() < &MaybeSelectedExpanded.front()) {
308       return SelectionTree::Unselected;
309     }
310 
311     // The eof token is used as a sentinel.
312     // In general, source range from an AST node should not claim the eof token,
313     // but it could occur for unmatched-bracket cases.
314     // FIXME: fix it in TokenBuffer, expandedTokens(SourceRange) should not
315     // return the eof token.
316     if (ExpandedTokens.back().kind() == tok::eof)
317       ExpandedTokens = ExpandedTokens.drop_back();
318 
319     SelectionTree::Selection Result = NoTokens;
320     while (!ExpandedTokens.empty()) {
321       // Take consecutive tokens from the same context together for efficiency.
322       SourceLocation Start = ExpandedTokens.front().location();
323       FileID FID = SM.getFileID(Start);
324       // Comparing SourceLocations against bounds is cheaper than getFileID().
325       SourceLocation Limit = SM.getComposedLoc(FID, SM.getFileIDSize(FID));
326       auto Batch = ExpandedTokens.take_while([&](const syntax::Token &T) {
327         return T.location() >= Start && T.location() < Limit;
328       });
329       assert(!Batch.empty());
330       ExpandedTokens = ExpandedTokens.drop_front(Batch.size());
331 
332       update(Result, testChunk(FID, Batch));
333     }
334     return Result;
335   }
336 
337   // Cheap check whether any of the tokens in R might be selected.
338   // If it returns false, test() will return NoTokens or Unselected.
339   // If it returns true, test() may return any value.
mayHit(SourceRange R) const340   bool mayHit(SourceRange R) const {
341     if (SelectedSpelled.empty() || MaybeSelectedExpanded.empty())
342       return false;
343     // If the node starts after the selection ends, it is not selected.
344     // Tokens a macro location might claim are >= its expansion start.
345     // So if the expansion start > last selected token, we can prune it.
346     // (This is particularly helpful for GTest's TEST macro).
347     if (auto B = offsetInSelFile(getExpansionStart(R.getBegin())))
348       if (*B > SelectedSpelled.back().Offset)
349         return false;
350     // If the node ends before the selection begins, it is not selected.
351     SourceLocation EndLoc = R.getEnd();
352     while (EndLoc.isMacroID())
353       EndLoc = SM.getImmediateExpansionRange(EndLoc).getEnd();
354     // In the rare case that the expansion range is a char range, EndLoc is
355     // ~one token too far to the right. We may fail to prune, that's OK.
356     if (auto E = offsetInSelFile(EndLoc))
357       if (*E < SelectedSpelled.front().Offset)
358         return false;
359     return true;
360   }
361 
362 private:
363   // Plausible expanded tokens that might be affected by the selection.
364   // This is an overestimate, it may contain tokens that are not selected.
365   // The point is to allow cheap pruning in test()
366   llvm::ArrayRef<syntax::Token>
computeMaybeSelectedExpandedTokens(const syntax::TokenBuffer & Toks)367   computeMaybeSelectedExpandedTokens(const syntax::TokenBuffer &Toks) {
368     if (SelectedSpelled.empty())
369       return {};
370 
371     auto LastAffectedToken = [&](SourceLocation Loc) {
372       auto Offset = offsetInSelFile(Loc);
373       while (Loc.isValid() && !Offset) {
374         Loc = Loc.isMacroID() ? SM.getImmediateExpansionRange(Loc).getEnd()
375                               : SM.getIncludeLoc(SM.getFileID(Loc));
376         Offset = offsetInSelFile(Loc);
377       }
378       return Offset;
379     };
380     auto FirstAffectedToken = [&](SourceLocation Loc) {
381       auto Offset = offsetInSelFile(Loc);
382       while (Loc.isValid() && !Offset) {
383         Loc = Loc.isMacroID() ? SM.getImmediateExpansionRange(Loc).getBegin()
384                               : SM.getIncludeLoc(SM.getFileID(Loc));
385         Offset = offsetInSelFile(Loc);
386       }
387       return Offset;
388     };
389 
390     const syntax::Token *Start = llvm::partition_point(
391         Toks.expandedTokens(),
392         [&, First = SelectedSpelled.front().Offset](const syntax::Token &Tok) {
393           if (Tok.kind() == tok::eof)
394             return false;
395           // Implausible if upperbound(Tok) < First.
396           if (auto Offset = LastAffectedToken(Tok.location()))
397             return *Offset < First;
398           // A prefix of the expanded tokens may be from an an implicit
399           // inclusion (e.g. preamble patch, or command-line -include).
400           return true;
401         });
402 
403     bool EndInvalid = false;
404     const syntax::Token *End = std::partition_point(
405         Start, Toks.expandedTokens().end(),
406         [&, Last = SelectedSpelled.back().Offset](const syntax::Token &Tok) {
407           if (Tok.kind() == tok::eof)
408             return false;
409           // Plausible if lowerbound(Tok) <= Last.
410           if (auto Offset = FirstAffectedToken(Tok.location()))
411             return *Offset <= Last;
412           // Shouldn't happen: once we've seen tokens traceable to the main
413           // file, there shouldn't be any more implicit inclusions.
414           assert(false && "Expanded token could not be resolved to main file!");
415           EndInvalid = true;
416           return true; // conservatively assume this token can overlap
417         });
418     if (EndInvalid)
419       End = Toks.expandedTokens().end();
420 
421     return llvm::makeArrayRef(Start, End);
422   }
423 
424   // Hit-test a consecutive range of tokens from a single file ID.
425   SelectionTree::Selection
testChunk(FileID FID,llvm::ArrayRef<syntax::Token> Batch) const426   testChunk(FileID FID, llvm::ArrayRef<syntax::Token> Batch) const {
427     assert(!Batch.empty());
428     SourceLocation StartLoc = Batch.front().location();
429     // There are several possible categories of FileID depending on how the
430     // preprocessor was used to generate these tokens:
431     //   main file, #included file, macro args, macro bodies.
432     // We need to identify the main-file tokens that represent Batch, and
433     // determine whether we want to exclusively claim them. Regular tokens
434     // represent one AST construct, but a macro invocation can represent many.
435 
436     // Handle tokens written directly in the main file.
437     if (FID == SelFile) {
438       return testTokenRange(*offsetInSelFile(Batch.front().location()),
439                             *offsetInSelFile(Batch.back().location()));
440     }
441 
442     // Handle tokens in another file #included into the main file.
443     // Check if the #include is selected, but don't claim it exclusively.
444     if (StartLoc.isFileID()) {
445       for (SourceLocation Loc = Batch.front().location(); Loc.isValid();
446            Loc = SM.getIncludeLoc(SM.getFileID(Loc))) {
447         if (auto Offset = offsetInSelFile(Loc))
448           // FIXME: use whole #include directive, not just the filename string.
449           return testToken(*Offset);
450       }
451       return NoTokens;
452     }
453 
454     assert(StartLoc.isMacroID());
455     // Handle tokens that were passed as a macro argument.
456     SourceLocation ArgStart = SM.getTopMacroCallerLoc(StartLoc);
457     if (auto ArgOffset = offsetInSelFile(ArgStart)) {
458       if (isFirstExpansion(FID, ArgStart, SM)) {
459         SourceLocation ArgEnd =
460             SM.getTopMacroCallerLoc(Batch.back().location());
461         return testTokenRange(*ArgOffset, *offsetInSelFile(ArgEnd));
462       } else { // NOLINT(llvm-else-after-return)
463         /* fall through and treat as part of the macro body */
464       }
465     }
466 
467     // Handle tokens produced by non-argument macro expansion.
468     // Check if the macro name is selected, don't claim it exclusively.
469     if (auto ExpansionOffset = offsetInSelFile(getExpansionStart(StartLoc)))
470       // FIXME: also check ( and ) for function-like macros?
471       return testToken(*ExpansionOffset);
472     return NoTokens;
473   }
474 
475   // Is the closed token range [Begin, End] selected?
testTokenRange(unsigned Begin,unsigned End) const476   SelectionTree::Selection testTokenRange(unsigned Begin, unsigned End) const {
477     assert(Begin <= End);
478     // Outside the selection entirely?
479     if (End < SelectedSpelled.front().Offset ||
480         Begin > SelectedSpelled.back().Offset)
481       return SelectionTree::Unselected;
482 
483     // Compute range of tokens.
484     auto B = llvm::partition_point(
485         SelectedSpelled, [&](const Tok &T) { return T.Offset < Begin; });
486     auto E = std::partition_point(B, SelectedSpelled.end(), [&](const Tok &T) {
487       return T.Offset <= End;
488     });
489 
490     // Aggregate selectedness of tokens in range.
491     bool ExtendsOutsideSelection = Begin < SelectedSpelled.front().Offset ||
492                                    End > SelectedSpelled.back().Offset;
493     SelectionTree::Selection Result =
494         ExtendsOutsideSelection ? SelectionTree::Unselected : NoTokens;
495     for (auto It = B; It != E; ++It)
496       update(Result, It->Selected);
497     return Result;
498   }
499 
500   // Is the token at `Offset` selected?
testToken(unsigned Offset) const501   SelectionTree::Selection testToken(unsigned Offset) const {
502     // Outside the selection entirely?
503     if (Offset < SelectedSpelled.front().Offset ||
504         Offset > SelectedSpelled.back().Offset)
505       return SelectionTree::Unselected;
506     // Find the token, if it exists.
507     auto It = llvm::partition_point(
508         SelectedSpelled, [&](const Tok &T) { return T.Offset < Offset; });
509     if (It != SelectedSpelled.end() && It->Offset == Offset)
510       return It->Selected;
511     return NoTokens;
512   }
513 
514   // Decomposes Loc and returns the offset if the file ID is SelFile.
offsetInSelFile(SourceLocation Loc) const515   llvm::Optional<unsigned> offsetInSelFile(SourceLocation Loc) const {
516     // Decoding Loc with SM.getDecomposedLoc is relatively expensive.
517     // But SourceLocations for a file are numerically contiguous, so we
518     // can use cheap integer operations instead.
519     if (Loc < SelFileBounds.getBegin() || Loc >= SelFileBounds.getEnd())
520       return llvm::None;
521     // FIXME: subtracting getRawEncoding() is dubious, move this logic into SM.
522     return Loc.getRawEncoding() - SelFileBounds.getBegin().getRawEncoding();
523   }
524 
getExpansionStart(SourceLocation Loc) const525   SourceLocation getExpansionStart(SourceLocation Loc) const {
526     while (Loc.isMacroID())
527       Loc = SM.getImmediateExpansionRange(Loc).getBegin();
528     return Loc;
529   }
530 
531   struct Tok {
532     unsigned Offset;
533     SelectionTree::Selection Selected;
534   };
535   std::vector<Tok> SelectedSpelled;
536   llvm::ArrayRef<syntax::Token> MaybeSelectedExpanded;
537   FileID SelFile;
538   SourceRange SelFileBounds;
539   const SourceManager &SM;
540 };
541 
542 // Show the type of a node for debugging.
printNodeKind(llvm::raw_ostream & OS,const DynTypedNode & N)543 void printNodeKind(llvm::raw_ostream &OS, const DynTypedNode &N) {
544   if (const TypeLoc *TL = N.get<TypeLoc>()) {
545     // TypeLoc is a hierarchy, but has only a single ASTNodeKind.
546     // Synthesize the name from the Type subclass (except for QualifiedTypeLoc).
547     if (TL->getTypeLocClass() == TypeLoc::Qualified)
548       OS << "QualifiedTypeLoc";
549     else
550       OS << TL->getType()->getTypeClassName() << "TypeLoc";
551   } else {
552     OS << N.getNodeKind().asStringRef();
553   }
554 }
555 
556 #ifndef NDEBUG
printNodeToString(const DynTypedNode & N,const PrintingPolicy & PP)557 std::string printNodeToString(const DynTypedNode &N, const PrintingPolicy &PP) {
558   std::string S;
559   llvm::raw_string_ostream OS(S);
560   printNodeKind(OS, N);
561   return std::move(OS.str());
562 }
563 #endif
564 
isImplicit(const Stmt * S)565 bool isImplicit(const Stmt *S) {
566   // Some Stmts are implicit and shouldn't be traversed, but there's no
567   // "implicit" attribute on Stmt/Expr.
568   // Unwrap implicit casts first if present (other nodes too?).
569   if (auto *ICE = llvm::dyn_cast<ImplicitCastExpr>(S))
570     S = ICE->getSubExprAsWritten();
571   // Implicit this in a MemberExpr is not filtered out by RecursiveASTVisitor.
572   // It would be nice if RAV handled this (!shouldTraverseImplicitCode()).
573   if (auto *CTI = llvm::dyn_cast<CXXThisExpr>(S))
574     if (CTI->isImplicit())
575       return true;
576   // Make sure implicit access of anonymous structs don't end up owning tokens.
577   if (auto *ME = llvm::dyn_cast<MemberExpr>(S)) {
578     if (auto *FD = llvm::dyn_cast<FieldDecl>(ME->getMemberDecl()))
579       if (FD->isAnonymousStructOrUnion())
580         // If Base is an implicit CXXThis, then the whole MemberExpr has no
581         // tokens. If it's a normal e.g. DeclRef, we treat the MemberExpr like
582         // an implicit cast.
583         return isImplicit(ME->getBase());
584   }
585   // Refs to operator() and [] are (almost?) always implicit as part of calls.
586   if (auto *DRE = llvm::dyn_cast<DeclRefExpr>(S)) {
587     if (auto *FD = llvm::dyn_cast<FunctionDecl>(DRE->getDecl())) {
588       switch (FD->getOverloadedOperator()) {
589       case OO_Call:
590       case OO_Subscript:
591         return true;
592       default:
593         break;
594       }
595     }
596   }
597   return false;
598 }
599 
600 // We find the selection by visiting written nodes in the AST, looking for nodes
601 // that intersect with the selected character range.
602 //
603 // While traversing, we maintain a parent stack. As nodes pop off the stack,
604 // we decide whether to keep them or not. To be kept, they must either be
605 // selected or contain some nodes that are.
606 //
607 // For simple cases (not inside macros) we prune subtrees that don't intersect.
608 class SelectionVisitor : public RecursiveASTVisitor<SelectionVisitor> {
609 public:
610   // Runs the visitor to gather selected nodes and their ancestors.
611   // If there is any selection, the root (TUDecl) is the first node.
collect(ASTContext & AST,const syntax::TokenBuffer & Tokens,const PrintingPolicy & PP,unsigned Begin,unsigned End,FileID File)612   static std::deque<Node> collect(ASTContext &AST,
613                                   const syntax::TokenBuffer &Tokens,
614                                   const PrintingPolicy &PP, unsigned Begin,
615                                   unsigned End, FileID File) {
616     SelectionVisitor V(AST, Tokens, PP, Begin, End, File);
617     V.TraverseAST(AST);
618     assert(V.Stack.size() == 1 && "Unpaired push/pop?");
619     assert(V.Stack.top() == &V.Nodes.front());
620     return std::move(V.Nodes);
621   }
622 
623   // We traverse all "well-behaved" nodes the same way:
624   //  - push the node onto the stack
625   //  - traverse its children recursively
626   //  - pop it from the stack
627   //  - hit testing: is intersection(node, selection) - union(children) empty?
628   //  - attach it to the tree if it or any children hit the selection
629   //
630   // Two categories of nodes are not "well-behaved":
631   //  - those without source range information, we don't record those
632   //  - those that can't be stored in DynTypedNode.
TraverseDecl(Decl * X)633   bool TraverseDecl(Decl *X) {
634     if (llvm::isa_and_nonnull<TranslationUnitDecl>(X))
635       return Base::TraverseDecl(X); // Already pushed by constructor.
636     // Base::TraverseDecl will suppress children, but not this node itself.
637     if (X && X->isImplicit())
638       return true;
639     return traverseNode(X, [&] { return Base::TraverseDecl(X); });
640   }
TraverseTypeLoc(TypeLoc X)641   bool TraverseTypeLoc(TypeLoc X) {
642     return traverseNode(&X, [&] { return Base::TraverseTypeLoc(X); });
643   }
TraverseTemplateArgumentLoc(const TemplateArgumentLoc & X)644   bool TraverseTemplateArgumentLoc(const TemplateArgumentLoc &X) {
645     return traverseNode(&X,
646                         [&] { return Base::TraverseTemplateArgumentLoc(X); });
647   }
TraverseNestedNameSpecifierLoc(NestedNameSpecifierLoc X)648   bool TraverseNestedNameSpecifierLoc(NestedNameSpecifierLoc X) {
649     return traverseNode(
650         &X, [&] { return Base::TraverseNestedNameSpecifierLoc(X); });
651   }
TraverseConstructorInitializer(CXXCtorInitializer * X)652   bool TraverseConstructorInitializer(CXXCtorInitializer *X) {
653     return traverseNode(
654         X, [&] { return Base::TraverseConstructorInitializer(X); });
655   }
TraverseCXXBaseSpecifier(const CXXBaseSpecifier & X)656   bool TraverseCXXBaseSpecifier(const CXXBaseSpecifier &X) {
657     return traverseNode(&X, [&] { return Base::TraverseCXXBaseSpecifier(X); });
658   }
TraverseAttr(Attr * X)659   bool TraverseAttr(Attr *X) {
660     return traverseNode(X, [&] { return Base::TraverseAttr(X); });
661   }
662   // Stmt is the same, but this form allows the data recursion optimization.
dataTraverseStmtPre(Stmt * X)663   bool dataTraverseStmtPre(Stmt *X) {
664     if (!X || isImplicit(X))
665       return false;
666     auto N = DynTypedNode::create(*X);
667     if (canSafelySkipNode(N))
668       return false;
669     push(std::move(N));
670     if (shouldSkipChildren(X)) {
671       pop();
672       return false;
673     }
674     return true;
675   }
dataTraverseStmtPost(Stmt * X)676   bool dataTraverseStmtPost(Stmt *X) {
677     pop();
678     return true;
679   }
680   // QualifiedTypeLoc is handled strangely in RecursiveASTVisitor: the derived
681   // TraverseTypeLoc is not called for the inner UnqualTypeLoc.
682   // This means we'd never see 'int' in 'const int'! Work around that here.
683   // (The reason for the behavior is to avoid traversing the nested Type twice,
684   // but we ignore TraverseType anyway).
TraverseQualifiedTypeLoc(QualifiedTypeLoc QX)685   bool TraverseQualifiedTypeLoc(QualifiedTypeLoc QX) {
686     return traverseNode<TypeLoc>(
687         &QX, [&] { return TraverseTypeLoc(QX.getUnqualifiedLoc()); });
688   }
TraverseObjCProtocolLoc(ObjCProtocolLoc PL)689   bool TraverseObjCProtocolLoc(ObjCProtocolLoc PL) {
690     return traverseNode(&PL, [&] { return Base::TraverseObjCProtocolLoc(PL); });
691   }
692   // Uninteresting parts of the AST that don't have locations within them.
TraverseNestedNameSpecifier(NestedNameSpecifier *)693   bool TraverseNestedNameSpecifier(NestedNameSpecifier *) { return true; }
TraverseType(QualType)694   bool TraverseType(QualType) { return true; }
695 
696   // The DeclStmt for the loop variable claims to cover the whole range
697   // inside the parens, this causes the range-init expression to not be hit.
698   // Traverse the loop VarDecl instead, which has the right source range.
TraverseCXXForRangeStmt(CXXForRangeStmt * S)699   bool TraverseCXXForRangeStmt(CXXForRangeStmt *S) {
700     return traverseNode(S, [&] {
701       return TraverseStmt(S->getInit()) && TraverseDecl(S->getLoopVariable()) &&
702              TraverseStmt(S->getRangeInit()) && TraverseStmt(S->getBody());
703     });
704   }
705   // OpaqueValueExpr blocks traversal, we must explicitly traverse it.
TraverseOpaqueValueExpr(OpaqueValueExpr * E)706   bool TraverseOpaqueValueExpr(OpaqueValueExpr *E) {
707     return traverseNode(E, [&] { return TraverseStmt(E->getSourceExpr()); });
708   }
709   // We only want to traverse the *syntactic form* to understand the selection.
TraversePseudoObjectExpr(PseudoObjectExpr * E)710   bool TraversePseudoObjectExpr(PseudoObjectExpr *E) {
711     return traverseNode(E, [&] { return TraverseStmt(E->getSyntacticForm()); });
712   }
TraverseTypeConstraint(const TypeConstraint * C)713   bool TraverseTypeConstraint(const TypeConstraint *C) {
714     if (auto *E = C->getImmediatelyDeclaredConstraint()) {
715       // Technically this expression is 'implicit' and not traversed by the RAV.
716       // However, the range is correct, so we visit expression to avoid adding
717       // an extra kind to 'DynTypeNode' that hold 'TypeConstraint'.
718       return TraverseStmt(E);
719     }
720     return Base::TraverseTypeConstraint(C);
721   }
722 
723 private:
724   using Base = RecursiveASTVisitor<SelectionVisitor>;
725 
SelectionVisitor(ASTContext & AST,const syntax::TokenBuffer & Tokens,const PrintingPolicy & PP,unsigned SelBegin,unsigned SelEnd,FileID SelFile)726   SelectionVisitor(ASTContext &AST, const syntax::TokenBuffer &Tokens,
727                    const PrintingPolicy &PP, unsigned SelBegin, unsigned SelEnd,
728                    FileID SelFile)
729       : SM(AST.getSourceManager()), LangOpts(AST.getLangOpts()),
730 #ifndef NDEBUG
731         PrintPolicy(PP),
732 #endif
733         TokenBuf(Tokens), SelChecker(Tokens, SelFile, SelBegin, SelEnd, SM),
734         UnclaimedExpandedTokens(Tokens.expandedTokens()) {
735     // Ensure we have a node for the TU decl, regardless of traversal scope.
736     Nodes.emplace_back();
737     Nodes.back().ASTNode = DynTypedNode::create(*AST.getTranslationUnitDecl());
738     Nodes.back().Parent = nullptr;
739     Nodes.back().Selected = SelectionTree::Unselected;
740     Stack.push(&Nodes.back());
741   }
742 
743   // Generic case of TraverseFoo. Func should be the call to Base::TraverseFoo.
744   // Node is always a pointer so the generic code can handle any null checks.
745   template <typename T, typename Func>
traverseNode(T * Node,const Func & Body)746   bool traverseNode(T *Node, const Func &Body) {
747     if (Node == nullptr)
748       return true;
749     auto N = DynTypedNode::create(*Node);
750     if (canSafelySkipNode(N))
751       return true;
752     push(DynTypedNode::create(*Node));
753     bool Ret = Body();
754     pop();
755     return Ret;
756   }
757 
758   // HIT TESTING
759   //
760   // We do rough hit testing on the way down the tree to avoid traversing
761   // subtrees that don't touch the selection (canSafelySkipNode), but
762   // fine-grained hit-testing is mostly done on the way back up (in pop()).
763   // This means children get to claim parts of the selection first, and parents
764   // are only selected if they own tokens that no child owned.
765   //
766   // Nodes *usually* nest nicely: a child's getSourceRange() lies within the
767   // parent's, and a node (transitively) owns all tokens in its range.
768   //
769   // Exception 1: when declarators nest, *inner* declarator is the *outer* type.
770   //              e.g. void foo[5](int) is an array of functions.
771   // To handle this case, declarators are careful to only claim the tokens they
772   // own, rather than claim a range and rely on claim ordering.
773   //
774   // Exception 2: siblings both claim the same node.
775   //              e.g. `int x, y;` produces two sibling VarDecls.
776   //                    ~~~~~ x
777   //                    ~~~~~~~~ y
778   // Here the first ("leftmost") sibling claims the tokens it wants, and the
779   // other sibling gets what's left. So selecting "int" only includes the left
780   // VarDecl in the selection tree.
781 
782   // An optimization for a common case: nodes outside macro expansions that
783   // don't intersect the selection may be recursively skipped.
canSafelySkipNode(const DynTypedNode & N)784   bool canSafelySkipNode(const DynTypedNode &N) {
785     SourceRange S = getSourceRange(N);
786     if (auto *TL = N.get<TypeLoc>()) {
787       // FIXME: TypeLoc::getBeginLoc()/getEndLoc() are pretty fragile
788       // heuristics. We should consider only pruning critical TypeLoc nodes, to
789       // be more robust.
790 
791       // AttributedTypeLoc may point to the attribute's range, NOT the modified
792       // type's range.
793       if (auto AT = TL->getAs<AttributedTypeLoc>())
794         S = AT.getModifiedLoc().getSourceRange();
795     }
796     // SourceRange often doesn't manage to accurately cover attributes.
797     // Fortunately, attributes are rare.
798     if (llvm::any_of(getAttributes(N),
799                      [](const Attr *A) { return !A->isImplicit(); }))
800       return false;
801     if (!SelChecker.mayHit(S)) {
802       dlog("{2}skip: {0} {1}", printNodeToString(N, PrintPolicy),
803            S.printToString(SM), indent());
804       return true;
805     }
806     return false;
807   }
808 
809   // There are certain nodes we want to treat as leaves in the SelectionTree,
810   // although they do have children.
shouldSkipChildren(const Stmt * X) const811   bool shouldSkipChildren(const Stmt *X) const {
812     // UserDefinedLiteral (e.g. 12_i) has two children (12 and _i).
813     // Unfortunately TokenBuffer sees 12_i as one token and can't split it.
814     // So we treat UserDefinedLiteral as a leaf node, owning the token.
815     return llvm::isa<UserDefinedLiteral>(X);
816   }
817 
818   // Pushes a node onto the ancestor stack. Pairs with pop().
819   // Performs early hit detection for some nodes (on the earlySourceRange).
push(DynTypedNode Node)820   void push(DynTypedNode Node) {
821     SourceRange Early = earlySourceRange(Node);
822     dlog("{2}push: {0} {1}", printNodeToString(Node, PrintPolicy),
823          Node.getSourceRange().printToString(SM), indent());
824     Nodes.emplace_back();
825     Nodes.back().ASTNode = std::move(Node);
826     Nodes.back().Parent = Stack.top();
827     Nodes.back().Selected = NoTokens;
828     Stack.push(&Nodes.back());
829     claimRange(Early, Nodes.back().Selected);
830   }
831 
832   // Pops a node off the ancestor stack, and finalizes it. Pairs with push().
833   // Performs primary hit detection.
pop()834   void pop() {
835     Node &N = *Stack.top();
836     dlog("{1}pop: {0}", printNodeToString(N.ASTNode, PrintPolicy), indent(-1));
837     claimTokensFor(N.ASTNode, N.Selected);
838     if (N.Selected == NoTokens)
839       N.Selected = SelectionTree::Unselected;
840     if (N.Selected || !N.Children.empty()) {
841       // Attach to the tree.
842       N.Parent->Children.push_back(&N);
843     } else {
844       // Neither N any children are selected, it doesn't belong in the tree.
845       assert(&N == &Nodes.back());
846       Nodes.pop_back();
847     }
848     Stack.pop();
849   }
850 
851   // Returns the range of tokens that this node will claim directly, and
852   // is not available to the node's children.
853   // Usually empty, but sometimes children cover tokens but shouldn't own them.
earlySourceRange(const DynTypedNode & N)854   SourceRange earlySourceRange(const DynTypedNode &N) {
855     if (const Decl *D = N.get<Decl>()) {
856       // We want the name in the var-decl to be claimed by the decl itself and
857       // not by any children. Ususally, we don't need this, because source
858       // ranges of children are not overlapped with their parent's.
859       // An exception is lambda captured var decl, where AutoTypeLoc is
860       // overlapped with the name loc.
861       //    auto fun = [bar = foo]() { ... }
862       //                ~~~~~~~~~   VarDecl
863       //                ~~~         |- AutoTypeLoc
864       if (const auto *DD = llvm::dyn_cast<VarDecl>(D))
865         return DD->getLocation();
866     }
867 
868     return SourceRange();
869   }
870 
871   // Claim tokens for N, after processing its children.
872   // By default this claims all unclaimed tokens in getSourceRange().
873   // We override this if we want to claim fewer tokens (e.g. there are gaps).
claimTokensFor(const DynTypedNode & N,SelectionTree::Selection & Result)874   void claimTokensFor(const DynTypedNode &N, SelectionTree::Selection &Result) {
875     // CXXConstructExpr often shows implicit construction, like `string s;`.
876     // Don't associate any tokens with it unless there's some syntax like {}.
877     // This prevents it from claiming 's', its primary location.
878     if (const auto *CCE = N.get<CXXConstructExpr>()) {
879       claimRange(CCE->getParenOrBraceRange(), Result);
880       return;
881     }
882     // ExprWithCleanups is always implicit. It often wraps CXXConstructExpr.
883     // Prevent it claiming 's' in the case above.
884     if (N.get<ExprWithCleanups>())
885       return;
886 
887     // Declarators nest "inside out", with parent types inside child ones.
888     // Instead of claiming the whole range (clobbering parent tokens), carefully
889     // claim the tokens owned by this node and non-declarator children.
890     // (We could manipulate traversal order instead, but this is easier).
891     //
892     // Non-declarator types nest normally, and are handled like other nodes.
893     //
894     // Example:
895     //   Vec<R<int>(*[2])(A<char>)> is a Vec of arrays of pointers to functions,
896     //                              which accept A<char> and return R<int>.
897     // The TypeLoc hierarchy:
898     //   Vec<R<int>(*[2])(A<char>)> m;
899     //   Vec<#####################>      TemplateSpecialization Vec
900     //       --------[2]----------       `-Array
901     //       -------*-------------         `-Pointer
902     //       ------(----)---------           `-Paren
903     //       ------------(#######)             `-Function
904     //       R<###>                              |-TemplateSpecialization R
905     //         int                               | `-Builtin int
906     //                    A<####>                `-TemplateSpecialization A
907     //                      char                   `-Builtin char
908     //
909     // In each row
910     //   --- represents unclaimed parts of the SourceRange.
911     //   ### represents parts that children already claimed.
912     if (const auto *TL = N.get<TypeLoc>()) {
913       if (auto PTL = TL->getAs<ParenTypeLoc>()) {
914         claimRange(PTL.getLParenLoc(), Result);
915         claimRange(PTL.getRParenLoc(), Result);
916         return;
917       }
918       if (auto ATL = TL->getAs<ArrayTypeLoc>()) {
919         claimRange(ATL.getBracketsRange(), Result);
920         return;
921       }
922       if (auto PTL = TL->getAs<PointerTypeLoc>()) {
923         claimRange(PTL.getStarLoc(), Result);
924         return;
925       }
926       if (auto FTL = TL->getAs<FunctionTypeLoc>()) {
927         claimRange(SourceRange(FTL.getLParenLoc(), FTL.getEndLoc()), Result);
928         return;
929       }
930     }
931     claimRange(getSourceRange(N), Result);
932   }
933 
934   // Perform hit-testing of a complete Node against the selection.
935   // This runs for every node in the AST, and must be fast in common cases.
936   // This is usually called from pop(), so we can take children into account.
937   // The existing state of Result is relevant.
claimRange(SourceRange S,SelectionTree::Selection & Result)938   void claimRange(SourceRange S, SelectionTree::Selection &Result) {
939     for (const auto &ClaimedRange :
940          UnclaimedExpandedTokens.erase(TokenBuf.expandedTokens(S)))
941       update(Result, SelChecker.test(ClaimedRange));
942 
943     if (Result && Result != NoTokens)
944       dlog("{1}hit selection: {0}", S.printToString(SM), indent());
945   }
946 
indent(int Offset=0)947   std::string indent(int Offset = 0) {
948     // Cast for signed arithmetic.
949     int Amount = int(Stack.size()) + Offset;
950     assert(Amount >= 0);
951     return std::string(Amount, ' ');
952   }
953 
954   SourceManager &SM;
955   const LangOptions &LangOpts;
956 #ifndef NDEBUG
957   const PrintingPolicy &PrintPolicy;
958 #endif
959   const syntax::TokenBuffer &TokenBuf;
960   std::stack<Node *> Stack;
961   SelectionTester SelChecker;
962   IntervalSet<syntax::Token> UnclaimedExpandedTokens;
963   std::deque<Node> Nodes; // Stable pointers as we add more nodes.
964 };
965 
966 } // namespace
967 
abbreviatedString(DynTypedNode N,const PrintingPolicy & PP)968 llvm::SmallString<256> abbreviatedString(DynTypedNode N,
969                                          const PrintingPolicy &PP) {
970   llvm::SmallString<256> Result;
971   {
972     llvm::raw_svector_ostream OS(Result);
973     N.print(OS, PP);
974   }
975   auto Pos = Result.find('\n');
976   if (Pos != llvm::StringRef::npos) {
977     bool MoreText = !llvm::all_of(Result.str().drop_front(Pos), llvm::isSpace);
978     Result.resize(Pos);
979     if (MoreText)
980       Result.append(" …");
981   }
982   return Result;
983 }
984 
print(llvm::raw_ostream & OS,const SelectionTree::Node & N,int Indent) const985 void SelectionTree::print(llvm::raw_ostream &OS, const SelectionTree::Node &N,
986                           int Indent) const {
987   if (N.Selected)
988     OS.indent(Indent - 1) << (N.Selected == SelectionTree::Complete ? '*'
989                                                                     : '.');
990   else
991     OS.indent(Indent);
992   printNodeKind(OS, N.ASTNode);
993   OS << ' ' << abbreviatedString(N.ASTNode, PrintPolicy) << "\n";
994   for (const Node *Child : N.Children)
995     print(OS, *Child, Indent + 2);
996 }
997 
kind() const998 std::string SelectionTree::Node::kind() const {
999   std::string S;
1000   llvm::raw_string_ostream OS(S);
1001   printNodeKind(OS, ASTNode);
1002   return std::move(OS.str());
1003 }
1004 
1005 // Decide which selections emulate a "point" query in between characters.
1006 // If it's ambiguous (the neighboring characters are selectable tokens), returns
1007 // both possibilities in preference order.
1008 // Always returns at least one range - if no tokens touched, and empty range.
1009 static llvm::SmallVector<std::pair<unsigned, unsigned>, 2>
pointBounds(unsigned Offset,const syntax::TokenBuffer & Tokens)1010 pointBounds(unsigned Offset, const syntax::TokenBuffer &Tokens) {
1011   const auto &SM = Tokens.sourceManager();
1012   SourceLocation Loc = SM.getComposedLoc(SM.getMainFileID(), Offset);
1013   llvm::SmallVector<std::pair<unsigned, unsigned>, 2> Result;
1014   // Prefer right token over left.
1015   for (const syntax::Token &Tok :
1016        llvm::reverse(spelledTokensTouching(Loc, Tokens))) {
1017     if (shouldIgnore(Tok))
1018       continue;
1019     unsigned Offset = Tokens.sourceManager().getFileOffset(Tok.location());
1020     Result.emplace_back(Offset, Offset + Tok.length());
1021   }
1022   if (Result.empty())
1023     Result.emplace_back(Offset, Offset);
1024   return Result;
1025 }
1026 
createEach(ASTContext & AST,const syntax::TokenBuffer & Tokens,unsigned Begin,unsigned End,llvm::function_ref<bool (SelectionTree)> Func)1027 bool SelectionTree::createEach(ASTContext &AST,
1028                                const syntax::TokenBuffer &Tokens,
1029                                unsigned Begin, unsigned End,
1030                                llvm::function_ref<bool(SelectionTree)> Func) {
1031   if (Begin != End)
1032     return Func(SelectionTree(AST, Tokens, Begin, End));
1033   for (std::pair<unsigned, unsigned> Bounds : pointBounds(Begin, Tokens))
1034     if (Func(SelectionTree(AST, Tokens, Bounds.first, Bounds.second)))
1035       return true;
1036   return false;
1037 }
1038 
createRight(ASTContext & AST,const syntax::TokenBuffer & Tokens,unsigned int Begin,unsigned int End)1039 SelectionTree SelectionTree::createRight(ASTContext &AST,
1040                                          const syntax::TokenBuffer &Tokens,
1041                                          unsigned int Begin, unsigned int End) {
1042   llvm::Optional<SelectionTree> Result;
1043   createEach(AST, Tokens, Begin, End, [&](SelectionTree T) {
1044     Result = std::move(T);
1045     return true;
1046   });
1047   return std::move(*Result);
1048 }
1049 
SelectionTree(ASTContext & AST,const syntax::TokenBuffer & Tokens,unsigned Begin,unsigned End)1050 SelectionTree::SelectionTree(ASTContext &AST, const syntax::TokenBuffer &Tokens,
1051                              unsigned Begin, unsigned End)
1052     : PrintPolicy(AST.getLangOpts()) {
1053   // No fundamental reason the selection needs to be in the main file,
1054   // but that's all clangd has needed so far.
1055   const SourceManager &SM = AST.getSourceManager();
1056   FileID FID = SM.getMainFileID();
1057   PrintPolicy.TerseOutput = true;
1058   PrintPolicy.IncludeNewlines = false;
1059 
1060   dlog("Computing selection for {0}",
1061        SourceRange(SM.getComposedLoc(FID, Begin), SM.getComposedLoc(FID, End))
1062            .printToString(SM));
1063   Nodes = SelectionVisitor::collect(AST, Tokens, PrintPolicy, Begin, End, FID);
1064   Root = Nodes.empty() ? nullptr : &Nodes.front();
1065   recordMetrics(*this, AST.getLangOpts());
1066   dlog("Built selection tree\n{0}", *this);
1067 }
1068 
commonAncestor() const1069 const Node *SelectionTree::commonAncestor() const {
1070   const Node *Ancestor = Root;
1071   while (Ancestor->Children.size() == 1 && !Ancestor->Selected)
1072     Ancestor = Ancestor->Children.front();
1073   // Returning nullptr here is a bit unprincipled, but it makes the API safer:
1074   // the TranslationUnitDecl contains all of the preamble, so traversing it is a
1075   // performance cliff. Callers can check for null and use root() if they want.
1076   return Ancestor != Root ? Ancestor : nullptr;
1077 }
1078 
getDeclContext() const1079 const DeclContext &SelectionTree::Node::getDeclContext() const {
1080   for (const Node *CurrentNode = this; CurrentNode != nullptr;
1081        CurrentNode = CurrentNode->Parent) {
1082     if (const Decl *Current = CurrentNode->ASTNode.get<Decl>()) {
1083       if (CurrentNode != this)
1084         if (auto *DC = dyn_cast<DeclContext>(Current))
1085           return *DC;
1086       return *Current->getLexicalDeclContext();
1087     }
1088   }
1089   llvm_unreachable("A tree must always be rooted at TranslationUnitDecl.");
1090 }
1091 
ignoreImplicit() const1092 const SelectionTree::Node &SelectionTree::Node::ignoreImplicit() const {
1093   if (Children.size() == 1 &&
1094       getSourceRange(Children.front()->ASTNode) == getSourceRange(ASTNode))
1095     return Children.front()->ignoreImplicit();
1096   return *this;
1097 }
1098 
outerImplicit() const1099 const SelectionTree::Node &SelectionTree::Node::outerImplicit() const {
1100   if (Parent && getSourceRange(Parent->ASTNode) == getSourceRange(ASTNode))
1101     return Parent->outerImplicit();
1102   return *this;
1103 }
1104 
1105 } // namespace clangd
1106 } // namespace clang
1107