Colour the editor view with relocated semantic tokens

The results document only ever got TextMate colouring, because a language
server will not serve semantic tokens for a `colored-refs` URI: there is no
file and no compilation behind it. So identifiers looked the way a C# file
does before the server has analysed it.

Instead of tokenizing anything ourselves, fetch the tokens for each *source*
file and move them:

- vscode.provideDocumentSemanticTokensLegend / provideDocumentSemanticTokens
  give the legend and the delta-encoded tokens for a real file. Neither is
  listed by getCommands(), so a test calls them to prove they exist rather
  than looking them up.
- Decode to absolute positions, keep the tokens on displayed lines, and shift
  each column by CODE_INDENT minus the stripped leading whitespace — the same
  mapping that places the highlight ranges. Re-encode with the source legend,
  so the type and modifier numbers stay meaningful.
- The legend is only known after a server answers, but
  registerDocumentSemanticTokensProvider wants it up front, so registration
  is deferred to the first search and redone if a later legend differs.
- Runs after the results are on screen and only changes colours, never text,
  so nothing waits on it. coloredReferences.semanticTokens turns it off;
  results over 40 files skip it.

Verified against DotRush: every relocated token covers exactly the text it
covered in the source file, and `Profiler` comes back typed as `class`.

README: drops the claim that semantic tokens would fix the panel's colours.
They supply the classification, never the colours, and webviews are not given
theme token colours — so only the editor view can be theme-exact. Feeding the
panel real token types is now roadmap item 5.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

src/extension.ts

@@ -1,6 +1,7 @@
 import * as vscode from 'vscode';
 import * as path from 'path';
 import { PanelManager, Placement, ResultsView } from './panel';
+import { SemanticOverlay } from './semantic';
 import {
     CODE_INDENT, Origin, ReferenceResults, RenderedDocument, ResultLine, SCHEME,
     applyKinds, findReferences, gather, invalidateSymbols, referenceKinds, relativePath,
@@ -60,6 +61,7 @@ class ResultsStore {
 }
 
 const store = new ResultsStore();
+const semantics = new SemanticOverlay();
 let panels: PanelManager;
 
 /** Open tabs showing a given results URI. */
@@ -244,7 +246,13 @@ async function showDocument(results: ReferenceResults, targetUri?: vscode.Uri):
         editor.revealRange(new vscode.Range(0, 0, firstCode, 0), vscode.TextEditorRevealType.AtTop);
     }
     applyDecorations(editor);
+
+    // Both refinements only need the results already on screen: the kinds change
+    // decorations, the tokens change colouring, and neither touches the text.
     void refineKinds(uri, stored);
+    if (cfg.get<boolean>('semanticTokens', true)) {
+        void semantics.build(uri, stored.lines, results).catch(() => undefined);
+    }
 }
 
 async function showPanel(results: ReferenceResults, target?: ResultsView): Promise<void> {
@@ -470,6 +478,7 @@ export function activate(context: vscode.ExtensionContext): void {
                     vscode.workspace.textDocuments.some(d => d.uri.toString() === uri.toString());
                 if (!stillOpen) {
                     store.delete(uri);
+                    semantics.forget(uri);
                 }
             }, 0);
         }),
@@ -480,6 +489,7 @@ export function activate(context: vscode.ExtensionContext): void {
         headerDecoration,
         titleDecoration,
         store.onDidChange,
+        semantics,
         { dispose: () => panels.dispose() },
     );
 

src/semantic.ts

