[clangd] Rethink how SelectionTree deals with macros and #includes.

[sammccall]

Summary

The exclusive-claim model is successful at resolving conflicts over tokens
between parent/child or siblings. However it's not the right model for selecting
AST nodes produced by a macro expansion, which can produce several independent
nodes that are equally associated with the macro invocation.
Additionally, any model that only uses the endpoints in a range can fail when
a macro invocation occurs inside the node.

To address this, we use the existing TokenBuffer in more depth. SourceRanges
are translated into an array of expanded tokens we can iterate over, and in
principle process token by token (in practice, batching related tokens helps).

For regular tokens (and macro-arg expansions) we claim the tokens as before,
but for tokens from macro bodies we merely check whether the macro name was
selected. Thus tokens in macro bodies are selected by selecting the macro name.

The aggregation of Selection is now more principled as we need to be able to
call claim()/peek() an arbitrary number of times.

One side-effect of iterating over the expanded tokens is that (usually) nothing
claims preprocessor tokens like macro names and directives.
Rather than fixing this I just left them unclaimed, and removed support for
determining the selectedness of TUDecl.
(That was originally implemented in 90a5bf92ff97b1, but doesn't seem to be very
important or worth the complexity any longer).

The expandedTokens(SourceLocation) helper could be added locally, but seems to
make sense on TokenBuffer.

[sammccall]

This fixes:

• https://github.com/clangd/clangd/issues/126
• https://github.com/clangd/clangd/issues/202

