Edit/Chat: Always expand to the nearest enclosing function, if available, before folding ranges

vscode/CHANGELOG.md

@@ -6,6 +6,8 @@ This is a log of all notable changes to Cody for VS Code. [Unreleased] changes a
 
 ### Added
 
+- Edit/Chat: Cody now expands the selection to the nearest enclosing function, if available, before attempting to expand to the nearest enclosing block. [pull/3507](https://github.com/sourcegraph/cody/pull/3507)
+
 ### Fixed
 
 ### Changed

vscode/src/code-actions/fixup.ts

@@ -35,7 +35,7 @@ export class FixupCodeAction implements vscode.CodeActionProvider {
         // const importDiagnostics = diagnostics.filter(diagnostic => diagnostic.message.includes('import'))
 
         // Expand range by getting the folding range contains the target (error) area
-        const targetAreaRange = await getSmartSelection(document.uri, range.start.line)
+        const targetAreaRange = await getSmartSelection(document.uri, range.start)
 
         const newRange = targetAreaRange
             ? new vscode.Range(targetAreaRange.start, targetAreaRange.end)

vscode/src/commands/context/selection.ts

@@ -29,7 +29,7 @@ export async function getContextFileFromCursor(): Promise<ContextItem[]> {
             // Else, use smart selection based on cursor position
             // Else, use visible range of the editor that contains the cursor as fallback
             const cursor = editor.active.selection
-            const smartSelection = await getSmartSelection(document?.uri, cursor?.start.line)
+            const smartSelection = await getSmartSelection(document?.uri, cursor?.start)
             const activeSelection = !cursor?.start.isEqual(cursor?.end) ? cursor : smartSelection
             const visibleRange = editor.active.visibleRanges.find(range => range.contains(cursor?.start))
             const selection = activeSelection ?? visibleRange

vscode/src/edit/utils/edit-selection.ts

@@ -41,15 +41,14 @@ export async function getEditSmartSelection(
     // If we find a new expanded selection position then we set it as the new start position
     // and if we don't then we fallback to the original selection made by the user
     const newSelectionStartingPosition =
-        (await getSmartSelection(document, activeCursorStartPosition.line))?.start ||
-        selectionRange.start
+        (await getSmartSelection(document, activeCursorStartPosition))?.start || selectionRange.start
 
     // Retrieve the ending line of the current selection
     const activeCursorEndPosition = selectionRange.end
     // If we find a new expanded selection position then we set it as the new ending position
     // and if we don't then we fallback to the original selection made by the user
     const newSelectionEndingPosition =
-        (await getSmartSelection(document, activeCursorEndPosition.line))?.end || selectionRange.end
+        (await getSmartSelection(document, activeCursorEndPosition))?.end || selectionRange.end
 
     // Create a new range that starts from the beginning of the folding range at the start position
     // and ends at the end of the folding range at the end position.

vscode/src/editor/utils/index.ts

@@ -1,5 +1,6 @@
 import * as vscode from 'vscode'
 
+import { execQueryWrapper } from '../../tree-sitter/query-sdk'
 import { getSelectionAroundLine } from './document-sections'
 
 /**
@@ -23,13 +24,34 @@ import { getSelectionAroundLine } from './document-sections'
  */
 export async function getSmartSelection(
     documentOrUri: vscode.TextDocument | vscode.Uri,
-    target: number
+    target: vscode.Position
 ): Promise<vscode.Selection | undefined> {
     const document =
         documentOrUri instanceof vscode.Uri
             ? await vscode.workspace.openTextDocument(documentOrUri)
             : documentOrUri
-    return getSelectionAroundLine(document, target)
+
+    const [enclosingFunction] = execQueryWrapper({
+        document,
+        position: target,
+        queryWrapper: 'getEnclosingFunction',
+    })
+
+    if (enclosingFunction) {
+        const { startPosition, endPosition } = enclosingFunction.node
+        // Regardless of the columns provided, we want to ensure the edit spans the full range of characters
+        // on the start and end lines. This helps improve the reliability of the output.
+        const adjustedStartColumn = document.lineAt(startPosition.row).firstNonWhitespaceCharacterIndex
+        const adjustedEndColumn = Number.MAX_SAFE_INTEGER
+        return new vscode.Selection(
+            startPosition.row,
+            adjustedStartColumn,
+            endPosition.row,
+            adjustedEndColumn
+        )
+    }
+
+    return getSelectionAroundLine(document, target.line)
 }
 
 /**

vscode/src/tree-sitter/queries.ts

@@ -8,6 +8,7 @@ export type QueryName =
     | 'intents'
     | 'documentableNodes'
     | 'graphContextIdentifiers'
+    | 'enclosingFunction'
 
 /**
  * Completion intents sorted by priority.

vscode/src/tree-sitter/queries/go.ts

@@ -40,6 +40,12 @@ const DOCUMENTABLE_NODES = dedent`
             name: (field_identifier) @symbol.identifier)  @range.identifier)
 `
 
+const ENCLOSING_FUNCTION_QUERY = dedent`
+    (function_declaration) @range.function
+    (method_declaration) @range.function
+    (func_literal) @range.function
+`
+
 export const goQueries = {
     [SupportedLanguage.go]: {
         singlelineTriggers: SINGLE_LINE_TRIGGERS,
@@ -51,5 +57,6 @@ export const goQueries = {
             (type_spec (type_identifier) @identifier)
             (selector_expression (field_identifier)) @identifier
         `,
+        enclosingFunction: ENCLOSING_FUNCTION_QUERY,
     },
 } satisfies Partial<Record<SupportedLanguage, Record<QueryName, string>>>

vscode/src/tree-sitter/queries/javascript.ts

@@ -216,29 +216,41 @@ const TSX_GRAPH_CONTEXT_IDENTIFIERS_QUERY = dedent`
     (jsx_attribute (property_identifier) @identifier)
 `
 
+const JS_ENCLOSING_FUNCTION_QUERY = dedent`
+    (function_declaration) @range.function
+    (generator_function_declaration) @range.function
+    (function_expression) @range.function
+    (arrow_function) @range.function
+    (method_definition) @range.function
+`
+
 export const javascriptQueries = {
     [SupportedLanguage.javascript]: {
         singlelineTriggers: '',
         intents: JS_INTENTS_QUERY,
         documentableNodes: JS_DOCUMENTABLE_NODES_QUERY,
         graphContextIdentifiers: JS_GRAPH_CONTEXT_IDENTIFIERS_QUERY,
+        enclosingFunction: JS_ENCLOSING_FUNCTION_QUERY,
     },
     [SupportedLanguage.javascriptreact]: {
         singlelineTriggers: '',
         intents: JSX_INTENTS_QUERY,
         documentableNodes: JS_DOCUMENTABLE_NODES_QUERY,
         graphContextIdentifiers: JSX_GRAPH_CONTEXT_IDENTIFIERS_QUERY,
+        enclosingFunction: JS_ENCLOSING_FUNCTION_QUERY,
     },
     [SupportedLanguage.typescript]: {
         singlelineTriggers: TS_SINGLELINE_TRIGGERS_QUERY,
         intents: TS_INTENTS_QUERY,
         documentableNodes: TS_DOCUMENTABLE_NODES_QUERY,
         graphContextIdentifiers: TS_GRAPH_CONTEXT_IDENTIFIERS_QUERY,
+        enclosingFunction: JS_ENCLOSING_FUNCTION_QUERY,
     },
     [SupportedLanguage.typescriptreact]: {
         singlelineTriggers: TS_SINGLELINE_TRIGGERS_QUERY,
         intents: TSX_INTENTS_QUERY,
         documentableNodes: TS_DOCUMENTABLE_NODES_QUERY,
         graphContextIdentifiers: TSX_GRAPH_CONTEXT_IDENTIFIERS_QUERY,
+        enclosingFunction: JS_ENCLOSING_FUNCTION_QUERY,
     },
 } satisfies Partial<Record<SupportedLanguage, Record<QueryName, string>>>

vscode/src/tree-sitter/queries/python.ts

@@ -65,11 +65,16 @@ const DOCUMENTABLE_NODES_QUERY = dedent`
         left: (identifier) @symbol.identifier) @range.identifier
 `
 
+const ENCLOSING_FUNCTION_QUERY = dedent`
+    (function_definition) @range.function
+`
+
 export const pythonQueries = {
     [SupportedLanguage.python]: {
         singlelineTriggers: '',
         intents: INTENTS_QUERY,
         documentableNodes: DOCUMENTABLE_NODES_QUERY,
         graphContextIdentifiers: '(call (identifier) @identifier)',
+        enclosingFunction: ENCLOSING_FUNCTION_QUERY,
     },
 } satisfies Partial<Record<SupportedLanguage, Record<QueryName, string>>>

vscode/src/tree-sitter/query-sdk.ts

@@ -113,6 +113,7 @@ export interface QueryWrappers {
               },
           ]
     getGraphContextIdentifiers: (node: SyntaxNode, start: Point, end?: Point) => QueryCapture[]
+    getEnclosingFunction: (node: SyntaxNode, start: Point, end?: Point) => QueryCapture[]
 }
 
 /**
@@ -217,6 +218,22 @@ function getLanguageSpecificQueryWrappers(
         getGraphContextIdentifiers: (root, start, end) => {
             return queries.graphContextIdentifiers.compiled.captures(root, start, end)
         },
+        getEnclosingFunction: (root, start, end) => {
+            const captures = queries.enclosingFunction.compiled.captures(root, start, end)
+
+            const firstEnclosingFunction = findLast(captures, ({ node }) => {
+                return (
+                    node.startPosition.row <= start.row &&
+                    (start.column <= node.endPosition.column || start.row < node.endPosition.row)
+                )
+            })
+
+            if (!firstEnclosingFunction) {
+                return []
+            }
+
+            return [firstEnclosingFunction]
+        },
     }
 }
 

vscode/src/tree-sitter/query-tests/enclosing-function.test.ts

@@ -0,0 +1,41 @@
+import { describe, it } from 'vitest'
+
+import { initTreeSitterSDK } from '../test-helpers'
+
+import { SupportedLanguage } from '../grammars'
+import { annotateAndMatchSnapshot } from './annotate-and-match-snapshot'
+
+describe('getEnclosingFunction', () => {
+    it('typescript', async () => {
+        const { language, parser, queries } = await initTreeSitterSDK(SupportedLanguage.typescript)
+
+        await annotateAndMatchSnapshot({
+            parser,
+            language,
+            captures: queries.getEnclosingFunction,
+            sourcesPath: 'test-data/enclosing-function.ts',
+        })
+    })
+
+    it('python', async () => {
+        const { language, parser, queries } = await initTreeSitterSDK(SupportedLanguage.python)
+
+        await annotateAndMatchSnapshot({
+            parser,
+            language,
+            captures: queries.getEnclosingFunction,
+            sourcesPath: 'test-data/enclosing-function.py',
+        })
+    })
+
+    it('go', async () => {
+        const { language, parser, queries } = await initTreeSitterSDK(SupportedLanguage.go)
+
+        await annotateAndMatchSnapshot({
+            parser,
+            language,
+            captures: queries.getEnclosingFunction,
+            sourcesPath: 'test-data/enclosing-function.go',
+        })
+    })
+})

vscode/src/tree-sitter/query-tests/test-data/enclosing-function.go

@@ -0,0 +1,49 @@
+func nestedVar() {
+    y := 4
+//  |
+}
+
+// ------------------------------------
+
+func greet() {
+	//  |
+}
+
+// ------------------------------------
+
+func (u User) DisplayName() string {
+    return u.FirstName + " " + u.LastName
+    //           |
+}
+
+// ------------------------------------
+
+func funcFactory(mystring string) func(before, after string) string {
+    return func(before, after string) string {
+        return fmt.Sprintf("%s %s %s", before, mystring, after)
+            //           |
+    }
+}
+
+// ------------------------------------
+
+func funcFactory(mystring string) func(before, after string) string {
+    //  |
+    return func(before, after string) string {
+        return fmt.Sprintf("%s %s %s", before, mystring, after)
+    }
+}
+
+// ------------------------------------
+
+func() {
+    fmt.Println("I'm an anonymous function!")
+    //  |
+}()
+
+// ------------------------------------
+
+var varFunction = func(name string) {
+    fmt.Println("Hello,", name)
+    //  |
+}

vscode/src/tree-sitter/query-tests/test-data/enclosing-function.py

@@ -0,0 +1,40 @@
+def wrapper():
+    print('wrapper')
+    def test():
+        pass
+        # |
+
+# ------------------------------------
+
+def test():
+    pass
+    # |
+
+# ------------------------------------
+
+def test_parameter(val):
+    #              |
+    wrapper()
+
+# ------------------------------------
+
+class Agent:
+    pass
+#   |
+
+# ------------------------------------
+
+class Agent:
+    def __init__(self, name):
+        self.name = name
+    #      |
+
+# ------------------------------------
+
+class Agent:
+    def __init__(self, name):
+        self.name = name
+
+    def test(self):
+        pass
+    #    |