[clangd] Bring back early-claim approach to fix a selection-tree regression.

[hokein]

Summary

The early-claim hack was removed in 96f5cc1ee417f863f85756d1e56b1bed1bd76a7e,
we see a regression about captured var-decl in lambda, see https://github.com/clangd/clangd/issues/990.

[sammccall]

clang-tools-extra/clangd/Selection.cpp

836

I don't think this is an e.g., but rather is suitable for specifically this case.

If there are other cases, we likely should be testing DeclaratorDecl here, and excluding Constructor/Destructor as before.

852

I wonder whether we want to keep the CXXConstructExpr hack now, up to you.

[hokein]

refine the comment.

[hokein]

clang-tools-extra/clangd/Selection.cpp

836

refined the comment.

Testing DeclaratorDecl is safer, but we miss opportunities to know other cases (hopefully, there should be no other cases). For now, I'm leaning towards merely handling this specific case. If it turns out more cases, we can definitely use the DeclaratorDecl.

852

if we remove it, than we need to add the CXXCtorInitializer case back to the earlySourceRange :(

clang-tools-extra/clangd/Selection.cpp

#endif
   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);
   }
 
   // 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));
     claimTokensFor(N.ASTNode, 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();
     }
     Stack.pop();
   }
 
+  // Returns the range of tokens that this node will claim directly, and
+  // is not available to the node's children.
+  // Usually empty, but sometimes children cover tokens but shouldn't own them.
+  SourceRange earlySourceRange(const DynTypedNode &N) {
+    if (const Decl *D = N.get<Decl>()) {
+      // We want the name in the var-decl to be claimed by the decl itself and
+      // not by any children. Ususally, we don't need this, because source
+      // ranges of children are not overlapped with their parent's.
+      // An exception is lambda captured var decl, where AutoTypeLoc is
+      // overlapped with the name loc.
+      //    auto fun = [bar = foo]() { ... }
+      //                ~~~~~~~~~   VarDecl
+      //                ~~~         |- AutoTypeLoc
+      if (const auto *DD = llvm::dyn_cast<VarDecl>(D))
+        return DD->getLocation();
+    }
+
+    return SourceRange();
+  }
+
   // Claim tokens for N, after processing its children.
   // By default this claims all unclaimed tokens in getSourceRange().
   // We override this if we want to claim fewer tokens (e.g. there are gaps).
   void claimTokensFor(const DynTypedNode &N, SelectionTree::Selection &Result) {
     // CXXConstructExpr often shows implicit construction, like `string s;`.
     // Don't associate any tokens with it unless there's some syntax like {}.
     // This prevents it from claiming 's', its primary location.
     if (const auto *CCE = N.get<CXXConstructExpr>()) {
       claimRange(CCE->getParenOrBraceRange(), Result);
       return;
     }
     // ExprWithCleanups is always implicit. It often wraps CXXConstructExpr.
     // Prevent it claiming 's' in the case above.

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

Case Cases[] = {
         enum Bar : [[Fo^o]] {};
       )cpp",
        "TypedefTypeLoc"},
       {R"cpp(
         typedef int Foo;
         enum Bar : [[Fo^o]];
       )cpp",
        "TypedefTypeLoc"},
+
+      // lambda captured var-decl
+      {R"cpp(
+        void test(int bar) {
+          auto l = [^[[foo = bar]]] { };
+        })cpp",
+       "VarDecl"},
   };
 
   for (const Case &C : Cases) {
     trace::TestTracer Tracer;
     Annotations Test(C.Code);
 
     TestTU TU;
     TU.Code = std::string(Test.code());