(I've amended the commit message)

[merge_guards_bot]

Build result: pass - 60204 tests passed, 0 failed and 732 were skipped.
Log files: console-log.txt, CMakeCache.txt

[hokein]

clang-tools-extra/clangd/Selection.cpp

49

I might be missing something, I don't understand why we set Partial if Result != New.

598

what's S for this case?

616

nit: maybe use [early-exist](https://llvm.org/docs/CodingStandards.html#use-early-exits-and-continue-to-simplify-code) here?

648

not sure what's the rational behavior, but I think for the following case, we just have the TUDecl in the selection tree, maybe use the whole macro range?

#define M() 123

int d = M(^); // now we only have the TUDecl in the selection tree.

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

139

could we have a testcase to cover the "tokens expanded from another #include file" code path?

clang/lib/Tooling/Syntax/Tokens.cpp

125

nit: for code readability, I'd use llvm::ArrayRef<syntax::Token>::iterator type here.

126

I think the parition_point requires the ExpandedTokens is partitioned, but I didn't found any documentation about this guarantee in the code, would be nice to have this in the comment (probably around the ExpandedTokens).

[sammccall]

Address review comments.

[merge_guards_bot]

Build result: fail - 60299 tests passed, 1 failed and 732 were skipped.

failed: Clangd Unit Tests._/ClangdTests/SelectionTest.PathologicalPreprocessor

Log files: console-log.txt, CMakeCache.txt

[sammccall]

No, this patch is busted, and the tests were too simple go catch it.

The issue is that with no exclusivity check, given {{{ MACRO }}} all the enclosing blocks get to count themselves selected.

We need an exclusivity check on the expanded token stream first. This will yield a sequence of slices of tokens not yet claimed. Then each slice gets partitioned by FileID, each partition gets mapped to spelled tokens and checked as in this patch.

I suspect the exclusivity/claiming at the spelled token level is no longer needed.

clang-tools-extra/clangd/Selection.cpp

49

Added a brief comment and an assertion.
Intuitive explanation: Consider Complete = black, Unselected = white, Partial = grey.
White + White -> white, Black + Black -> black, every other mixture is grey.

(But honestly I wrote up the state transition table first and then extracted the logic by staring at it)

598

S is the function parameter, the implication is that it's the range of the vardecl in the example. Expanded the comment a bit.

648

This is definitely better, but isn't trivial to do, and isn't a very important case.
Added a FIXME rather than clutter/delay this patch with it.

clang/lib/Tooling/Syntax/Tokens.cpp

125

I think that hurts readability on two counts:

• it obscures the actual type: a pointer is concrete and familiar, so it's easier to realize that e.g. > is well-defined here
• it makes it harder to understand return {Begin, End} which relies on the fact that the actual type here is Token*

[sammccall]

Rewrote patch with a better approach (claiming expanded tokens rather than spelled).
Added more tests, including one showing a problem with multiple arg expansion.

[merge_guards_bot]

Build result: pass - 60302 tests passed, 0 failed and 732 were skipped.

Log files: console-log.txt, CMakeCache.txt

[hokein]

Thanks, the patch looks good! please also update the patch description.

clang-tools-extra/clangd/Selection.cpp

63

nit: this comment seems not reflect to the code now.

126

nit: remove the redundant start.

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

468

I think cases like below are broken:

#define greater(x, y) x > y? x : y
#define abs(x) x > 0 ? x : -x

Selecting the first element for a macro arg seems good to me.

[thakis]

hover.test is failing on Mac: http://45.33.8.238/mac/3308/step_7.txt

Please take a look, and revert if it takes a while.

[sammccall]

Reverted as 905b002c139f039a32ab9bf1fad63d745d12423f

@thakis any information available about that bot configuration?
I'm not seeing any failures on "official" ones, I can't reproduce locally, and the failures don't make a lot of sense to me (lots of off-by-ones everywhere maybe?)

[thakis]

That particular bot does GN builds, but these test failures repro in the cmake build for me on both macs I tried. (Remember that the official mac bots are on greendragon, not on buildbot -- I'd guess it shows there too.)

[thakis]

Mac is happy now, but it fails to build on Windows: http://45.33.8.238/win/3327/step_4.txt

clang-tools-extra/clangd/Selection.h

SelectionTree(ASTContext &AST, const syntax::TokenBuffer &Tokens,
                 unsigned Offset);
   // Creates a selection tree for the given range in the main file.
   // The range includes bytes [Start, End).
   // If Start == End, uses the same heuristics as SelectionTree(AST, Start).
   SelectionTree(ASTContext &AST, const syntax::TokenBuffer &Tokens,
                 unsigned Start, unsigned End);
 
   // Describes to what extent an AST node is covered by the selection.
-  enum Selection {
+  enum Selection : unsigned char {
     // The AST node owns no characters covered by the selection.
     // Note that characters owned by children don't count:
     //   if (x == 0) scream();
     //       ^^^^^^
     // The IfStmt would be Unselected because all the selected characters are
     // associated with its children.
     // (Invisible nodes like ImplicitCastExpr are always unselected).
     Unselected,

clang-tools-extra/clangd/Selection.cpp

 #include <string>
 
 namespace clang {
 namespace clangd {
 namespace {
 using Node = SelectionTree::Node;
 using ast_type_traits::DynTypedNode;
 
-// Identifies which tokens are selected, and evaluates claims of source ranges
-// by AST nodes. Tokens may be claimed only once: first-come, first-served.
-class SelectedTokens {
+// An IntervalSet maintains a set of disjoint subranges of an array.
+//
+// Initially, it contains the entire array.
+//           [-----------------------------------------------------------]
+//
+// When a range is erased(), it will typically split the array in two.
+//  Claim:                     [--------------------]
+//  after:   [----------------]                      [-------------------]
+//
+// erase() returns the segments actually erased. Given the state above:
+//  Claim:          [---------------------------------------]
+//  Out:            [---------]                      [------]
+//  After:   [-----]                                         [-----------]
+//
+// It is used to track (expanded) tokens not yet associated with an AST node.
+// On traversing an AST node, its token range is erased from the unclaimed set.
+// The tokens actually removed are associated with that node, and hit-tested
+// against the selection to determine whether the node is selected.
+template <typename T>
+class IntervalSet {
 public:
-  SelectedTokens(llvm::ArrayRef<syntax::Token> Spelled, const SourceManager &SM,
-                 unsigned SelBegin, unsigned SelEnd)
-      : SelBegin(SelBegin), SelEnd(SelEnd) {
-    // Extract bounds and selected-ness for all tokens spelled in the file.
-    Tokens.reserve(Spelled.size());
-    for (const auto& Tok : Spelled) {
+  IntervalSet(llvm::ArrayRef<T> Range) : UnclaimedRanges(&rangeLess) {
+    UnclaimedRanges.insert(Range);
+  }
+
+  // Removes the elements of Claim from the set, modifying or removing ranges
+  // that overlap it.
+  // Returns the continuous subranges of Claim that were actually removed.
+  llvm::SmallVector<llvm::ArrayRef<T>, 4> erase(llvm::ArrayRef<T> Claim) {
+    llvm::SmallVector<llvm::ArrayRef<T>, 4> Out;
+    if (Claim.empty())
+      return Out;
+    // equal_range finds overlapping ranges, because of how we chose <.
+    auto Overlap = UnclaimedRanges.equal_range(Claim);
+    if (Overlap.first == Overlap.second)
+      return Out;
+
+    // General case:
+    // Claim:                   [-----------------]
+    // UnclaimedRanges: [-A-] [-B-] [-C-] [-D-] [-E-] [-F-] [-G-]
+    // Overlap:               ^first                  ^second
+    // Ranges C and D are fully included. Ranges B and E must be trimmed.
+
+    // First, copy all overlapping ranges into the output.
+    auto OutFirst = Out.insert(Out.end(), Overlap.first, Overlap.second);
+    // If any of the overlapping ranges were sliced by the claim, split them:
+    //  - restrict the returned range to the claimed part
+    //  - save the unclaimed part so it can be reinserted
+    llvm::ArrayRef<T> RemainingHead, RemainingTail;
+    if (Claim.begin() > OutFirst->begin()) {
+      RemainingHead = {OutFirst->begin(), Claim.begin()};
+      *OutFirst = {Claim.begin(), OutFirst->end()};
+    }
+    if (Claim.end() < Out.back().end()) {
+      RemainingTail = {Claim.end(), Out.back().end()};
+      Out.back() = {Out.back().begin(), Claim.end()};
+    }
+
+    // Erase all the overlapping ranges (invalidating all iterators).
+    UnclaimedRanges.erase(Overlap.first, Overlap.second);
+    // Reinsert ranges that were merely trimmed.
+    if (!RemainingHead.empty())
+      UnclaimedRanges.insert(RemainingHead);
+    if (!RemainingTail.empty())
+      UnclaimedRanges.insert(RemainingTail);
+
+    return Out;
+  }
+
+private:
+  using TokenRange = llvm::ArrayRef<T>;
+  // Given that the ranges we insert are disjoint, there are several ways to
+  // legally define range < range.
+  // We choose to define it so overlapping ranges compare equal.
+  static bool rangeLess(llvm::ArrayRef<T> L, llvm::ArrayRef<T> R) {
+    return L.end() <= R.begin();
+  }
+
+  // Disjoint sorted unclaimed ranges of expanded tokens.
+  std::set<llvm::ArrayRef<T>, decltype(&rangeLess)> UnclaimedRanges;
+};
+
+// Sentinel value for the selectedness of a node where we've seen no tokens yet.
+// This resolves to Unselected if no tokens are ever seen.
+// But Unselected + Complete -> Partial, while NoTokens + Complete --> Complete.
+// This value is never exposed publicly.
+constexpr SelectionTree::Selection NoTokens =
+    static_cast<SelectionTree::Selection>(
+        static_cast<unsigned char>(SelectionTree::Complete + 1));
+
+// Nodes start with NoTokens, and then use this function to aggregate the
+// selectedness as more tokens are found.
+void update(SelectionTree::Selection &Result, SelectionTree::Selection New) {
+  if (New == NoTokens)
+    return;
+  if (Result == NoTokens)
+    Result = New;
+  else if (Result != New)
+    // Can only be completely selected (or unselected) if all tokens are.
+    Result = SelectionTree::Partial;
+}
+
+
+// SelectionTester can determine whether a range of tokens from the PP-expanded
+// stream (corresponding to an AST node) is considered selected.
+//
+// When the tokens result from macro expansions, the appropriate tokens in the
+// main file are examined (macro invocation or args). Similarly for #includes.
+//
+// It tests each token in the range (not just the endpoints) as contiguous
+// expanded tokens may not have contiguous spellings (with macros).
+//
+// Non-token text, and tokens not modeled in the AST (comments, semicolons)
+// are ignored when determining selectedness.
+class SelectionTester {
+public:
+  // The selection is offsets [SelBegin, SelEnd) in SelFile.
+  SelectionTester(const syntax::TokenBuffer &Buf, FileID SelFile,
+                  unsigned SelBegin, unsigned SelEnd, const SourceManager &SM)
+      : SelFile(SelFile), SM(SM) {
+    // Find all tokens (partially) selected in the file.
+    auto AllSpelledTokens = Buf.spelledTokens(SelFile);
+    const syntax::Token *SelFirst =
+        llvm::partition_point(AllSpelledTokens, [&](const syntax::Token &Tok) {
+          return SM.getFileOffset(Tok.endLocation()) <= SelBegin;
+        });
+    const syntax::Token *SelLimit = std::partition_point(
+        SelFirst, AllSpelledTokens.end(), [&](const syntax::Token &Tok) {
+          return SM.getFileOffset(Tok.location()) < SelEnd;
+        });
+    // Precompute selectedness and offset for selected spelled tokens.
+    for (const syntax::Token *T = SelFirst; T < SelLimit; ++T) {
       // As well as comments, don't count semicolons as real tokens.
       // They're not properly claimed as expr-statement is missing from the AST.
-      if (Tok.kind() == tok::comment || Tok.kind() == tok::semi)
+      if (T->kind() == tok::comment || T->kind() == tok::semi)
         continue;
-
-      Tokens.emplace_back();
-      TokInfo &S = Tokens.back();
-      S.StartOffset = SM.getFileOffset(Tok.location());
-      S.EndOffset = S.StartOffset + Tok.length();
-      if (S.StartOffset >= SelBegin && S.EndOffset <= SelEnd)
+      SpelledTokens.emplace_back();
+      Tok &S = SpelledTokens.back();
+      S.Offset = SM.getFileOffset(T->location());
+      if (S.Offset >= SelBegin && S.Offset + T->length() <= SelEnd)
         S.Selected = SelectionTree::Complete;
-      else if (S.EndOffset > SelBegin && S.StartOffset < SelEnd)
-        S.Selected = SelectionTree::Partial;
       else
-        S.Selected = SelectionTree::Unselected;
-      S.Claimed = false;
+        S.Selected = SelectionTree::Partial;
     }
   }
 
-  // Associates any tokens overlapping [Begin, End) with an AST node.
-  // Tokens that were already claimed by another AST node are not claimed again.
-  // Updates Result if the node is selected in the sense of SelectionTree.
-  void claim(unsigned Begin, unsigned End, SelectionTree::Selection &Result) {
-    assert(Begin <= End);
-
-    // Fast-path for missing the selection entirely.
-    if (Begin >= SelEnd || End <= SelBegin)
-      return;
-
-    // We will consider the range (at least partially) selected if it hit any
-    // selected and previously unclaimed token.
-    bool ClaimedAnyToken = false;
-    // The selection is (at most) partial if:
-    // - any claimed token is partially selected
-    // - any token in the range is unselected
-    bool PartialSelection = false;
-
-    // Find the first token that (maybe) overlaps the claimed range.
-    auto Start = llvm::partition_point(Tokens, [&](const TokInfo &Tok) {
-      return Tok.EndOffset <= Begin;
+  // Test whether a consecutive range of tokens is selected.
+  // The tokens are taken from the expanded token stream.
+  SelectionTree::Selection
+  test(llvm::ArrayRef<syntax::Token> ExpandedTokens) const {
+    if (SpelledTokens.empty())
+      return NoTokens;
+    SelectionTree::Selection Result = NoTokens;
+    while (!ExpandedTokens.empty()) {
+      // Take consecutive tokens from the same context together for efficiency.
+      FileID FID = SM.getFileID(ExpandedTokens.front().location());
+      auto Batch = ExpandedTokens.take_while([&](const syntax::Token &T) {
+        return SM.getFileID(T.location()) == FID;
       });
-    // Iterate over every token that overlaps the range.
-    // Claim selected tokens, and update the two result flags.
-    for (auto It = Start; It != Tokens.end() && It->StartOffset < End; ++It) {
-      if (It->Selected) {
-        if (!It->Claimed) {
-          // Token is selected, in the node's range, and unclaimed; claim it.
-          It->Claimed = true;
-          ClaimedAnyToken = true;
-          // If the token was only partially selected, so is the node.
-          PartialSelection |= (It->Selected == SelectionTree::Partial);
-        }
-      } else {
-        // If the node covers an unselected token, it's not completely selected.
-        PartialSelection = true;
-      }
-    }
+      assert(!Batch.empty());
+      ExpandedTokens = ExpandedTokens.drop_front(Batch.size());
+
+      update(Result, testChunk(FID, Batch));
+    return Result;
+
+  // Cheap check whether any of the tokens in R might be selected.
+  // If it returns false, test() will return NoTokens or Unselected.
+  // If it returns true, test() may return any value.
+  bool mayHit(SourceRange R) const {
+    if (SpelledTokens.empty())
+      return false;
+    auto B = SM.getDecomposedLoc(R.getBegin());
+    auto E = SM.getDecomposedLoc(R.getEnd());
+    if (B.first == SelFile && E.first == SelFile)
+      if (E.second < SpelledTokens.front().Offset ||
+          B.second > SpelledTokens.back().Offset)
+        return false;
+    return true;
 
-    // If some tokens were previously claimed (Result != Unselected), we may
-    // upgrade from Partial->Complete, even if no new tokens were claimed.
-    // Important for [[int a]].
-    if (ClaimedAnyToken || Result) {
-      Result = std::max(Result, PartialSelection ? SelectionTree::Partial
-                                                 : SelectionTree::Complete);
-    }
+private:
+  // Hit-test a consecutive range of tokens from a single file ID.
+  SelectionTree::Selection
+  testChunk(FileID FID, llvm::ArrayRef<syntax::Token> Batch) const {
+    assert(!Batch.empty());
+    SourceLocation StartLoc = Batch.front().location();
+    // There are several possible categories of FileID depending on how the
+    // preprocessor was used to generate these tokens:
+    //   main file, #included file, macro args, macro bodies.
+    // We need to identify the main-file tokens that represent Batch, and
+    // determine whether we want to exclusively claim them. Regular tokens
+    // represent one AST construct, but a macro invocation can represent many.
+
+    // Handle tokens written directly in the main file.
+    if (FID == SelFile) {
+      return testTokenRange(SM.getFileOffset(Batch.front().location()),
+                            SM.getFileOffset(Batch.back().location()));
+    }
+
+    // Handle tokens in another file #included into the main file.
+    // Check if the #include is selected, but don't claim it exclusively.
+    if (StartLoc.isFileID()) {
+      for (SourceLocation Loc = Batch.front().location(); Loc.isValid();
+           Loc = SM.getIncludeLoc(SM.getFileID(Loc))) {
+        if (SM.getFileID(Loc) == SelFile)
+          // FIXME: use whole #include directive, not just the filename string.
+          return testToken(SM.getFileOffset(Loc));
+      }
+      return NoTokens;
+    }
+
+    assert(StartLoc.isMacroID());
+    // Handle tokens that were passed as a macro argument.
+    SourceLocation ArgStart = SM.getTopMacroCallerLoc(StartLoc);
+    if (SM.getFileID(ArgStart) == SelFile) {
+      SourceLocation ArgEnd = SM.getTopMacroCallerLoc(Batch.back().location());
+      return testTokenRange(SM.getFileOffset(ArgStart),
+                            SM.getFileOffset(ArgEnd));
+    }
+
+    // Handle tokens produced by non-argument macro expansion.
+    // Check if the macro name is selected, don't claim it exclusively.
+    auto Expansion = SM.getDecomposedExpansionLoc(StartLoc);
+    if (Expansion.first == SelFile)
+      // FIXME: also check ( and ) for function-like macros?
+      return testToken(Expansion.second);
+    else
+      return NoTokens;
+
+  // Is the closed token range [Begin, End] selected?
+  SelectionTree::Selection testTokenRange(unsigned Begin, unsigned End) const {
+    assert(Begin <= End);
+    // Outside the selection entirely?
+    if (End < SpelledTokens.front().Offset ||
+        Begin > SpelledTokens.back().Offset)
+      return SelectionTree::Unselected;
+
+    // Compute range of tokens.
+    auto B = llvm::partition_point(
+        SpelledTokens, [&](const Tok &T) { return T.Offset < Begin; });
+    auto E = std::partition_point(
+        B, SpelledTokens.end(), [&](const Tok &T) { return T.Offset <= End; });
+
+    // Aggregate selectedness of tokens in range.
+    bool ExtendsOutsideSelection = Begin < SpelledTokens.front().Offset ||
+                                   End > SpelledTokens.back().Offset;
+    SelectionTree::Selection Result =
+        ExtendsOutsideSelection ? SelectionTree::Unselected : NoTokens;
+    for (auto It = B; It != E; ++It)
+      update(Result, It->Selected);
+    return Result;
+  }
+
+  // Is the token at `Offset` selected?
+  SelectionTree::Selection testToken(unsigned Offset) const {
+    // Outside the selection entirely?
+    if (Offset < SpelledTokens.front().Offset ||
+        Offset > SpelledTokens.back().Offset)
+      return SelectionTree::Unselected;
+    // Find the token, if it exists.
+    auto It = llvm::partition_point(
+        SpelledTokens, [&](const Tok &T) { return T.Offset < Offset; });
+    if (It != SpelledTokens.end() && It->Offset == Offset)
+      return It->Selected;
+    return NoTokens;
   }
 
-private:
-  struct TokInfo {
-    unsigned StartOffset;
-    unsigned EndOffset;
+  struct Tok {
+    unsigned Offset;
     SelectionTree::Selection Selected;
-    bool Claimed;
-    bool operator<(const TokInfo &Other) const {
-      return StartOffset < Other.StartOffset;
-    }
   };
-  std::vector<TokInfo> Tokens;
-  unsigned SelBegin, SelEnd;
+  std::vector<Tok> SpelledTokens;
+  FileID SelFile;
+  const SourceManager &SM;
 };
 
 // Show the type of a node for debugging.
 void printNodeKind(llvm::raw_ostream &OS, const DynTypedNode &N) {
   if (const TypeLoc *TL = N.get<TypeLoc>()) {
     // TypeLoc is a hierarchy, but has only a single ASTNodeKind.
     // Synthesize the name from the Type subclass (except for QualifiedTypeLoc).
     if (TL->getTypeLocClass() == TypeLoc::Qualified)
public:
   static std::deque<Node> collect(ASTContext &AST,
                                   const syntax::TokenBuffer &Tokens,
                                   const PrintingPolicy &PP, unsigned Begin,
                                   unsigned End, FileID File) {
     SelectionVisitor V(AST, Tokens, PP, Begin, End, File);
     V.TraverseAST(AST);
     assert(V.Stack.size() == 1 && "Unpaired push/pop?");
     assert(V.Stack.top() == &V.Nodes.front());
-    // We selected TUDecl if tokens were unclaimed (or the file is empty).
-    SelectionTree::Selection UnclaimedTokens = SelectionTree::Unselected;
-    V.Claimed.claim(Begin, End, UnclaimedTokens);
-    if (UnclaimedTokens || V.Nodes.size() == 1) {
-      StringRef FileContent = AST.getSourceManager().getBufferData(File);
-      // Don't require the trailing newlines to be selected.
-      bool SelectedAll = Begin == 0 && End >= FileContent.rtrim().size();
-      V.Stack.top()->Selected =
-          SelectedAll ? SelectionTree::Complete : SelectionTree::Partial;
-    }
     return std::move(V.Nodes);
   }
 
   // We traverse all "well-behaved" nodes the same way:
   //  - push the node onto the stack
   //  - traverse its children recursively
   //  - pop it from the stack
   //  - hit testing: is intersection(node, selection) - union(children) empty?
private:
 
   SelectionVisitor(ASTContext &AST, const syntax::TokenBuffer &Tokens,
                    const PrintingPolicy &PP, unsigned SelBegin, unsigned SelEnd,
                    FileID SelFile)
       : SM(AST.getSourceManager()), LangOpts(AST.getLangOpts()),
 #ifndef NDEBUG
         PrintPolicy(PP),
 #endif
-        Claimed(Tokens.spelledTokens(SelFile), SM, SelBegin, SelEnd),
-        SelFile(SelFile),
-        SelBeginTokenStart(SM.getFileOffset(Lexer::GetBeginningOfToken(
-            SM.getComposedLoc(SelFile, SelBegin), SM, LangOpts))),
-        SelEnd(SelEnd) {
+        TokenBuf(Tokens), SelChecker(Tokens, SelFile, SelBegin, SelEnd, SM),
+        UnclaimedExpandedTokens(Tokens.expandedTokens()) {
     // Ensure we have a node for the TU decl, regardless of traversal scope.
     Nodes.emplace_back();
     Nodes.back().ASTNode = DynTypedNode::create(*AST.getTranslationUnitDecl());
     Nodes.back().Parent = nullptr;
     Nodes.back().Selected = SelectionTree::Unselected;
     Stack.push(&Nodes.back());
   }
 
#endif
   // Here the first ("leftmost") sibling claims the tokens it wants, and the
   // other sibling gets what's left. So selecting "int" only includes the left
   // VarDecl in the selection tree.
 
   // An optimization for a common case: nodes outside macro expansions that
   // don't intersect the selection may be recursively skipped.
   bool canSafelySkipNode(const DynTypedNode &N) {
     SourceRange S = N.getSourceRange();
-    auto B = SM.getDecomposedLoc(S.getBegin());
-    auto E = SM.getDecomposedLoc(S.getEnd());
-    // Node lies in a macro expansion?
-    if (B.first != SelFile || E.first != SelFile)
-      return false;
-    // Node intersects selection tokens?
-    if (B.second < SelEnd && E.second >= SelBeginTokenStart)
-      return false;
-    // Otherwise, allow skipping over the node.
+    if (!SelChecker.mayHit(S)) {
       dlog("{1}skip: {0}", printNodeToString(N, PrintPolicy), indent());
       dlog("{1}skipped range = {0}", S.printToString(SM), indent(1));
       return true;
     }
+    return false;
+  }
 
   // There are certain nodes we want to treat as leaves in the SelectionTree,
   // although they do have children.
   bool shouldSkipChildren(const Stmt *X) const {
     // UserDefinedLiteral (e.g. 12_i) has two children (12 and _i).
     // Unfortunately TokenBuffer sees 12_i as one token and can't split it.
     // So we treat UserDefinedLiteral as a leaf node, owning the token.
     return llvm::isa<UserDefinedLiteral>(X);
   }
 
   // Pushes a node onto the ancestor stack. Pairs with pop().
   // Performs early hit detection for some nodes (on the earlySourceRange).
   void push(DynTypedNode Node) {
     SourceRange Early = earlySourceRange(Node);
     dlog("{1}push: {0}", printNodeToString(Node, PrintPolicy), indent());
     Nodes.emplace_back();
     Nodes.back().ASTNode = std::move(Node);
     Nodes.back().Parent = Stack.top();
+    Nodes.back().Selected = NoTokens;
     Stack.push(&Nodes.back());
     claimRange(Early, Nodes.back().Selected);
-    // Early hit detection never selects the whole node.
-    if (Nodes.back().Selected)
-      Nodes.back().Selected = SelectionTree::Partial;
   }
 
   // Pops a node off the ancestor stack, and finalizes it. Pairs with push().
   // Performs primary hit detection.
   void pop() {
     Node &N = *Stack.top();
     dlog("{1}pop: {0}", printNodeToString(N.ASTNode, PrintPolicy), indent(-1));
     claimRange(N.ASTNode.getSourceRange(), N.Selected);
+    if (N.Selected == NoTokens)
+      N.Selected = SelectionTree::Unselected;
     if (N.Selected || !N.Children.empty()) {
       // Attach to the tree.
       N.Parent->Children.push_back(&N);
     } else {
       // Neither N any children are selected, it doesn't belong in the tree.
       assert(&N == &Nodes.back());
       Nodes.pop_back();
     }
SourceRange earlySourceRange(const DynTypedNode &N) {
     return SourceRange();
   }
 
   // Perform hit-testing of a complete Node against the selection.
   // This runs for every node in the AST, and must be fast in common cases.
   // This is usually called from pop(), so we can take children into account.
   // The existing state of Result is relevant (early/late claims can interact).
   void claimRange(SourceRange S, SelectionTree::Selection &Result) {
-    if (!S.isValid())
-      return;
-    // toHalfOpenFileRange() allows selection of constructs in macro args. e.g:
-    //   #define LOOP_FOREVER(Body) for(;;) { Body }
-    //   void IncrementLots(int &x) {
-    //     LOOP_FOREVER( ++x; )
-    //   }
-    // Selecting "++x" or "x" will do the right thing.
-    auto Range = toHalfOpenFileRange(SM, LangOpts, S);
-    assert(Range && "We should be able to get the File Range");
-    dlog("{1}claimRange: {0}", Range->printToString(SM), indent());
-    auto B = SM.getDecomposedLoc(Range->getBegin());
-    auto E = SM.getDecomposedLoc(Range->getEnd());
-    // Otherwise, nodes in macro expansions can't be selected.
-    if (B.first != SelFile || E.first != SelFile)
-      return;
-    // Attempt to claim the remaining range. If there's nothing to claim, only
-    // children were selected.
-    Claimed.claim(B.second, E.second, Result);
-    if (Result)
-      dlog("{1}hit selection: {0}",
-           SourceRange(SM.getComposedLoc(B.first, B.second),
-                       SM.getComposedLoc(E.first, E.second))
-               .printToString(SM),
-           indent());
+    for (const auto &ClaimedRange :
+         UnclaimedExpandedTokens.erase(TokenBuf.expandedTokens(S)))
+      update(Result, SelChecker.test(ClaimedRange));
+
+    if (Result && Result != NoTokens)
+      dlog("{1}hit selection: {0}", S.printToString(SM), indent());
   }
 
   std::string indent(int Offset = 0) {
     // Cast for signed arithmetic.
     int Amount = int(Stack.size()) + Offset;
     assert(Amount >= 0);
     return std::string(Amount, ' ');
   }
 
   SourceManager &SM;
   const LangOptions &LangOpts;
 #ifndef NDEBUG
   const PrintingPolicy &PrintPolicy;
 #endif
+  const syntax::TokenBuffer &TokenBuf;
   std::stack<Node *> Stack;
-  SelectedTokens Claimed;
+  SelectionTester SelChecker;
+  IntervalSet<syntax::Token> UnclaimedExpandedTokens;
   std::deque<Node> Nodes; // Stable pointers as we add more nodes.
-  FileID SelFile;
-  // If the selection start slices a token in half, the beginning of that token.
-  // This is useful for checking whether the end of a token range overlaps
-  // the selection: range.end < SelBeginTokenStart is equivalent to
-  // range.end + measureToken(range.end) < SelBegin (assuming range.end points
-  // to a token), and it saves a lex every time.
-  unsigned SelBeginTokenStart;
-  unsigned SelEnd;
 };
 
 } // namespace
 
 void SelectionTree::print(llvm::raw_ostream &OS, const SelectionTree::Node &N,
                           int Indent) const {
   if (N.Selected)
     OS.indent(Indent - 1) << (N.Selected == SelectionTree::Complete ? '*'

 std::string SelectionTree::Node::kind() const {
   std::string S;
   llvm::raw_string_ostream OS(S);
   printNodeKind(OS, ASTNode);
   return std::move(OS.str());
 }
 
 // Decide which selection emulates a "point" query in between characters.
 static std::pair<unsigned, unsigned> pointBounds(unsigned Offset, FileID FID,
                                                  ASTContext &AST) {
   StringRef Buf = AST.getSourceManager().getBufferData(FID);
   // Edge-cases where the choice is forced.
   if (Buf.size() == 0)
     return {0, 0};
   if (Offset == 0)
     return {0, 1};
   if (Offset == Buf.size())

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

Case Cases[] = {
       {
           R"cpp(
             void foo() { [[if (1^11) { return; } else {^ }]] }
           )cpp",
           "IfStmt",
       },
       {
           R"cpp(
+            int x(int);
+            #define M(foo) x(foo)
+            int a = 42;
+            int b = M([[^a]]);
+          )cpp",
+          "DeclRefExpr",
+      },
+      {
+          R"cpp(
             void foo();
             #define CALL_FUNCTION(X) X()
             void bar() { CALL_FUNCTION([[f^o^o]]); }
           )cpp",
           "DeclRefExpr",
       },
       {
           R"cpp(
const char *Cases[] = {
       R"cpp( int abc, xyz = [[a^bc^]]; )cpp",
       R"cpp( int abc, xyz = $C[[^abc^]]; )cpp",
       R"cpp(
         void foo() {
           [[if ([[1^11]]) $C[[{
             $C[[return]];
           }]] else [[{^
           }]]]]
+          char z;
         }
       )cpp",
       R"cpp(
           template <class T>
           struct unique_ptr {};
           void foo(^$C[[unique_ptr<$C[[unique_ptr<$C[[int]]>]]>]]^ a) {}
       )cpp",
       R"cpp(int a = [[5 >^> 1]];)cpp",
-      R"cpp([[
+      R"cpp(
         #define ECHO(X) X
-        ECHO(EC^HO([[$C[[int]]) EC^HO(a]]));
-      ]])cpp",
+        ECHO(EC^HO($C[[int]]) EC^HO(a));
+      )cpp",
       R"cpp( $C[[^$C[[int]] a^]]; )cpp",
       R"cpp( $C[[^$C[[int]] a = $C[[5]]^]]; )cpp",
   };
   for (const char *C : Cases) {
     Annotations Test(C);
     auto AST = TestTU::withCode(Test.code()).build();
     auto T = makeSelectionTree(C, AST);
 
#include "Expand.inc"
   auto AST = TU.build();
   EXPECT_THAT(AST.getDiagnostics(), ::testing::IsEmpty());
   auto T = makeSelectionTree(Case, AST);
 
   EXPECT_EQ("BreakStmt", T.commonAncestor()->kind());
   EXPECT_EQ("WhileStmt", T.commonAncestor()->Parent->kind());
 }
 
+TEST(SelectionTest, IncludedFile) {
+  const char *Case = R"cpp(
+    void test() {
+#include "Exp^and.inc"
+        break;
+    }
+  )cpp";
+  Annotations Test(Case);
+  auto TU = TestTU::withCode(Test.code());
+  TU.AdditionalFiles["Expand.inc"] = "while(1)\n";
+  auto AST = TU.build();
+  auto T = makeSelectionTree(Case, AST);
+
+  EXPECT_EQ("WhileStmt", T.commonAncestor()->kind());
+}
+
+TEST(SelectionTest, MacroArgExpansion) {
+  // If a macro arg is expanded several times, we consider them all selected.
+  const char *Case = R"cpp(
+    int mul(int, int);
+    #define SQUARE(X) mul(X, X);
+    int nine = SQUARE(^3);
+  )cpp";
+  Annotations Test(Case);
+  auto AST = TestTU::withCode(Test.code()).build();
+  auto T = makeSelectionTree(Case, AST);
+  // Unfortunately, this makes the common ancestor the CallExpr...
+  // FIXME: hack around this by picking one?
+  EXPECT_EQ("CallExpr", T.commonAncestor()->kind());
+  EXPECT_FALSE(T.commonAncestor()->Selected);
+  EXPECT_EQ(2u, T.commonAncestor()->Children.size());
+  for (const auto* N : T.commonAncestor()->Children) {
+    EXPECT_EQ("IntegerLiteral", N->kind());
+    EXPECT_TRUE(N->Selected);
+  }
+
+  // Verify that the common assert() macro doesn't suffer from this.
+  // (This is because we don't associate the stringified token with the arg).
+  Case = R"cpp(
+    void die(const char*);
+    #define assert(x) (x ? (void)0 : die(#x)
+    void foo() { assert(^42); }
+  )cpp";
+  Test = Annotations(Case);
+  AST = TestTU::withCode(Test.code()).build();
+  T = makeSelectionTree(Case, AST);
+
+  EXPECT_EQ("IntegerLiteral", T.commonAncestor()->kind());
+}
+
 TEST(SelectionTest, Implicit) {
   const char* Test = R"cpp(
     struct S { S(const char*); };
     int f(S);
     int x = f("^");
   )cpp";
   auto AST = TestTU::withCode(Annotations(Test).code()).build();
   auto T = makeSelectionTree(Test, AST);

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

void f(int b = [[1]]) {
       goto label;
       label:
         a = [[1]];
     }
   )cpp";
   EXPECT_UNAVAILABLE(UnavailableCases);
 
   // vector of pairs of input and output strings
-  const std::vector<std::pair<llvm::StringLiteral, llvm::StringLiteral>>
+  const std::vector<std::pair<std::string, std::string>>
       InputOutputs = {
           // extraction from variable declaration/assignment
           {R"cpp(void varDecl() {
                    int a = 5 * (4 + (3 [[- 1)]]);
                  })cpp",
            R"cpp(void varDecl() {
                    auto dummy = (3 - 1); int a = 5 * (4 + dummy);
                  })cpp"},
const std::vector<std::pair<std::string, std::string>>
                    auto dummy = PLUS(1+a); int y = dummy;
                  })cpp"},
           // ensure InsertionPoint isn't inside a macro
           {R"cpp(#define LOOP(x) while (1) {a = x;}
                  void f(int a) {
                    if(1)
                     LOOP(5 + [[3]])
                  })cpp",
-           /*FIXME: It should be extracted like this. SelectionTree needs to be
-             * fixed for macros.
             R"cpp(#define LOOP(x) while (1) {a = x;}
                  void f(int a) {
                    auto dummy = 3; if(1)
                     LOOP(5 + dummy)
-                })cpp"},*/
-           R"cpp(#define LOOP(x) while (1) {a = x;}
-                 void f(int a) {
-                   auto dummy = LOOP(5 + 3); if(1)
-                    dummy
                  })cpp"},
           {R"cpp(#define LOOP(x) do {x;} while(1);
                  void f(int a) {
                    if(1)
                     LOOP(5 + [[3]])
                  })cpp",
            R"cpp(#define LOOP(x) do {x;} while(1);
                  void f(int a) {
void f(const int c) {
   std::string TemplateFailInput = R"cpp(
     template<typename T>
     void f() {
       [[int x;]]
     }
   )cpp";
   EXPECT_EQ(apply(TemplateFailInput), "unavailable");
 
-  // FIXME: This should be extractable after selectionTree works correctly for
-  // macros (currently it doesn't select anything for the following case)
-  std::string MacroFailInput = R"cpp(
+  std::string MacroInput = R"cpp(
     #define F(BODY) void f() { BODY }
     F ([[int x = 0;]])
   )cpp";
-  EXPECT_EQ(apply(MacroFailInput), "unavailable");
+  std::string MacroOutput = R"cpp(
+    #define F(BODY) void f() { BODY }
+    void extracted() {
+int x = 0;
+}
+F (extracted();)
+  )cpp";
+  EXPECT_EQ(apply(MacroInput), MacroOutput);
 
   // Shouldn't crash.
   EXPECT_EQ(apply("void f([[int a]]);"), "unavailable");
   // Don't extract if we select the entire function body (CompoundStmt).
   std::string CompoundFailInput = R"cpp(
     void f() [[{
       int a;
     }]]

clang/include/clang/Tooling/Syntax/Tokens.h

 /// FIXME: allow to map from spelled to expanded tokens when use-case shows up.
 /// FIXME: allow mappings into macro arguments.
 class TokenBuffer {
 public:
   TokenBuffer(const SourceManager &SourceMgr) : SourceMgr(&SourceMgr) {}
   /// All tokens produced by the preprocessor after all macro replacements,
   /// directives, etc. Source locations found in the clang AST will always
   /// point to one of these tokens.
+  /// Tokens are in TU order (per SourceManager::isBeforeInTranslationUnit()).
   /// FIXME: figure out how to handle token splitting, e.g. '>>' can be split
   ///        into two '>' tokens by the parser. However, TokenBuffer currently
   ///        keeps it as a single '>>' token.
   llvm::ArrayRef<syntax::Token> expandedTokens() const {
     return ExpandedTokens;
   }
 
+  /// Returns the subrange of expandedTokens() corresponding to the closed
+  /// token range R.
+  llvm::ArrayRef<syntax::Token> expandedTokens(SourceRange R) const;
+
   /// Find the subrange of spelled tokens that produced the corresponding \p
   /// Expanded tokens.
   ///
   /// EXPECTS: \p Expanded is a subrange of expandedTokens().
   ///
   /// Will fail if the expanded tokens do not correspond to a
   /// sequence of spelled tokens. E.g. for the following example:
   ///

clang/lib/Tooling/Syntax/Tokens.cpp

llvm::StringRef FileRange::text(const SourceManager &SM) const {
   StringRef Text = SM.getBufferData(File, &Invalid);
   if (Invalid)
     return "";
   assert(Begin <= Text.size());
   assert(End <= Text.size());
   return Text.substr(Begin, length());
 }
 
+llvm::ArrayRef<syntax::Token> TokenBuffer::expandedTokens(SourceRange R) const {
+  if (R.isInvalid())
+    return {};
+  const Token *Begin =
+      llvm::partition_point(expandedTokens(), [&](const syntax::Token &T) {
+        return SourceMgr->isBeforeInTranslationUnit(T.location(), R.getBegin());
+      });
+  const Token *End =
+      llvm::partition_point(expandedTokens(), [&](const syntax::Token &T) {
+        return !SourceMgr->isBeforeInTranslationUnit(R.getEnd(), T.location());
+      });
+  if (Begin > End)
+    return {};
+  return {Begin, End};
+}
+
 std::pair<const syntax::Token *, const TokenBuffer::Mapping *>
 TokenBuffer::spelledForExpandedToken(const syntax::Token *Expanded) const {
   assert(Expanded);
   assert(ExpandedTokens.data() <= Expanded &&
          Expanded < ExpandedTokens.data() + ExpandedTokens.size());
 
   auto FileIt = Files.find(
       SourceMgr->getFileID(SourceMgr->getExpansionLoc(Expanded->location())));

clang/unittests/Tooling/Syntax/TokensTest.cpp

 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/VirtualFileSystem.h"
 #include "llvm/Support/raw_os_ostream.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Testing/Support/Annotations.h"
 #include "llvm/Testing/Support/SupportHelpers.h"
+#include "gmock/gmock.h"
 #include <cassert>
 #include <cstdlib>
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <memory>
 #include <ostream>
 #include <string>
 
recordTokens(R"cpp(
     ID(1)
     #pragma lalala
     not_mapped
   )cpp");
   EXPECT_THAT(Buffer.spelledForExpanded(findExpanded("not_mapped")),
               ValueIs(SameRange(findSpelled("not_mapped"))));
 }
 
+TEST_F(TokenBufferTest, ExpandedTokensForRange) {
+  recordTokens(R"cpp(
+    #define SIGN(X) X##_washere
+    A SIGN(B) C SIGN(D) E SIGN(F) G
+  )cpp");
+
+  SourceRange R(findExpanded("C").front().location(),
+                findExpanded("F_washere").front().location());
+  // Sanity check: expanded and spelled tokens are stored separately.
+  EXPECT_THAT(Buffer.expandedTokens(R),
+              SameRange(findExpanded("C D_washere E F_washere")));
+  EXPECT_THAT(Buffer.expandedTokens(SourceRange()), testing::IsEmpty());
+}
+
 TEST_F(TokenBufferTest, ExpansionStartingAt) {
   // Object-like macro expansions.
   recordTokens(R"cpp(
     #define FOO 3+4
     int a = FOO 1;
     int b = FOO 2;
   )cpp");