From 579dbafadbf9a1cb20c23f760e32812987fe750f Mon Sep 17 00:00:00 2001 From: Kun Ren Date: Sat, 24 Oct 2026 16:12:34 +0800 Subject: [PATCH 1/3] fix: resolve session pipeline completion from expression context --- src/completions.ts | 54 +++----- src/pipelineCompletion.ts | 262 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 277 insertions(+), 39 deletions(-) create mode 100644 src/pipelineCompletion.ts diff --git a/src/completions.ts b/src/completions.ts index ac5ec8c8f..b8a63ec28 100644 --- a/src/completions.ts +++ b/src/completions.ts @@ -2,8 +2,7 @@ import * as vscode from 'vscode'; import * as session from './session'; -import { extendSelection } from './selection'; -import { cleanLine } from './lineCache'; +import { getPipelineSource } from './pipelineCompletion'; import { globalRHelp } from './extension'; import { config } from './util'; import { getChunks } from './rmarkdown'; @@ -225,6 +224,7 @@ export class LiveCompletionItemProvider implements vscode.CompletionItemProvider return items; } + let codeStartLine = 0; if (document.languageId === 'rmd') { const chunks = getChunks(document); const chunk = chunks.find( @@ -236,6 +236,7 @@ export class LiveCompletionItemProvider implements vscode.CompletionItemProvider if (!chunk) { return items; } + codeStartLine = chunk.startLine + 1; } const trigger = completionContext.triggerCharacter; @@ -313,7 +314,15 @@ export class LiveCompletionItemProvider implements vscode.CompletionItemProvider } if (trigger === undefined || trigger === '(' || trigger === ',') { - items.push(...getPipelineCompletionItems(document, position, token, activeSession)); + items.push( + ...getPipelineCompletionItems( + document, + position, + token, + activeSession, + codeStartLine, + ), + ); } return items; @@ -429,44 +438,11 @@ function getPipelineCompletionItems( position: vscode.Position, token: vscode.CancellationToken, activeSession: session.Session, + startLine: number, ): vscode.CompletionItem[] { const items: vscode.CompletionItem[] = []; - const range = extendSelection( - position.line, - (x) => document.lineAt(x).text, - document.lineCount, - ); - let symbol: string | undefined = undefined; - - for (let i = range.startLine; i <= range.endLine; i++) { - if (token.isCancellationRequested) { - break; - } - - const line = document.lineAt(i); - if (line.isEmptyOrWhitespace) { - continue; - } - - const cleanedLine = cleanLine(line.text); - if (cleanedLine.length === 0) { - continue; - } - - const pipeSymbolIndex = line.text.search(/([\w_.]+)\s*(%.+%|\|>)/); - if (pipeSymbolIndex < 0) { - break; - } - - const symbolPosition = new vscode.Position(i, pipeSymbolIndex); - const symbolRange = document.getWordRangeAtPosition(symbolPosition); - - if (symbolRange !== undefined) { - symbol = document.getText(symbolRange); - } - - break; - } + const text = document.getText(new vscode.Range(new vscode.Position(startLine, 0), position)); + const symbol = getPipelineSource(text); if (!token.isCancellationRequested && symbol !== undefined) { const obj = activeSession.workspaceData.globalenv[symbol]; diff --git a/src/pipelineCompletion.ts b/src/pipelineCompletion.ts new file mode 100644 index 000000000..eb52884ad --- /dev/null +++ b/src/pipelineCompletion.ts @@ -0,0 +1,262 @@ +interface Token { + text: string; + symbol?: string; +} + +interface Expression { + source?: string; + // null denotes a pipeline whose source cannot be resolved. It must not + // inherit suggestions from a surrounding pipeline. + pipeline?: string | null; +} + +const pipes = new Set(['|>', '%>%', '%<>%', '%T>%', '%$%']); +const precedence = new Map([ + ['?', 1], + ['=', 2], + ['<-', 3], + ['<<-', 3], + ['->', 4], + ['->>', 4], + ['~', 5], + ['||', 6], + ['|', 6], + ['&&', 7], + ['&', 7], + ['<', 9], + ['>', 9], + ['<=', 9], + ['>=', 9], + ['==', 9], + ['!=', 9], + ['+', 10], + ['-', 10], + ['*', 11], + ['/', 11], + ['|>', 12], + [':', 13], + ['^', 15], + ['$', 16], + ['@', 16], + ['::', 17], + [':::', 17], +]); + +function tokenize(text: string): Token[] | undefined { + const tokens: Token[] = []; + for (let i = 0; i < text.length;) { + const rest = text.slice(i); + const char = text[i]; + if (char === '\n') { + tokens.push({ text: '\n' }); + i++; + } else if (/\s/.test(char)) { + i++; + } else if (char === '#') { + const end = text.indexOf('\n', i); + if (end < 0) { + return undefined; + } + i = end; + } else { + const raw = /^[rR](["'])(-*)([([{])/.exec(rest); + if (raw) { + const close = { '(': ')', '[': ']', '{': '}' }[raw[3]]; + const ending = close + raw[2] + raw[1]; + const end = text.indexOf(ending, i + raw[0].length); + if (end < 0) { + return undefined; + } + tokens.push({ text: 'literal' }); + i = end + ending.length; + } else if (char === '"' || char === "'" || char === '`') { + let end = i + 1; + let symbol = ''; + while (end < text.length && text[end] !== char) { + if (text[end] === '\\') { + end++; + const escaped = text[end]; + const escapes: Record = { + a: '\x07', + b: '\b', + f: '\f', + n: '\n', + r: '\r', + t: '\t', + v: '\v', + }; + const encoded = + /^(?:x([\da-fA-F]{1,2})|u([\da-fA-F]{4})|U([\da-fA-F]{8})|u\{([\da-fA-F]{1,4})\}|U\{([\da-fA-F]{1,8})\}|([0-7]{1,3}))/.exec( + text.slice(end), + ); + if (encoded) { + const digits = encoded.slice(1).find((value) => value !== undefined)!; + const value = parseInt(digits, encoded[6] ? 8 : 16); + if (value === 0 || value > 0x10ffff) { + return undefined; + } + symbol += String.fromCodePoint(value); + end += encoded[0].length - 1; + } else if (escaped in escapes || ['\\', '`', '"', "'"].includes(escaped)) { + symbol += escapes[escaped] ?? escaped; + } else { + return undefined; + } + } else { + symbol += text[end]; + } + end++; + } + if (end >= text.length) { + return undefined; + } + tokens.push(char === '`' ? { text: 'symbol', symbol } : { text: 'literal' }); + i = end + 1; + } else { + const number = /^(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?[Li]?/.exec(rest); + const name = /^[\p{L}._][\p{L}\p{N}._]*/u.exec(rest); + const operator = + /^(?:%[^%\n]*%|:::|::|<<-|->>|<-|->|\|>|\|\||&&|<=|>=|==|!=|\[\[)/.exec(rest); + if (number) { + tokens.push({ text: 'literal' }); + i += number[0].length; + } else if (name) { + tokens.push({ text: 'symbol', symbol: name[0] }); + i += name[0].length; + } else if (operator) { + tokens.push({ text: operator[0] }); + i += operator[0].length; + } else { + tokens.push({ text: char }); + i++; + } + } + } + } + return tokens; +} + +/** + * Recover a pipeline's initial session symbol from the code before the cursor. + * Only expression structure is inspected: calls, members and computed values + * are never evaluated. Incomplete calls and grouping are expected while typing. + */ +export function getPipelineSource(text: string): string | undefined { + const tokens = tokenize(text); + if (!tokens) { + return undefined; + } + let index = 0; + let depth = 0; + const current = () => tokens[index]?.text; + const skipNewlines = () => { + while (current() === '\n') { + index++; + } + }; + const binding = (operator: string | undefined) => + operator?.startsWith('%') ? 12 : precedence.get(operator ?? ''); + + function sequence(close?: string, continuation = false): Expression { + let result: Expression = {}; + while (index < tokens!.length && current() !== close) { + if (current() === '\n' || current() === ';' || current() === ',') { + index++; + result = {}; + continue; + } + const start = index; + result = expression(0, continuation); + if (index === start) { + index++; + } + } + return result; + } + + function expression(minimum: number, continuation: boolean): Expression { + // Bound recovery for malformed or deeply nested editor input. + if (++depth > 128) { + index = tokens!.length; + depth--; + return { pipeline: null }; + } + skipNewlines(); + const token = tokens![index++]; + let left: Expression = {}; + if (token?.symbol !== undefined) { + left.source = token.symbol; + } else if (token?.text === '(' || token?.text === '{') { + const close = token.text === '(' ? ')' : '}'; + left = sequence(close, token.text === '('); + if (current() === close) { + index++; + left.pipeline = undefined; + } + if (token.text === '{') { + left.source = undefined; + } + } else if (token && ['+', '-', '!', '~', '?'].includes(token.text)) { + left = expression(token.text === '!' ? 8 : token.text === '~' ? 5 : 14, continuation); + left.source = undefined; + } + + while (index < tokens!.length) { + if (current() === '\n' && continuation) { + let next = index; + while (tokens![next]?.text === '\n') { + next++; + } + if (binding(tokens![next]?.text) !== undefined) { + index = next; + } + } + const operator = current(); + if (operator === '(' || operator === '[' || operator === '[[') { + if (18 < minimum) { + break; + } + index++; + const close = operator === '(' ? ')' : ']'; + const args = sequence(close, true); + if (current() === close) { + index++; + if (operator === '[[' && current() === ']') { + index++; + } + left = {}; + } else { + left = { pipeline: args.pipeline }; + } + continue; + } + const power = binding(operator); + if (power === undefined || power < minimum) { + break; + } + index++; + const right = expression( + power + (['^', '<-', '<<-', '='].includes(operator!) ? 0 : 1), + continuation, + ); + if (pipes.has(operator!)) { + // A nested magrittr pipeline can start from the surrounding + // pipeline's dot placeholder rather than another session object. + const pipeline = + right.pipeline === '.' && ['%>%', '%<>%', '%T>%'].includes(operator!) + ? (left.source ?? null) + : right.pipeline; + left = { + source: left.source, + pipeline: pipeline !== undefined ? pipeline : (left.source ?? null), + }; + } else { + left = { pipeline: right.pipeline }; + } + } + depth--; + return left; + } + + return sequence().pipeline ?? undefined; +} From 74def8a971ec5e816b1bedbab1778ae88b591b7b Mon Sep 17 00:00:00 2001 From: Kun Ren Date: Sat, 24 Oct 2026 16:12:34 +0800 Subject: [PATCH 2/3] test: cover grouped and nested session pipelines --- src/test/node/pipelineCompletion.test.ts | 134 ++++++++++++++++++ src/test/suite/completions.test.ts | 164 +++++++++++++++++++++++ 2 files changed, 298 insertions(+) create mode 100644 src/test/node/pipelineCompletion.test.ts diff --git a/src/test/node/pipelineCompletion.test.ts b/src/test/node/pipelineCompletion.test.ts new file mode 100644 index 000000000..f2af530f5 --- /dev/null +++ b/src/test/node/pipelineCompletion.test.ts @@ -0,0 +1,134 @@ +import * as assert from 'assert'; +import { getPipelineSource } from '../../pipelineCompletion'; + +suite('Session pipeline source', () => { + const source = (code: string) => getPipelineSource(code.split('')[0]); + + for (const pipe of ['%>%', '|>', '%<>%', '%T>%', '%$%']) { + for (const code of [ + `raw ${pipe} select()`, + `(raw ${pipe} select())`, + `(\n raw ${pipe} select()\n)`, + `(\n (\n raw ${pipe} select()\n )\n)`, + `(\n raw\n ${pipe} select()\n)`, + `((raw)) ${pipe}\n select()`, + `(raw ${pipe} head()) ${pipe} select()`, + `result <- (\n raw ${pipe} select()\n)`, + `(\n raw ${pipe} select()\n) -> result`, + `print(\n raw ${pipe} select()\n)`, + `raw ${pipe} select(a, )`, + `raw ${pipe} mutate(x = sum())`, + ]) { + test(`finds the source in ${JSON.stringify(code)}`, () => { + assert.strictEqual(source(code), 'raw'); + }); + } + } + + const nested = [ + 'raw %>% mutate(x = other %>% select())', + 'list(raw %>% select(a), other %>% select())', + 'raw %>% mutate(x = (other |> select()))', + 'raw %>% mutate(x = other %>% select(a), y = other2 |> select())', + ]; + for (const code of nested) { + test(`selects the nearest pipeline in ${JSON.stringify(code)}`, () => { + assert.strictEqual(source(code), code.includes('other2') ? 'other2' : 'other'); + }); + } + + for (const code of [ + 'raw %>% mutate(x = other %>% select(a), y = sum())', + 'raw %>% mutate(x = other %>% select(a), )', + 'raw %>% mutate(x = list(other %>% select(a), sum()))', + 'raw %>% mutate(x = { other %>% select(a); sum() })', + 'raw %>% { . %>% select() }', + 'raw %<>% mutate(x = . %>% select())', + 'raw %T>% { . %>% select() }', + ]) { + test(`returns to the surrounding pipeline in ${JSON.stringify(code)}`, () => { + assert.strictEqual(source(code), 'raw'); + }); + } + + test('binds nested dot placeholders to their nearest magrittr input', () => { + assert.strictEqual( + source('raw %>% mutate(x = other %>% { . %>% select() })'), + 'other', + ); + assert.strictEqual( + source('raw %>% mutate(x = identity(other) %>% { . %>% select() })'), + undefined, + ); + }); + + for (const code of [ + '( # grouping\n # another comment\n raw %>% # pipe\n select())', + 'raw %>% mutate(note = "other %>% select(", )', + 'raw %>% mutate(note = "escaped \\\" other %>% select(", )', + "raw %>% mutate(note = 'other %>% select(', )", + 'raw %>% mutate(note = r"(other %>% select(")", )', + 'raw %>% mutate(note = R"--[other %>% select("]--", )', + ]) { + test(`ignores comments and string contents in ${JSON.stringify(code)}`, () => { + assert.strictEqual(source(code), 'raw'); + }); + } + + for (const [name, spelling] of [ + ['raw data', '`raw data`'], + ['数据', '数据'], + ['a`b', '`a\\`b`'], + ['a\\b', '`a\\\\b`'], + ['a\nb', '`a\\nb`'], + ['ab', '`a\\x62`'], + ['ab', '`a\\u0062`'], + ['ab', '`a\\U00000062`'], + ['ab', '`a\\u{62}`'], + ['ab', '`a\\U{62}`'], + ['ab', '`a\\142`'], + ]) { + test(`recognizes session symbol ${JSON.stringify(name)}`, () => { + assert.strictEqual(source(`${spelling} %>% select()`), name); + }); + } + + for (const code of [ + 'raw %>% select(a); select()', + 'raw %>% select(a)\nselect()', + 'list(raw %>% select(a), select())', + 'raw %>% select(a) + select()', + 'raw %>% select(a) %in% select()', + 'identity(raw) %>% select()', + 'raw$other %>% select()', + 'raw[[1]] %>% select()', + '(raw + other) %>% select()', + 'pkg::raw %>% select()', + 'raw %>% mutate(x = identity(other) %>% select())', + 'raw %>% mutate(x = (other + other2) %>% select())', + 'raw %>% mutate(note = "other %>% select()")', + 'raw %>% mutate(note = r"(other %>% select())")', + 'raw %>% mutate( # other %>% select()\n)', + 'raw %in% select()', + 'raw %*% select()', + 'select(); raw %>% select(a)', + '`a\\q` %>% select()', + '`a\\UFFFFFFFF` %>% select()', + ]) { + test(`does not borrow unrelated session names in ${JSON.stringify(code)}`, () => { + assert.strictEqual(source(code), undefined); + }); + } + + test('bounds recovery of deeply nested incomplete input', () => { + assert.strictEqual(getPipelineSource('('.repeat(1000) + 'raw %>% select('), undefined); + }); + + test('ignores earlier statements and completed function bodies', () => { + assert.strictEqual( + source('f <- function(x) { other %>% select(a) }\n(raw\n |> select())'), + 'raw', + ); + assert.strictEqual(source('other %>% select(a)\nselect()'), undefined); + }); +}); diff --git a/src/test/suite/completions.test.ts b/src/test/suite/completions.test.ts index 5b86ad516..aa2168db1 100644 --- a/src/test/suite/completions.test.ts +++ b/src/test/suite/completions.test.ts @@ -64,3 +64,167 @@ suite('Session hover and completion', () => { sinon.assert.notCalled(sessionRequest); }); }); + +suite('Session pipeline completion', () => { + let sandbox: sinon.SinonSandbox; + let cancellation: vscode.CancellationTokenSource; + let request: sinon.SinonStub; + + setup(() => { + sandbox = sinon.createSandbox(); + cancellation = new vscode.CancellationTokenSource(); + const object = (names: string[]) => ({ + class: ['data.frame'], + type: 'list', + length: names.length, + str: 'data.frame', + names, + }); + const workspaceData: session.WorkspaceData = { + search: [], + loaded_namespaces: [], + globalenv: { + raw: object(['sample', 'umi1']), + other: object(['other_column']), + 'raw data': object(['quoted_column']), + 数据: object(['unicode_column']), + }, + }; + sandbox.stub(session, 'sessionForDocument').returns({ workspaceData } as session.Session); + request = sandbox.stub(session, 'sessionRequest'); + }); + + teardown(() => { + cancellation.dispose(); + sandbox.restore(); + }); + + const cases: [string, string[]][] = [ + ['raw %>% select()', ['sample', 'umi1']], + ['(\n raw %>% select()\n)', ['sample', 'umi1']], + ['(\n (\n raw\n |> select()\n )\n)', ['sample', 'umi1']], + ['result <- (\n raw %<>% select()\n)', ['sample', 'umi1']], + ['print(raw %T>% select())', ['sample', 'umi1']], + ['raw %>% select(a, )', ['sample', 'umi1']], + ['raw %>% mutate(x = other |> select())', ['other_column']], + ['raw %>% { . %>% select() }', ['sample', 'umi1']], + ['raw %>% mutate(x = other %>% { . %>% select() })', ['other_column']], + ['list(raw %>% select(a), other %>% select())', ['other_column']], + ['raw %>% mutate(x = other %>% select(a), y = sum())', ['sample', 'umi1']], + ['`raw data` %>% select()', ['quoted_column']], + ['数据 |> select()', ['unicode_column']], + ['raw %>% select(a); select()', []], + ['list(raw %>% select(a), select())', []], + ['raw %>% select(a) + select()', []], + ['identity(raw) %>% select()', []], + ['raw %>% mutate(x = identity(other) %>% select())', []], + ['raw %>% mutate(note = "other %>% select()")', []], + ['raw %>% mutate( # comment \n)', []], + ['select(); raw %>% select(a)', []], + [ + 'f <- function(x) { other %>% select(a) }\n(raw\n |> select())', + ['sample', 'umi1'], + ], + ]; + + for (const [content, expected] of cases) { + test(`completes ${JSON.stringify(content)}`, async () => { + const offset = content.indexOf(''); + const document = await vscode.workspace.openTextDocument({ + language: 'r', + content: content.replace('', ''), + }); + const before = content.slice(0, offset).trimEnd(); + const trigger = before.endsWith(',') ? ',' : '('; + const items = await new LiveCompletionItemProvider().provideCompletionItems( + document, + document.positionAt(offset), + cancellation.token, + { + triggerKind: vscode.CompletionTriggerKind.TriggerCharacter, + triggerCharacter: trigger, + }, + ); + assert.deepStrictEqual( + items.map((item) => item.label), + expected, + ); + assert.ok(items.every((item) => item.detail === '[session]')); + sinon.assert.notCalled(request); + }); + } + + for (const [code, expected] of [ + ['(\n raw\n |> select()\n)', ['sample', 'umi1']], + ['select()', []], + ] as [string, string[]][]) { + test(`completes within the current R Markdown chunk: ${JSON.stringify(code)}`, async () => { + const content = + 'Text with an unmatched " quote.\n\n```{r}\nother %>% select(a)\n```\n\n```{r}\n' + + code + + '\n```'; + const offset = content.indexOf(''); + const document = await vscode.workspace.openTextDocument({ + language: 'rmd', + content: content.replace('', ''), + }); + const items = await new LiveCompletionItemProvider().provideCompletionItems( + document, + document.positionAt(offset), + cancellation.token, + { + triggerKind: vscode.CompletionTriggerKind.TriggerCharacter, + triggerCharacter: '(', + }, + ); + assert.deepStrictEqual( + items.map((item) => item.label), + expected, + ); + sinon.assert.notCalled(request); + }); + } + + test('VS Code completion requests include the pipeline columns', async () => { + const content = '(\n raw\n |> select(\n)\n)'; + const document = await vscode.workspace.openTextDocument({ language: 'r', content }); + const provider = vscode.languages.registerCompletionItemProvider( + { language: 'r', scheme: 'untitled' }, + new LiveCompletionItemProvider(), + '(', + ); + try { + const result = await vscode.commands.executeCommand( + 'vscode.executeCompletionItemProvider', + document.uri, + document.positionAt(content.indexOf('select(') + 'select('.length), + '(', + ); + assert.ok(result); + const columns = result.items.filter((item) => item.detail === '[session]'); + assert.deepStrictEqual( + columns.map((item) => item.label), + ['sample', 'umi1'], + ); + sinon.assert.notCalled(request); + } finally { + provider.dispose(); + } + }); + + test('cancellation suppresses pipeline completion', async () => { + cancellation.cancel(); + const document = await vscode.workspace.openTextDocument({ + language: 'r', + content: 'raw |> select(', + }); + const items = await new LiveCompletionItemProvider().provideCompletionItems( + document, + document.positionAt(document.getText().length), + cancellation.token, + { triggerKind: vscode.CompletionTriggerKind.TriggerCharacter, triggerCharacter: '(' }, + ); + assert.deepStrictEqual(items, []); + sinon.assert.notCalled(request); + }); +}); From bddd864386baa52035bf2f90c99d98090df3034e Mon Sep 17 00:00:00 2001 From: Kun Ren Date: Sat, 24 Oct 2026 17:33:55 +0800 Subject: [PATCH 3/3] perf: cache pipeline context and yield during cold completion scans --- scripts/benchmark-pipeline-completion.cjs | 66 ++++++++ src/completions.ts | 32 +++- src/pipelineCompletion.ts | 184 +++++++++++++++++++--- src/test/node/pipelineCompletion.test.ts | 92 ++++++++++- src/test/suite/completions.test.ts | 41 +++++ 5 files changed, 384 insertions(+), 31 deletions(-) create mode 100644 scripts/benchmark-pipeline-completion.cjs diff --git a/scripts/benchmark-pipeline-completion.cjs b/scripts/benchmark-pipeline-completion.cjs new file mode 100644 index 000000000..8f3e9f1de --- /dev/null +++ b/scripts/benchmark-pipeline-completion.cjs @@ -0,0 +1,66 @@ +// Run after compiling TypeScript: node scripts/benchmark-pipeline-completion.cjs +// Optional argument: absolute path to a compiled earlier resolver for comparison. +const assert = require('node:assert/strict'); +const { performance } = require('node:perf_hooks'); +const implementation = require(process.argv[2] || '../out/pipelineCompletion'); +const createResolver = () => implementation.PipelineSourceCache + ? new implementation.PipelineSourceCache() + : { resolve: implementation.getPipelineSource }; + +async function measure(run) { + const samples = []; + let maxBlock = 0; + for (let i = 0; i < 9; i++) { + const start = performance.now(); + let previous = start; + let heartbeat; + const beat = () => { + const now = performance.now(); + maxBlock = Math.max(maxBlock, now - previous); + previous = now; + heartbeat = setImmediate(beat); + }; + heartbeat = setImmediate(beat); + try { + await run(i); + maxBlock = Math.max(maxBlock, performance.now() - previous); + } finally { + clearImmediate(heartbeat); + } + samples.push(performance.now() - start); + } + samples.sort((a, b) => a - b); + return { + median_ms: +samples[4].toFixed(3), + max_ms: +samples[8].toFixed(3), + max_block_ms: +maxBlock.toFixed(3), + }; +} + +async function main() { + for (const lines of [1000, 10000, 50000]) { + // An earlier pipeline prevents the benchmark from relying on a global + // string search to skip ordinary completion requests. + const prefix = 'other |> head()\n' + 'value <- mean(1:10)\n'.repeat(lines); + for (const tail of ['select(', '(raw\n |> select(']) { + const expected = tail.includes('raw') ? 'raw' : undefined; + const text = prefix + tail; + const cold = await measure(async () => + assert.equal(await createResolver().resolve(text), expected)); + const cache = createResolver(); + await cache.resolve(text); + const warm = await measure(async () => + assert.equal(await cache.resolve(text), expected)); + const edited = await measure(async (i) => + assert.equal(await cache.resolve(text + 'column_' + i), expected)); + console.log(JSON.stringify({ + lines, pipeline: expected !== undefined, cold, warm, edited, + })); + } + } +} + +main().catch((error) => { + console.error(error); + process.exitCode = 1; +}); diff --git a/src/completions.ts b/src/completions.ts index b8a63ec28..0a394f053 100644 --- a/src/completions.ts +++ b/src/completions.ts @@ -2,7 +2,7 @@ import * as vscode from 'vscode'; import * as session from './session'; -import { getPipelineSource } from './pipelineCompletion'; +import { PipelineSourceCache } from './pipelineCompletion'; import { globalRHelp } from './extension'; import { config } from './util'; import { getChunks } from './rmarkdown'; @@ -212,6 +212,8 @@ export class StaticCompletionItemProvider implements vscode.CompletionItemProvid } export class LiveCompletionItemProvider implements vscode.CompletionItemProvider { + private readonly pipelineSources = new WeakMap(); + async provideCompletionItems( document: vscode.TextDocument, position: vscode.Position, @@ -314,18 +316,24 @@ export class LiveCompletionItemProvider implements vscode.CompletionItemProvider } if (trigger === undefined || trigger === '(' || trigger === ',') { + let cache = this.pipelineSources.get(document); + if (!cache) { + cache = new PipelineSourceCache(); + this.pipelineSources.set(document, cache); + } items.push( - ...getPipelineCompletionItems( + ...(await getPipelineCompletionItems( document, position, token, activeSession, codeStartLine, - ), + cache, + )), ); } - return items; + return token.isCancellationRequested ? [] : items; } } @@ -433,18 +441,26 @@ function getBracketCompletionItems( return items; } -function getPipelineCompletionItems( +async function getPipelineCompletionItems( document: vscode.TextDocument, position: vscode.Position, token: vscode.CancellationToken, activeSession: session.Session, startLine: number, -): vscode.CompletionItem[] { + cache: PipelineSourceCache, +): Promise { const items: vscode.CompletionItem[] = []; + if (token.isCancellationRequested) { + return items; + } + const version = document.version; const text = document.getText(new vscode.Range(new vscode.Position(startLine, 0), position)); - const symbol = getPipelineSource(text); + const symbol = await cache.resolve( + text, + () => token.isCancellationRequested || document.version !== version, + ); - if (!token.isCancellationRequested && symbol !== undefined) { + if (!token.isCancellationRequested && document.version === version && symbol !== undefined) { const obj = activeSession.workspaceData.globalenv[symbol]; if (obj !== undefined && obj.names !== undefined) { const doc = new vscode.MarkdownString('Element of `' + symbol + '`'); diff --git a/src/pipelineCompletion.ts b/src/pipelineCompletion.ts index eb52884ad..50364c9de 100644 --- a/src/pipelineCompletion.ts +++ b/src/pipelineCompletion.ts @@ -10,6 +10,27 @@ interface Expression { pipeline?: string | null; } +interface ScanContext { + offset: number; + tokens: Token[]; + brackets: { close: string; contentStart: number }[]; + previous?: string; +} + +interface ScanResult { + tokens?: Token[]; + checkpoint: ScanContext; +} + +const scanQuantum = 8192; +const maxExpressionTokens = 4096; +const emptyContext = (): ScanContext => ({ offset: 0, tokens: [], brackets: [] }); +const copyContext = (context: ScanContext): ScanContext => ({ + ...context, + tokens: [...context.tokens], + brackets: context.brackets.map((bracket) => ({ ...bracket })), +}); + const pipes = new Set(['|>', '%>%', '%<>%', '%T>%', '%$%']); const precedence = new Map([ ['?', 1], @@ -42,9 +63,29 @@ const precedence = new Map([ [':::', 17], ]); -function tokenize(text: string): Token[] | undefined { - const tokens: Token[] = []; - for (let i = 0; i < text.length;) { +function binding(operator: string | undefined): number | undefined { + return operator?.startsWith('%') ? 12 : precedence.get(operator ?? ''); +} + +function* tokenize( + text: string, + initial = emptyContext(), + canceled = () => false, +): Generator { + const context = copyContext(initial); + const tokens = context.tokens; + let checkpoint = initial; + let nextYield = initial.offset + scanQuantum; + const unknown = () => ({ checkpoint }); + for (let i = initial.offset; i < text.length;) { + if (i >= nextYield) { + yield; + nextYield = i + scanQuantum; + } + if (canceled()) { + return undefined; + } + const tokenStart = tokens.length; const rest = text.slice(i); const char = text[i]; if (char === '\n') { @@ -55,7 +96,7 @@ function tokenize(text: string): Token[] | undefined { } else if (char === '#') { const end = text.indexOf('\n', i); if (end < 0) { - return undefined; + return unknown(); } i = end; } else { @@ -65,7 +106,7 @@ function tokenize(text: string): Token[] | undefined { const ending = close + raw[2] + raw[1]; const end = text.indexOf(ending, i + raw[0].length); if (end < 0) { - return undefined; + return unknown(); } tokens.push({ text: 'literal' }); i = end + ending.length; @@ -73,6 +114,13 @@ function tokenize(text: string): Token[] | undefined { let end = i + 1; let symbol = ''; while (end < text.length && text[end] !== char) { + if (end >= nextYield) { + yield; + nextYield = end + scanQuantum; + if (canceled()) { + return undefined; + } + } if (text[end] === '\\') { end++; const escaped = text[end]; @@ -93,14 +141,14 @@ function tokenize(text: string): Token[] | undefined { const digits = encoded.slice(1).find((value) => value !== undefined)!; const value = parseInt(digits, encoded[6] ? 8 : 16); if (value === 0 || value > 0x10ffff) { - return undefined; + return unknown(); } symbol += String.fromCodePoint(value); end += encoded[0].length - 1; } else if (escaped in escapes || ['\\', '`', '"', "'"].includes(escaped)) { symbol += escapes[escaped] ?? escaped; } else { - return undefined; + return unknown(); } } else { symbol += text[end]; @@ -108,7 +156,7 @@ function tokenize(text: string): Token[] | undefined { end++; } if (end >= text.length) { - return undefined; + return unknown(); } tokens.push(char === '`' ? { text: 'symbol', symbol } : { text: 'literal' }); i = end + 1; @@ -132,8 +180,43 @@ function tokenize(text: string): Token[] | undefined { } } } + if (tokens.length === tokenStart) { + continue; + } + const token = tokens[tokens.length - 1].text; + const close = { '(': ')', '[': ']', '[[': ']]', '{': '}' }[token]; + if (close) { + context.brackets.push({ close, contentStart: tokens.length }); + } else if ([')', ']', '}'].includes(token)) { + const bracket = context.brackets.at(-1); + if (token === ']' && bracket?.close === ']]') { + bracket.close = ']'; + } else if (bracket?.close === token) { + context.brackets.pop(); + } + } + const bracket = context.brackets.at(-1); + // Completed statements and earlier call arguments cannot supply the + // active pipeline. Retain their enclosing expression, not their contents. + const boundary = + token === ';' || + token === ',' || + (token === '\n' && + binding(context.previous) === undefined && + (!bracket || bracket.close === '}')); + if (token !== '\n') { + context.previous = token; + } + if (boundary) { + tokens.length = bracket?.contentStart ?? 0; + context.offset = i; + checkpoint = copyContext(context); + } + if (context.brackets.length > 128) { + return unknown(); + } } - return tokens; + return { tokens, checkpoint }; } /** @@ -141,9 +224,24 @@ function tokenize(text: string): Token[] | undefined { * Only expression structure is inspected: calls, members and computed values * are never evaluated. Incomplete calls and grouping are expected while typing. */ -export function getPipelineSource(text: string): string | undefined { - const tokens = tokenize(text); - if (!tokens) { +export function getPipelineSource(text: string, canceled = () => false): string | undefined { + if (canceled()) { + return undefined; + } + const scan = tokenize(text, emptyContext(), canceled); + let result = scan.next(); + while (!result.done) { + result = scan.next(); + } + return result.value?.tokens ? resolveTokens(result.value.tokens, canceled) : undefined; +} + +function resolveTokens(tokens: Token[], canceled: () => boolean): string | undefined { + if ( + canceled() || + tokens.length > maxExpressionTokens || + !tokens.some((token) => pipes.has(token.text)) + ) { return undefined; } let index = 0; @@ -154,12 +252,13 @@ export function getPipelineSource(text: string): string | undefined { index++; } }; - const binding = (operator: string | undefined) => - operator?.startsWith('%') ? 12 : precedence.get(operator ?? ''); function sequence(close?: string, continuation = false): Expression { let result: Expression = {}; - while (index < tokens!.length && current() !== close) { + while (index < tokens.length && current() !== close) { + if (canceled()) { + return { pipeline: null }; + } if (current() === '\n' || current() === ';' || current() === ',') { index++; result = {}; @@ -176,13 +275,13 @@ export function getPipelineSource(text: string): string | undefined { function expression(minimum: number, continuation: boolean): Expression { // Bound recovery for malformed or deeply nested editor input. - if (++depth > 128) { - index = tokens!.length; + if (++depth > 128 || canceled()) { + index = tokens.length; depth--; return { pipeline: null }; } skipNewlines(); - const token = tokens![index++]; + const token = tokens[index++]; let left: Expression = {}; if (token?.symbol !== undefined) { left.source = token.symbol; @@ -201,13 +300,16 @@ export function getPipelineSource(text: string): string | undefined { left.source = undefined; } - while (index < tokens!.length) { + while (index < tokens.length) { + if (canceled()) { + return { pipeline: null }; + } if (current() === '\n' && continuation) { let next = index; - while (tokens![next]?.text === '\n') { + while (tokens[next]?.text === '\n') { next++; } - if (binding(tokens![next]?.text) !== undefined) { + if (binding(tokens[next]?.text) !== undefined) { index = next; } } @@ -258,5 +360,43 @@ export function getPipelineSource(text: string): string | undefined { return left; } - return sequence().pipeline ?? undefined; + const source = sequence().pipeline; + return canceled() ? undefined : (source ?? undefined); +} + +/** + * Cache only an unchanged lexical prefix and its enclosing pipeline context. + * Edits before that prefix or cursor moves into it restart the scan. The cold + * scan yields between chunks; source resolution is bounded to 4096 retained + * tokens and deliberately does not perform general R analysis or evaluation. + */ +export class PipelineSourceCache { + private prefix = ''; + private checkpoint = emptyContext(); + private request = 0; + + async resolve(text: string, canceled = () => false): Promise { + const request = ++this.request; + if (canceled()) { + return undefined; + } + const initial = text.startsWith(this.prefix) ? this.checkpoint : emptyContext(); + const scan = tokenize(text, initial, canceled); + let result = scan.next(); + while (!result.done) { + await new Promise((resolve) => setImmediate(resolve)); + if (canceled()) { + return undefined; + } + result = scan.next(); + } + if (!result.value || canceled()) { + return undefined; + } + if (request === this.request) { + this.checkpoint = result.value.checkpoint; + this.prefix = text.slice(0, this.checkpoint.offset); + } + return result.value.tokens ? resolveTokens(result.value.tokens, canceled) : undefined; + } } diff --git a/src/test/node/pipelineCompletion.test.ts b/src/test/node/pipelineCompletion.test.ts index f2af530f5..fdddd600b 100644 --- a/src/test/node/pipelineCompletion.test.ts +++ b/src/test/node/pipelineCompletion.test.ts @@ -1,5 +1,5 @@ import * as assert from 'assert'; -import { getPipelineSource } from '../../pipelineCompletion'; +import { getPipelineSource, PipelineSourceCache } from '../../pipelineCompletion'; suite('Session pipeline source', () => { const source = (code: string) => getPipelineSource(code.split('')[0]); @@ -12,6 +12,7 @@ suite('Session pipeline source', () => { `(\n (\n raw ${pipe} select()\n )\n)`, `(\n raw\n ${pipe} select()\n)`, `((raw)) ${pipe}\n select()`, + `raw ${pipe}\n\n # comment\n select()`, `(raw ${pipe} head()) ${pipe} select()`, `result <- (\n raw ${pipe} select()\n)`, `(\n raw ${pipe} select()\n) -> result`, @@ -132,3 +133,92 @@ suite('Session pipeline source', () => { assert.strictEqual(source('other %>% select(a)\nselect()'), undefined); }); }); + +suite('Cached session pipeline source', () => { + test('retains enclosing pipeline context across argument and block checkpoints', async () => { + const cache = new PipelineSourceCache(); + const prefix = 'value <- 1\n'.repeat(10000) + 'raw %>% mutate(x = other %>% select(a), '; + assert.strictEqual(await cache.resolve(prefix + 'y = sum('), 'raw'); + assert.strictEqual(await cache.resolve(prefix + 'y = third %>% select('), 'third'); + assert.strictEqual( + await cache.resolve(prefix + 'y = identity(third) %>% select('), + undefined, + ); + // Large ordinary function bodies keep only their current statement. + const body = 'f <- function(x) {\n' + 'value <- 1\n'.repeat(10000) + '(raw\n |> select('; + assert.strictEqual(await cache.resolve(body), 'raw'); + assert.strictEqual(await cache.resolve(body + 'a, '), 'raw'); + assert.strictEqual( + await cache.resolve('raw %>% { other %>% select(a); . %>% select('), + 'raw', + ); + }); + + test('invalidates cached context after earlier edits, cursor moves, and region changes', async () => { + const cache = new PipelineSourceCache(); + assert.strictEqual(await cache.resolve('raw %>% mutate(a, '), 'raw'); + assert.strictEqual(await cache.resolve('other %>% mutate(a, '), 'other'); + assert.strictEqual(await cache.resolve('other %>% mutate('), 'other'); + assert.strictEqual(await cache.resolve('select('), undefined); + assert.strictEqual(await cache.resolve('raw %>% mutate(a,\n b = sum('), 'raw'); + }); + + test('retains completion recovery after unrelated malformed code', async () => { + const cache = new PipelineSourceCache(); + assert.strictEqual(await cache.resolve('stray )\n(raw\n |> select('), 'raw'); + assert.strictEqual(await cache.resolve('broken <- c(1]\n(raw\n |> select('), 'raw'); + }); + + test('does not create checkpoints inside strings or comments', async () => { + const cache = new PipelineSourceCache(); + const text = 'raw %>% mutate(note = r"(text,\n fake %>% select()\n)"'; + assert.strictEqual(await cache.resolve(text + ', '), 'raw'); + assert.strictEqual(await cache.resolve(text + ', b = other %>% select('), 'other'); + assert.strictEqual(await cache.resolve(text.slice(0, -2)), undefined); + assert.strictEqual(await cache.resolve('raw %>% mutate( # , ;\n x = sum('), 'raw'); + assert.strictEqual(await cache.resolve('note <- "one,\nother %>% select('), undefined); + }); + + test('yields during cold scans so cancellation can interrupt completion', async () => { + const cache = new PipelineSourceCache(); + const text = 'value <- mean(1:10)\n'.repeat(50000) + 'raw %>% select('; + let canceled = false; + setImmediate(() => { + canceled = true; + }); + assert.strictEqual(await cache.resolve(text, () => canceled), undefined); + assert.ok(canceled); + assert.strictEqual( + await cache.resolve(text), + 'raw', + 'a canceled request must not poison the cache', + ); + }); + + test('yields during long quoted strings', async () => { + const cache = new PipelineSourceCache(); + let canceled = false; + setImmediate(() => { + canceled = true; + }); + assert.strictEqual( + await cache.resolve( + 'note <- "' + 'x'.repeat(100000) + '"\nraw |> select(', + () => canceled, + ), + undefined, + ); + assert.ok(canceled); + }); + + test('bounds active expressions while ignoring oversized completed statements', async () => { + const cache = new PipelineSourceCache(); + const expression = 'raw' + ' |> identity()'.repeat(2000); + assert.strictEqual(await cache.resolve(expression + ' |> select('), undefined); + assert.strictEqual(await cache.resolve(expression + '\nother |> select('), 'other'); + assert.strictEqual( + getPipelineSource('raw |> select(', () => true), + undefined, + ); + }); +}); diff --git a/src/test/suite/completions.test.ts b/src/test/suite/completions.test.ts index aa2168db1..f800162ca 100644 --- a/src/test/suite/completions.test.ts +++ b/src/test/suite/completions.test.ts @@ -227,4 +227,45 @@ suite('Session pipeline completion', () => { assert.deepStrictEqual(items, []); sinon.assert.notCalled(request); }); + + test('cancellation can interrupt a large document completion', async () => { + const content = 'value <- mean(1:10)\n'.repeat(10000) + 'raw |> select('; + const document = await vscode.workspace.openTextDocument({ language: 'r', content }); + const pending = new LiveCompletionItemProvider().provideCompletionItems( + document, + document.positionAt(content.length), + cancellation.token, + { triggerKind: vscode.CompletionTriggerKind.TriggerCharacter, triggerCharacter: '(' }, + ); + setImmediate(() => cancellation.cancel()); + assert.deepStrictEqual(await pending, []); + assert.ok(cancellation.token.isCancellationRequested); + sinon.assert.notCalled(request); + }); + + test('a document edit during scanning suppresses stale pipeline columns', async () => { + const content = 'value <- mean(1:10)\n'.repeat(10000) + 'raw |> select('; + const realDocument = await vscode.workspace.openTextDocument({ language: 'r', content }); + let version = 1; + const document = { + languageId: 'r', + uri: realDocument.uri, + get version() { + return version; + }, + getText: (range?: vscode.Range) => realDocument.getText(range), + } as vscode.TextDocument; + const pending = new LiveCompletionItemProvider().provideCompletionItems( + document, + realDocument.positionAt(content.length), + cancellation.token, + { triggerKind: vscode.CompletionTriggerKind.TriggerCharacter, triggerCharacter: '(' }, + ); + setImmediate(() => { + version++; + }); + assert.deepStrictEqual(await pending, []); + assert.strictEqual(version, 2); + sinon.assert.notCalled(request); + }); });