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>

README.md

@@ -3,7 +3,9 @@
 Find All References results shown the way Visual Studio does it — in two views.
 
 **Editor view** (default) writes the results into a read-only virtual document in the *same language as your
-source file*, so your theme's grammar colors every line for free. File headers show the containing project
+source file*, so your theme colors every line for free — including identifiers, because the semantic tokens
+are fetched from the real source files and relocated into the results (see *Semantic tokens* below). File
+headers show the containing project
 (`.csproj`/`.fsproj`/`.vbproj`) and reference count; line numbers are shown in the gutter. Read references are
 highlighted like search matches; **writes are bold** on the theme's stronger write-occurrence background, the
 same colour the editor itself uses for a write.
@@ -34,12 +36,29 @@ The **All / Reads / Writes** buttons filter by kind, and the summary line counts
 the text filter, so typing `write` narrows to writes as well. See *Known limitations* for where the kind comes
 from.
 
+## Semantic tokens
+
+VS Code colors code twice: a TextMate grammar handles keywords, strings and comments from the text alone, then
+the language server's *semantic tokens* recolor identifiers by what they actually are — a class, a method, a
+parameter. That second pass is why a C# file settles into its final colors a moment after you open it.
+
+Language servers will not serve semantic tokens for the results document: its `colored-refs` URI has no file
+and no compilation behind it. So the tokens are fetched for each *source* file instead and relocated into
+results-document coordinates — the same line and column mapping the extension already uses to place the
+highlights. The colouring you get is the server's own answer, moved, rather than a guess.
+
+This happens after the results are on screen and only ever changes colors, never text, so nothing waits on it.
+Set `coloredReferences.semanticTokens` to `false` for TextMate-only coloring. Results spanning more than 40
+files skip it.
+
 ## Settings
 
 - `coloredReferences.view` — which view `Find All References (Colored)` opens: `document` (default) or `panel`
 - `coloredReferences.panelLocation` — where the panel opens: `bottom` (default, docked next to Terminal /
   Problems and draggable to a side bar), `beside` (editor group to the side), or `below` (editor group
   underneath, so the table is wide and short)
+- `coloredReferences.semanticTokens` — color the editor view with semantic tokens fetched from the source
+  files (default `true`)
 - `coloredReferences.openBeside` — open results beside the current editor (default `true`)
 - `coloredReferences.showProject` — show the containing project in file headers (default `true`)
 - `coloredReferences.reuseTab` — reuse one results tab/panel instead of opening a new one per search (default `true`)
@@ -67,11 +86,10 @@ the expected results stay deterministic.
 
 ## Known limitations
 
-- Coloring in the editor view is TextMate-only (no semantic tokens), since language servers only serve semantic
-  tokens for real files. Types and identifiers therefore look like they do in a freshly opened file before the
-  server has analyzed it.
 - The panel's code column cannot use your theme's token colors: webviews are not given them as CSS variables.
-  It approximates the stock Dark+/Light+ hues instead. Semantic tokens (roadmap 2) would replace this.
+  It classifies code with a small tokenizer of its own and paints it with approximated Dark+/Light+ hues, so
+  it will not match your theme exactly. Semantic tokens do not fix this — they supply the *classification*
+  ("this is a class"), never the colors — so only the editor view can be theme-exact.
 - Read/write kind is not part of the reference request, so it is derived two ways. Every reference is first
   classified from the surrounding text (assignment and compound-assignment operators, `++`/`--`, `ref`/`out`
   arguments), then `textDocument/documentHighlight` is asked per file and its `Read`/`Write` kinds override
@@ -86,6 +104,8 @@ the expected results stay deterministic.
 
 1. ~~Webview panel with resizable, sortable columns~~ — done; missing: virtualized rendering for very large
    result sets, and remembering column layout per workspace rather than per panel
-2. Semantic token overlay via `vscode.provideDocumentSemanticTokens`
+2. ~~Semantic token overlay~~ — done for the editor view; the panel still uses its own tokenizer
 3. Filter by project / exclude tests
 4. ~~Read/write kind~~ — done in both views
+5. Feed the panel's code column the real token types instead of its regex tokenizer (colors would still be
+   the approximated palette, but the classification would be right)

package.json

@@ -160,6 +160,11 @@
           ],
           "default": "bottom",
           "description": "Where the results panel opens."
+        },
+        "coloredReferences.semanticTokens": {
+          "type": "boolean",
+          "default": true,
+          "description": "Colour the editor view with semantic tokens fetched from the source files, so identifiers look the way they do in the real file. Turn off to use TextMate grammar colouring only."
         }
       }
     },

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(', ')}`);
+    });
+});