Snap to nodes directly behind the cursor, create special rules for ParameterNodes

src/services/selectionRange.ts

@@ -17,7 +17,7 @@ namespace ts.SelectionRange {
                     break outer;
                 }
 
-                if (positionBelongsToNode(node, pos, sourceFile)) {
+                if (positionShouldSnapToNode(pos, node, nextNode, sourceFile)) {
                     // Blocks are effectively redundant with SyntaxLists.
                     // TemplateSpans, along with the SyntaxLists containing them,
                     // are a somewhat unintuitive grouping of things that should be
@@ -69,6 +69,28 @@ namespace ts.SelectionRange {
         }
     }
 
+    /**
+     * Like `ts.positionBelongsToNode`, except positions immediately after nodes
+     * count too, unless that position belongs to the next node. In effect, makes
+     * selections able to snap to preceding tokens when the cursor is on the tail
+     * end of them with only whitespace ahead.
+     * @param pos The position to check.
+     * @param node The candidate node to snap to.
+     * @param nextNode The next sibling node in the tree.
+     * @param sourceFile The source file containing the nodes.
+     */
+    function positionShouldSnapToNode(pos: number, node: Node, nextNode: Node | undefined, sourceFile: SourceFile) {
+        if (positionBelongsToNode(node, pos, sourceFile)) {
+            return true;
+        }
+        const nodeEnd = node.getEnd();
+        const nextNodeStart = nextNode && nextNode.getStart();
+        if (nodeEnd === pos) {
+            return pos !== nextNodeStart;
+        }
+        return false;
+    }
+
     const isImport = or(isImportDeclaration, isImportEqualsDeclaration);
 
     /**
@@ -100,34 +122,38 @@ namespace ts.SelectionRange {
             const closeBraceToken = Debug.assertDefined(children.pop());
             Debug.assertEqual(openBraceToken.kind, SyntaxKind.OpenBraceToken);
             Debug.assertEqual(closeBraceToken.kind, SyntaxKind.CloseBraceToken);
-            const [leftOfColon, ...rest] = splitChildren(children, child => child.kind === SyntaxKind.ColonToken);
             // Group `-/+readonly` and `-/+?`
-            const leftChildren = groupChildren(getChildrenOrSingleNode(leftOfColon), child =>
+            const groupedWithPlusMinusTokens = groupChildren(children, child =>
                 child === node.readonlyToken || child.kind === SyntaxKind.ReadonlyKeyword ||
                 child === node.questionToken || child.kind === SyntaxKind.QuestionToken);
+            // Group type parameter with surrounding brackets
+            const groupedWithBrackets = groupChildren(groupedWithPlusMinusTokens, ({ kind }) =>
+                kind === SyntaxKind.OpenBracketToken ||
+                kind === SyntaxKind.TypeParameter ||
+                kind === SyntaxKind.CloseBracketToken
+            );
             return [
                 openBraceToken,
-                createSyntaxList([
-                    // Group type parameter with surrounding brackets
-                    createSyntaxList(groupChildren(leftChildren, ({ kind }) =>
-                        kind === SyntaxKind.OpenBracketToken ||
-                        kind === SyntaxKind.TypeParameter ||
-                        kind === SyntaxKind.CloseBracketToken
-                    )),
-                    ...rest,
-                ]),
+                // Pivot on `:`
+                createSyntaxList(splitChildren(groupedWithBrackets, ({ kind }) => kind === SyntaxKind.ColonToken)),
                 closeBraceToken,
             ];
         }
 
-        // Split e.g. `readonly foo?: string` into left and right sides of the colon,
-        // the group `readonly foo` without the QuestionToken.
+        // Group modifiers and property name, then pivot on `:`.
         if (isPropertySignature(node)) {
-            const [leftOfColon, ...rest] = splitChildren(node.getChildren(), child => child.kind === SyntaxKind.ColonToken);
-            return [
-                createSyntaxList(groupChildren(getChildrenOrSingleNode(leftOfColon), child => child !== node.questionToken)),
-                ...rest,
-            ];
+            const children = groupChildren(node.getChildren(), child =>
+                child === node.name || contains(node.modifiers, child));
+            return splitChildren(children, ({ kind }) => kind === SyntaxKind.ColonToken);
+        }
+
+        // Group the parameter name with its `...`, then that group with its `?`, then pivot on `=`.
+        if (isParameter(node)) {
+            const groupedDotDotDotAndName = groupChildren(node.getChildren(), child =>
+                child === node.dotDotDotToken || child === node.name);
+            const groupedWithQuestionToken = groupChildren(groupedDotDotDotAndName, child =>
+                child === groupedDotDotDotAndName[0] || child === node.questionToken);
+            return splitChildren(groupedWithQuestionToken, ({ kind }) => kind === SyntaxKind.EqualsToken);
         }
 
         return node.getChildren();
@@ -194,10 +220,6 @@ namespace ts.SelectionRange {
         return separateLastToken ? result.concat(lastToken) : result;
     }
 
-    function getChildrenOrSingleNode(node: Node): Node[] {
-        return isSyntaxList(node) ? node.getChildren() : [node];
-    }
-
     function createSyntaxList(children: Node[]): SyntaxList {
         Debug.assertGreaterThanOrEqual(children.length, 1);
         const syntaxList = createNode(SyntaxKind.SyntaxList, children[0].pos, last(children).end) as SyntaxList;

src/testRunner/unittests/tsserver/selectionRange.ts

@@ -471,5 +471,75 @@ type M = { -readonly [K in keyof any]-?: any };`);
                     parent: leftOfColonUp },
             });
         });
+
+        it("works for parameters", () => {
+            const getSelectionRange = setup("/file.ts", `
+function f(p, q?, ...r: any[] = []) {}`);
+
+            const locations = getSelectionRange([
+                { line: 2, offset: 12 }, // p
+                { line: 2, offset: 15 }, // q
+                { line: 2, offset: 19 }, // ...
+            ]);
+
+            const allParamsUp: protocol.SelectionRange = {
+                textSpan: { // just inside parens
+                    start: { line: 2, offset: 12 },
+                    end: { line: 2, offset: 35 } },
+                parent: {
+                    textSpan: {
+                        start: { line: 2, offset: 1 },
+                        end: { line: 2, offset: 39 } },
+                    parent: {
+                        textSpan: {
+                            start: { line: 1, offset: 1 },
+                            end: { line: 2, offset: 39 } } } } };
+
+            assert.deepEqual(locations![0], {
+                textSpan: { // p
+                    start: { line: 2, offset: 12 },
+                    end: { line: 2, offset: 13 } },
+                parent: allParamsUp,
+            });
+
+            assert.deepEqual(locations![1], {
+                textSpan: { // q
+                    start: { line: 2, offset: 15 },
+                    end: { line: 2, offset: 16 } },
+                parent: {
+                    textSpan: { // q?
+                        start: { line: 2, offset: 15 },
+                        end: { line: 2, offset: 17 } },
+                    parent: allParamsUp },
+            });
+
+            assert.deepEqual(locations![2], {
+                textSpan: { // ...
+                    start: { line: 2, offset: 19 },
+                    end: { line: 2, offset: 22 } },
+                parent: {
+                    textSpan: { // ...r
+                        start: { line: 2, offset: 19 },
+                        end: { line: 2, offset: 23 } },
+                    parent: {
+                        textSpan: { // ...r: any[]
+                            start: { line: 2, offset: 19 },
+                            end: { line: 2, offset: 30 } },
+                        parent: {
+                            textSpan: { // ...r: any[] = []
+                                start: { line: 2, offset: 19 },
+                                end: { line: 2, offset: 35 } },
+                            parent: allParamsUp } } },
+            });
+        });
+
+        it("snaps to nodes directly behind the cursor instead of trivia ahead of the cursor", () => {
+            const getSelectionRange = setup("/file.ts", `let x: string`);
+            const locations = getSelectionRange([{ line: 1, offset: 4 }]);
+            assert.deepEqual(locations![0].textSpan, {
+                start: { line: 1, offset: 1 },
+                end: { line: 1, offset: 4 },
+            });
+        });
     });
 }