[clangd] Fix selection on multi-dimensional array.

This involves separating out the concepts of "which tokens should we
descend into this node for" vs "which tokens should this node claim".

Reviewed By: sammccall

Differential Revision: https://reviews.llvm.org/D116218

2 files changed, 32 insertions, 1 deletions

diff --git a/clang-tools-extra/clangd/Selection.cpp b/clang-tools-extra/clangd/Selection.cpp
index 0b10c7a3a6f9..2024228e2b58 100644
--- a/clang-tools-extra/clangd/Selection.cpp
+++ b/clang-tools-extra/clangd/Selection.cpp
@@ -58,6 +58,7 @@ void recordMetrics(const SelectionTree &S, const LangOptions &Lang) {
     SelectionUsedRecovery.record(0, LanguageLabel); // unused.
 }
 
+// Return the range covering a node and all its children.
 SourceRange getSourceRange(const DynTypedNode &N) {
   // MemberExprs to implicitly access anonymous fields should not claim any
   // tokens for themselves. Given:
@@ -702,7 +703,7 @@ private:
   void pop() {
     Node &N = *Stack.top();
     dlog("{1}pop: {0}", printNodeToString(N.ASTNode, PrintPolicy), indent(-1));
-    claimRange(getSourceRange(N.ASTNode), N.Selected);
+    claimTokensFor(N.ASTNode, N.Selected);
     if (N.Selected == NoTokens)
       N.Selected = SelectionTree::Unselected;
     if (N.Selected || !N.Children.empty()) {
@@ -744,6 +745,28 @@ private:
     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) {
+    if (const auto *TL = N.get<TypeLoc>()) {
+      // e.g. EltType Foo[OuterSize][InnerSize];
+      //      ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ArrayTypeLoc (Outer)
+      //      ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ |-ArrayTypeLoc (Inner)
+      //      ~~~~~~~                           | |-RecordType
+      //                             ~~~~~~~~~  | `-Expr (InnerSize)
+      //                  ~~~~~~~~~             `-Expr (OuterSize)
+      // Inner ATL must not claim its whole SourceRange, or it clobbers Outer.
+      if (TL->getAs<ArrayTypeLoc>()) {
+        claimRange(TL->getLocalSourceRange(), Result);
+        return;
+      }
+      // FIXME: maybe LocalSourceRange is a better default for TypeLocs.
+      // It doesn't seem to be usable for FunctionTypeLocs.
+    }
+    claimRange(getSourceRange(N), Result);
+  }
+
   // 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.
diff --git a/clang-tools-extra/clangd/unittests/SelectionTests.cpp b/clang-tools-extra/clangd/unittests/SelectionTests.cpp
index 6c6782a097db..971487c9cd27 100644
--- a/clang-tools-extra/clangd/unittests/SelectionTests.cpp
+++ b/clang-tools-extra/clangd/unittests/SelectionTests.cpp
@@ -318,6 +318,14 @@ TEST(SelectionTest, CommonAncestor) {
       {"[[st^ruct {int x;}]] y;", "CXXRecordDecl"},
       {"[[struct {int x;} ^y]];", "VarDecl"},
       {"struct {[[int ^x]];} y;", "FieldDecl"},
+
+      // Tricky case: nested ArrayTypeLocs have the same token range.
+      {"const int x = 1, y = 2; int array[^[[x]]][10][y];", "DeclRefExpr"},
+      {"const int x = 1, y = 2; int array[x][10][^[[y]]];", "DeclRefExpr"},
+      {"const int x = 1, y = 2; int array[x][^[[10]]][y];", "IntegerLiteral"},
+      {"const int x = 1, y = 2; [[i^nt]] array[x][10][y];", "BuiltinTypeLoc"},
+      {"void func(int x) { int v_array[^[[x]]][10]; }", "DeclRefExpr"},
+
       // FIXME: the AST has no location info for qualifiers.
       {"const [[a^uto]] x = 42;", "AutoTypeLoc"},
       {"[[co^nst auto x = 42]];", "VarDecl"},