[clangd] SelectionTree should mark a node as fully-selected if the only claimed tokens were early-claimed.

[sammccall]

Summary

Previously they would be marked as partially-selected based on the early claim,
and never updated as no more tokens were claimed.
This affects simple VarDecls like "int x".

[sammccall]

remove printfs

[sammccall]

remove special-case handling of VarDecl in ExtractFunction that is no longer needed.

[SureYeaah]

Sorry about the delay. LGTM. Thanks!

clang-tools-extra/clangd/Selection.cpp

for (const auto& Tok : Spelled) {
       else
         S.Selected = SelectionTree::Unselected;
       S.Claimed = false;
     }
   }
 
   // Associates any tokens overlapping [Begin, End) with an AST node.
   // Tokens that were already claimed by another AST node are not claimed again.
-  // Returns whether the node is selected in the sense of SelectionTree.
-  SelectionTree::Selection claim(unsigned Begin, unsigned End) {
+  // 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 SelectionTree::Unselected;
+      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;
for (auto It = Start; It != Tokens.end() && It->StartOffset < End; ++It) {
           PartialSelection |= (It->Selected == SelectionTree::Partial);
         }
       } else {
         // If the node covers an unselected token, it's not completely selected.
         PartialSelection = true;
       }
     }
 
-    if (!ClaimedAnyToken)
-      return SelectionTree::Unselected;
-    return PartialSelection ? SelectionTree::Partial : SelectionTree::Complete;
+    // 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:
   struct TokInfo {
     unsigned StartOffset;
     unsigned EndOffset;
     SelectionTree::Selection Selected;
     bool Claimed;
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).
-    if (V.Nodes.size() == 1 || V.Claimed.claim(Begin, End)) {
+    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);
   }
#endif
   // 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();
-    // Early hit detection never selects the whole node.
     Stack.push(&Nodes.back());
-    Nodes.back().Selected =
-        claimRange(Early) ? SelectionTree::Partial : SelectionTree::Unselected;
+    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));
-    if (auto Sel = claimRange(N.ASTNode.getSourceRange()))
-      N.Selected = Sel;
+    claimRange(N.ASTNode.getSourceRange(), N.Selected);
     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();
     }
if (const Decl *D = N.get<Decl>()) {
         return VD->getLocation();
     }
     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.
-  SelectionTree::Selection claimRange(SourceRange S) {
+  // The existing state of Result is relevant (early/late claims can interact).
+  void claimRange(SourceRange S, SelectionTree::Selection &Result) {
     if (!S.isValid())
-      return SelectionTree::Unselected;
+      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 SelectionTree::Unselected;
+      return;
     // Attempt to claim the remaining range. If there's nothing to claim, only
     // children were selected.
-    SelectionTree::Selection Result = Claimed.claim(B.second, E.second);
+    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());
-    return Result;
   }
 
   std::string indent(int Offset = 0) {
     // Cast for signed arithmetic.
     int Amount = int(Stack.size()) + Offset;
     assert(Amount >= 0);
     return std::string(Amount, ' ');
   }

clang-tools-extra/clangd/refactor/tweaks/ExtractFunction.cpp

 // having to change the selection range. Also, this means that any scope that
 // begins in selection range, ends in selection range and any scope that begins
 // outside the selection range, ends outside as well.
 const Node *getParentOfRootStmts(const Node *CommonAnc) {
   if (!CommonAnc)
     return nullptr;
   switch (CommonAnc->Selected) {
   case SelectionTree::Selection::Unselected:
+    // Typicaly a block, with the { and } unselected, could also be ForStmt etc
     // Ensure all Children are RootStmts.
     return llvm::all_of(CommonAnc->Children, isRootStmt) ? CommonAnc : nullptr;
   case SelectionTree::Selection::Partial:
-    // Treat Partially selected VarDecl as completely selected since
-    // SelectionTree doesn't always select VarDecls correctly.
-    // FIXME: Remove this after D66872 is upstream)
-    if (!CommonAnc->ASTNode.get<VarDecl>())
-      return nullptr;
-    LLVM_FALLTHROUGH;
+    // Only a fully-selected single statement can be selected.
   case SelectionTree::Selection::Complete:
     // If the Common Ancestor is completely selected, then it's a root statement
     // and its parent will be unselected.
     const Node *Parent = CommonAnc->Parent;
     // If parent is a DeclStmt, even though it's unselected, we consider it a
     // root statement and return its parent. This is done because the VarDecls
     // claim the entire selection range of the Declaration and DeclStmt is
     // always unselected.

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

const char *Cases[] = {
           struct unique_ptr {};
           void foo(^$C[[unique_ptr<$C[[unique_ptr<$C[[int]]>]]>]]^ a) {}
       )cpp",
       R"cpp(int a = [[5 >^> 1]];)cpp",
       R"cpp([[
         #define ECHO(X) X
         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);
 
     std::vector<Range> Complete, Partial;
     for (const SelectionTree::Node *N : allNodes(T))

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

TEST_F(ExtractFunctionTest, FunctionTest) {
   Context = Function;
 
   // Root statements should have common parent.
   EXPECT_EQ(apply("for(;;) [[1+2; 1+2;]]"), "unavailable");
   // Expressions aren't extracted.
   EXPECT_EQ(apply("int x = 0; [[x++;]]"), "unavailable");
   // We don't support extraction from lambdas.
   EXPECT_EQ(apply("auto lam = [](){ [[int x;]] }; "), "unavailable");
+  // Partial statements aren't extracted.
+  EXPECT_THAT(apply("int [[x = 0]];"), "unavailable");
 
   // Ensure that end of Zone and Beginning of PostZone being adjacent doesn't
   // lead to break being included in the extraction zone.
   EXPECT_THAT(apply("for(;;) { [[int x;]]break; }"), HasSubstr("extracted"));
-  // FIXME: This should be unavailable since partially selected but
-  // selectionTree doesn't always work correctly for VarDecls.
-  EXPECT_THAT(apply("int [[x = 0]];"), HasSubstr("extracted"));
   // FIXME: ExtractFunction should be unavailable inside loop construct
   // initalizer/condition.
   EXPECT_THAT(apply(" for([[int i = 0;]];);"), HasSubstr("extracted"));
   // Don't extract because needs hoisting.
   EXPECT_THAT(apply(" [[int a = 5;]] a++; "), StartsWith("fail"));
   // Don't extract return
   EXPECT_THAT(apply(" if(true) [[return;]] "), StartsWith("fail"));
-  
 }
 
 TEST_F(ExtractFunctionTest, FileTest) {
   // Check all parameters are in order
   std::string ParameterCheckInput = R"cpp(
 struct Foo {
   int x;
 };