@@ -0,0 +1,194 @@
+import * as vscode from 'vscode';
+import { CODE_INDENT, ReferenceResults, ResultLine, SCHEME } from './references';
+
+/**
+ * Semantic token overlay for the virtual results document.
+ *
+ * Language servers only serve semantic tokens for real files — the results document has
+ * no compilation behind it — so the tokens are fetched for each *source* file and
+ * relocated into results-document coordinates. The colouring is then exactly what the
+ * server would paint in the real file, rather than a guess of our own.
+ */
+
+/** One token, in absolute coordinates. */
+export interface Token {
+    line: number;
+    char: number;
+    length: number;
+    type: number;
+    modifiers: number;
+}
+
+/** Asking for whole-document tokens per file gets expensive on huge result sets. */
+const SEMANTIC_FILE_LIMIT = 40;
+
+function timeout<T>(promise: Thenable<T>, ms: number): Promise<T | undefined> {
+    return Promise.race([
+        Promise.resolve(promise),
+        new Promise<undefined>(resolve => setTimeout(() => resolve(undefined), ms)),
+    ]);
+}
+
+/** Decodes the delta-encoded `SemanticTokens.data` into absolute tokens. */
+export function decodeTokens(data: Uint32Array): Token[] {
+    const tokens: Token[] = [];
+    let line = 0;
+    let char = 0;
+    for (let i = 0; i + 4 < data.length; i += 5) {
+        const deltaLine = data[i];
+        const deltaChar = data[i + 1];
+        if (deltaLine === 0) {
+            char += deltaChar;
+        } else {
+            line += deltaLine;
+            char = deltaChar;
+        }
+        tokens.push({ line, char, length: data[i + 2], type: data[i + 3], modifiers: data[i + 4] });
+    }
+    return tokens;
+}
+
+async function sourceLegend(uri: vscode.Uri): Promise<vscode.SemanticTokensLegend | undefined> {
+    const legend = await timeout(vscode.commands.executeCommand<vscode.SemanticTokensLegend>(
+        'vscode.provideDocumentSemanticTokensLegend', uri), 4000);
+    return legend?.tokenTypes?.length ? legend : undefined;
+}
+
+async function sourceTokens(uri: vscode.Uri): Promise<Token[]> {
+    try {
+        // The server needs the document open to answer.
+        await vscode.workspace.openTextDocument(uri);
+    } catch {
+        return [];
+    }
+    const tokens = await timeout(vscode.commands.executeCommand<vscode.SemanticTokens>(
+        'vscode.provideDocumentSemanticTokens', uri), 8000);
+    return tokens?.data ? decodeTokens(tokens.data) : [];
+}
+
+/**
+ * Where a source line ended up in the results document, and by how much its columns
+ * shifted when the leading whitespace was replaced with {@link CODE_INDENT}.
+ */
+function lineIndex(lines: ResultLine[], results: ReferenceResults): Map<string, { row: number; shift: number }> {
+    const indent = CODE_INDENT.length;
+    const text = new Map<string, string>();
+    for (const file of results.files) {
+        for (const line of file.lines) {
+            text.set(`${file.uri.toString()}|${line.line}`, line.text);
+        }
+    }
+
+    const index = new Map<string, { row: number; shift: number }>();
+    lines.forEach((line, row) => {
+        if (line.kind !== 'code' || !line.file || line.sourceLine === undefined) {
+            return;
+        }
+        const key = `${line.file.toString()}|${line.sourceLine}`;
+        const raw = text.get(key) ?? '';
+        const removed = raw.length - raw.trimStart().length;
+        index.set(key, { row, shift: indent - removed });
+    });
+    return index;
+}
+
+/**
+ * Serves relocated semantic tokens for results documents.
+ *
+ * The legend is only known once a language server has answered, and
+ * `registerDocumentSemanticTokensProvider` wants it up front, so registration happens
+ * on the first search and is redone if a later search brings a different legend.
+ */
+export class SemanticOverlay {
+    private readonly tokens = new Map<string, vscode.SemanticTokens>();
+    private readonly changed = new vscode.EventEmitter<void>();
+    private registration: vscode.Disposable | undefined;
+    private legend: vscode.SemanticTokensLegend | undefined;
+
+    /**
+     * Fetches tokens for every source file in `results` and stores the relocated set
+     * for `uri`. Returns false when no server served any, so nothing changed.
+     */
+    async build(uri: vscode.Uri, lines: ResultLine[], results: ReferenceResults): Promise<boolean> {
+        if (results.files.length > SEMANTIC_FILE_LIMIT) {
+            return false;
+        }
+
+        const legend = await sourceLegend(results.origin.uri);
+        if (!legend) {
+            return false; // this language server does not do semantic tokens
+        }
+
+        const index = lineIndex(lines, results);
+        const relocated: Token[] = [];
+
+        await Promise.all(results.files.map(async file => {
+            const displayed = new Set(file.lines.map(l => l.line));
+            for (const token of await sourceTokens(file.uri)) {
+                if (!displayed.has(token.line)) {
+                    continue;
+                }
+                const target = index.get(`${file.uri.toString()}|${token.line}`);
+                if (!target) {
+                    continue;
+                }
+                const char = token.char + target.shift;
+                if (char < 0) {
+                    continue; // inside the stripped indentation
+                }
+                relocated.push({ ...token, line: target.row, char });
+            }
+        }));
+
+        if (relocated.length === 0) {
+            return false;
+        }
+
+        // The builder delta-encodes, so tokens have to arrive in document order.
+        relocated.sort((a, b) => a.line - b.line || a.char - b.char);
+        const builder = new vscode.SemanticTokensBuilder();
+        for (const token of relocated) {
+            builder.push(token.line, token.char, token.length, token.type, token.modifiers);
+        }
+
+        this.tokens.set(uri.toString(), builder.build());
+        this.ensureRegistered(legend);
+        this.changed.fire();
+        return true;
+    }
+
+    private ensureRegistered(legend: vscode.SemanticTokensLegend): void {
+        const same = this.legend &&
+            this.legend.tokenTypes.join() === legend.tokenTypes.join() &&
+            this.legend.tokenModifiers.join() === legend.tokenModifiers.join();
+        if (this.registration && same) {
+            return;
+        }
+
+        this.registration?.dispose();
+        this.legend = legend;
+        this.registration = vscode.languages.registerDocumentSemanticTokensProvider(
+            { scheme: SCHEME },
+            {
+                onDidChangeSemanticTokens: this.changed.event,
+                provideDocumentSemanticTokens: document => this.tokens.get(document.uri.toString()),
+            },
+            legend);
+    }
+
+    forget(uri: vscode.Uri): void {
+        this.tokens.delete(uri.toString());
+    }
+
+    /** Test seam: the tokens currently served for a results document. */
+    peek(uri: vscode.Uri): vscode.SemanticTokens | undefined {
+        return this.tokens.get(uri.toString());
+    }
+
+    dispose(): void {
+        this.registration?.dispose();
+        this.registration = undefined;
+        this.changed.dispose();
+        this.tokens.clear();
+    }
+}

src/test/suite/index.ts

@@ -13,6 +13,7 @@ export function run(): Promise<void> {
     mocha.addFile(path.resolve(__dirname, 'references.test.js'));
     mocha.addFile(path.resolve(__dirname, 'panel.test.js'));
     mocha.addFile(path.resolve(__dirname, 'kinds.test.js'));
+    mocha.addFile(path.resolve(__dirname, 'semantic.test.js'));
 
     return new Promise((resolve, reject) => {
         try {

src/test/suite/semantic.test.ts

@@ -0,0 +1,160 @@
+import * as assert from 'assert';
+import * as path from 'path';
+import * as vscode from 'vscode';
+
+import { SemanticOverlay, decodeTokens } from '../../semantic';
+import { CODE_INDENT, gather, renderDocument } from '../../references';
+
+const TARGET_FILE = path.join('Nerfed.Runtime', 'Profiler.cs');
+const TARGET_SYMBOL = 'Profiler';
+const TARGET_DECL = 'public static class Profiler';
+
+const LEGEND_COMMAND = 'vscode.provideDocumentSemanticTokensLegend';
+const TOKENS_COMMAND = 'vscode.provideDocumentSemanticTokens';
+
+suite('Semantic token overlay', () => {
+    let sourceUri: vscode.Uri;
+    let position: vscode.Position;
+
+    suiteSetup(async () => {
+        const ext = vscode.extensions.getExtension('local.colored-references');
+        assert.ok(ext);
+        await ext.activate();
+
+        const folder = vscode.workspace.workspaceFolders?.[0];
+        assert.ok(folder, 'no workspace folder');
+        sourceUri = vscode.Uri.joinPath(folder.uri, ...TARGET_FILE.split(path.sep));
+        const doc = await vscode.workspace.openTextDocument(sourceUri);
+        for (let i = 0; i < doc.lineCount; i++) {
+            const at = doc.lineAt(i).text.indexOf(TARGET_DECL);
+            if (at >= 0) {
+                position = new vscode.Position(i, doc.lineAt(i).text.indexOf(TARGET_SYMBOL, at) + 1);
+                break;
+            }
+        }
+        assert.ok(position, 'declaration not found');
+    });
+
+    test('the built-in semantic token commands are callable', async () => {
+        // These API commands are not enumerated by getCommands(), so the only way to
+        // check they exist is to call them: a missing command rejects with "not found".
+        const commands = await vscode.commands.getCommands(true);
+        assert.ok(!commands.includes(TOKENS_COMMAND),
+            'this assertion documents that the command is unlisted — if it now appears, ' +
+            'simplify this test to a getCommands() check');
+
+        for (const command of [LEGEND_COMMAND, TOKENS_COMMAND]) {
+            try {
+                await vscode.commands.executeCommand(command, sourceUri);
+            } catch (error) {
+                assert.fail(`${command} is not available: ${error}`);
+            }
+        }
+    });
+
+    test('the language server serves a legend and tokens for a real file', async () => {
+        const legend = await vscode.commands.executeCommand<vscode.SemanticTokensLegend>(
+            LEGEND_COMMAND, sourceUri);
+        assert.ok(legend?.tokenTypes?.length, 'no legend from the language server');
+        console.log(`[test] token types: ${legend.tokenTypes.slice(0, 20).join(', ')}` +
+            `${legend.tokenTypes.length > 20 ? ', …' : ''}`);
+
+        const tokens = await vscode.commands.executeCommand<vscode.SemanticTokens>(
+            TOKENS_COMMAND, sourceUri);
+        assert.ok(tokens?.data?.length, 'no semantic tokens for the source file');
+        assert.strictEqual(tokens.data.length % 5, 0, 'token data must be 5-tuples');
+    });
+
+    test('decodeTokens turns deltas back into absolute positions', () => {
+        // two tokens on line 3, one on line 5
+        const data = new Uint32Array([3, 4, 6, 1, 0, 0, 10, 3, 2, 0, 2, 7, 5, 4, 1]);
+        assert.deepStrictEqual(decodeTokens(data), [
+            { line: 3, char: 4, length: 6, type: 1, modifiers: 0 },
+            { line: 3, char: 14, length: 3, type: 2, modifiers: 0 },
+            { line: 5, char: 7, length: 5, type: 4, modifiers: 1 },
+        ]);
+    });
+
+    test('relocated tokens cover the same text in the results document', async () => {
+        const locations = await vscode.commands.executeCommand<vscode.Location[]>(
+            'vscode.executeReferenceProvider', sourceUri, position);
+        assert.ok(locations && locations.length > 5, 'not enough references to test with');
+
+        const results = await gather(
+            TARGET_SYMBOL, 'csharp', { uri: sourceUri, position }, locations, true);
+        const rendered = renderDocument(results);
+
+        const overlay = new SemanticOverlay();
+        const uri = vscode.Uri.from({ scheme: 'colored-refs', path: '/semantic-test.cs', query: 't' });
+        const built = await overlay.build(uri, rendered.lines, results);
+        assert.ok(built, 'the overlay produced no tokens');
+
+        const served = overlay.peek(uri);
+        assert.ok(served, 'no tokens stored for the results document');
+        const tokens = decodeTokens(served.data);
+        assert.ok(tokens.length > 10, `only ${tokens.length} tokens relocated`);
+
+        const displayLines = rendered.content.split('\n');
+        const sources = new Map<string, vscode.TextDocument>();
+        for (const file of results.files) {
+            sources.set(file.uri.toString(), await vscode.workspace.openTextDocument(file.uri));
+        }
+
+        // Every token must sit on a code line and cover exactly the text it covered in
+        // the source file — that is the whole correctness claim of the relocation.
+        const failures: string[] = [];
+        for (const token of tokens) {
+            const line = rendered.lines[token.line];
+            if (!line || line.kind !== 'code' || !line.file || line.sourceLine === undefined) {
+                failures.push(`token at row ${token.line} is not on a code line`);
+                continue;
+            }
+            const source = sources.get(line.file.toString());
+            if (!source) {
+                failures.push(`token at row ${token.line} has no source document`);
+                continue;
+            }
+            const sourceText = source.lineAt(line.sourceLine).text;
+            const shift = CODE_INDENT.length - (sourceText.length - sourceText.trimStart().length);
+
+            const displayed = displayLines[token.line].substr(token.char, token.length);
+            const original = sourceText.substr(token.char - shift, token.length);
+            if (displayed !== original) {
+                failures.push(`row ${token.line} col ${token.char}: results show ` +
+                    `${JSON.stringify(displayed)} but the source has ${JSON.stringify(original)}`);
+            }
+            if (displayed.trim().length === 0) {
+                failures.push(`row ${token.line} col ${token.char}: token covers whitespace`);
+            }
+        }
+        assert.deepStrictEqual(failures.slice(0, 10), [],
+            `${failures.length} mismatched token(s):\n${failures.slice(0, 10).join('\n')}`);
+    });
+
+    test('the searched symbol is classified as a type', async () => {
+        const legend = await vscode.commands.executeCommand<vscode.SemanticTokensLegend>(
+            LEGEND_COMMAND, sourceUri);
+        assert.ok(legend);
+
+        const locations = await vscode.commands.executeCommand<vscode.Location[]>(
+            'vscode.executeReferenceProvider', sourceUri, position);
+        const results = await gather(
+            TARGET_SYMBOL, 'csharp', { uri: sourceUri, position }, locations ?? [], true);
+        const rendered = renderDocument(results);
+
+        const overlay = new SemanticOverlay();
+        const uri = vscode.Uri.from({ scheme: 'colored-refs', path: '/semantic-kinds.cs', query: 't' });
+        assert.ok(await overlay.build(uri, rendered.lines, results));
+        const tokens = decodeTokens(overlay.peek(uri)!.data);
+
+        const displayLines = rendered.content.split('\n');
+        const onSymbol = tokens.filter(t =>
+            displayLines[t.line]?.substr(t.char, t.length) === TARGET_SYMBOL);
+        assert.ok(onSymbol.length > 0, `no token covers '${TARGET_SYMBOL}'`);
+
+        const names = [...new Set(onSymbol.map(t => legend.tokenTypes[t.type]))];
+        console.log(`[test] '${TARGET_SYMBOL}' token types: ${names.join(', ')}`);
+        assert.ok(names.every(n => /class|type|struct|enum|interface/i.test(n)),
+            `'${TARGET_SYMBOL}' is a static class but was typed ${names.join(', ')}`);
+    });
+});