vs-code-colored-references/src/references.ts

import * as vscode from 'vscode';
import * as path from 'path';

export const SCHEME = 'colored-refs';
export const CODE_INDENT = '    '; // room for the line-number decoration

// ---------------------------------------------------------------------------
// Data model
// ---------------------------------------------------------------------------

/** All references that share one source line. */
export interface ReferenceLine {
    /** 0-based source line. */
    line: number;
    /** Reference ranges on this line, sorted by column. */
    ranges: vscode.Range[];
    /** The raw source line. */
    text: string;
}

export interface ReferenceFile {
    uri: vscode.Uri;
    /** Nearest .csproj / .fsproj / .vbproj, without extension. */
    project: string | undefined;
    lines: ReferenceLine[];
    /** Reference count in this file (a line can hold several). */
    count: number;
}

export interface Origin {
    uri: vscode.Uri;
    position: vscode.Position;
}

/** The result of one Find All References, independent of how it is displayed. */
export interface ReferenceResults {
    symbol: string;
    languageId: string;
    /** Where the search was started, so it can be re-run. */
    origin: Origin;
    files: ReferenceFile[];
    /** Total references, after de-duplication. */
    total: number;
}

/** One displayed line in the virtual results document. */
export interface ResultLine {
    kind: 'code' | 'header' | 'title' | 'blank';
    /** Source file this line belongs to (code + header). */
    file?: vscode.Uri;
    /** 0-based source line (code) */
    sourceLine?: number;
    /** 0-based source column of the first reference on this line (code) */
    sourceCol?: number;
    /** Column ranges of the referenced symbol(s) within the displayed text (code) */
    symbolRanges?: [number, number][];
}

export interface RenderedDocument {
    content: string;
    lines: ResultLine[];
}

// ---------------------------------------------------------------------------
// Reading source files
// ---------------------------------------------------------------------------

const COMMENT_PREFIX: Record<string, string> = {
    python: '#',
    ruby: '#',
    shellscript: '#',
    perl: '#',
    r: '#',
    yaml: '#',
    powershell: '#',
    lua: '--',
    sql: '--',
    haskell: '--',
    vb: "'",
    fsharp: '//',
    csharp: '//',
};

export function commentPrefix(languageId: string): string {
    return COMMENT_PREFIX[languageId] ?? '//';
}

const decoder = new TextDecoder('utf-8');

/** Reads a file's lines, preferring the in-memory (possibly unsaved) version. */
async function readLines(uri: vscode.Uri): Promise<string[]> {
    const open = vscode.workspace.textDocuments.find(d => d.uri.toString() === uri.toString());
    if (open) {
        return open.getText().split(/\r?\n/);
    }
    try {
        const bytes = await vscode.workspace.fs.readFile(uri);
        return decoder.decode(bytes).split(/\r?\n/);
    } catch {
        return [];
    }
}

const projectCache = new Map<string, Promise<string | undefined>>();
const PROJECT_EXTENSIONS = new Set(['.csproj', '.fsproj', '.vbproj']);

/** Walks up from the file to the workspace root looking for a project file. */
function findProject(fileUri: vscode.Uri): Promise<string | undefined> {
    const dir = path.dirname(fileUri.fsPath);
    const cached = projectCache.get(dir);
    if (cached) {
        return cached;
    }

    const promise = (async () => {
        const root = vscode.workspace.getWorkspaceFolder(fileUri)?.uri.fsPath;
        let current = dir;
        for (let depth = 0; depth < 30; depth++) {
            try {
                const entries = await vscode.workspace.fs.readDirectory(vscode.Uri.file(current));
                const proj = entries.find(([name, type]) =>
                    type === vscode.FileType.File && PROJECT_EXTENSIONS.has(path.extname(name).toLowerCase()));
                if (proj) {
                    return path.basename(proj[0], path.extname(proj[0]));
                }
            } catch {
                return undefined;
            }
            if (!root || current === root || path.dirname(current) === current) {
                return undefined;
            }
            current = path.dirname(current);
        }
        return undefined;
    })();

    projectCache.set(dir, promise);
    return promise;
}

export function relativePath(uri: vscode.Uri): string {
    return vscode.workspace.asRelativePath(uri, vscode.workspace.workspaceFolders?.length !== 1);
}

// ---------------------------------------------------------------------------
// Gathering
// ---------------------------------------------------------------------------

/** Runs the reference request through every registered provider. */
export async function findReferences(origin: Origin, symbol: string): Promise<vscode.Location[]> {
    const locations = await vscode.window.withProgress(
        { location: vscode.ProgressLocation.Window, title: `Finding references to '${symbol}'…` },
        () => vscode.commands.executeCommand<vscode.Location[]>(
            'vscode.executeReferenceProvider', origin.uri, origin.position),
    );
    return locations ?? [];
}

/**
 * Groups raw locations by file and source line, reads the source text, and resolves
 * the containing project. Duplicate ranges (two language servers answering the same
 * request) collapse into one.
 */
export async function gather(
    symbol: string,
    languageId: string,
    origin: Origin,
    locations: vscode.Location[],
    showProject: boolean,
): Promise<ReferenceResults> {
    const byFile = new Map<string, { uri: vscode.Uri; byLine: Map<number, vscode.Range[]> }>();
    for (const loc of locations) {
        const key = loc.uri.toString();
        let entry = byFile.get(key);
        if (!entry) {
            entry = { uri: loc.uri, byLine: new Map() };
            byFile.set(key, entry);
        }
        const line = loc.range.start.line;
        const ranges = entry.byLine.get(line) ?? [];
        // De-duplicate identical ranges.
        if (!ranges.some(r => r.isEqual(loc.range))) {
            ranges.push(loc.range);
        }
        entry.byLine.set(line, ranges);
    }

    const sorted = [...byFile.values()].sort((a, b) => a.uri.fsPath.localeCompare(b.uri.fsPath));

    const files = await Promise.all(sorted.map(async (entry): Promise<ReferenceFile> => {
        const [sourceLines, project] = await Promise.all([
            readLines(entry.uri),
            showProject ? findProject(entry.uri) : Promise.resolve(undefined),
        ]);

        const lines = [...entry.byLine.keys()].sort((a, b) => a - b).map((line): ReferenceLine => ({
            line,
            ranges: entry.byLine.get(line)!.sort((a, b) => a.start.character - b.start.character),
            text: sourceLines[line] ?? '',
        }));

        return {
            uri: entry.uri,
            project,
            lines,
            count: lines.reduce((n, l) => n + l.ranges.length, 0),
        };
    }));

    return {
        symbol,
        languageId,
        origin,
        files,
        total: files.reduce((n, f) => n + f.count, 0),
    };
}

// ---------------------------------------------------------------------------
// Rendering the virtual document
// ---------------------------------------------------------------------------

/**
 * Trims a source line for display and maps the reference ranges onto the trimmed text.
 * Displayed offsets differ from source offsets by a constant per line.
 */
export function displayLine(
    source: ReferenceLine, indent = '',
): { text: string; symbolRanges: [number, number][] } {
    const trimmed = source.text.trimStart();
    const removed = source.text.length - trimmed.length;
    const limit = indent.length + trimmed.length;

    const symbolRanges: [number, number][] = [];
    for (const range of source.ranges) {
        const start = Math.max(0, range.start.character - removed) + indent.length;
        const endChar = range.end.line === source.line ? range.end.character : source.text.length;
        const end = Math.max(start + 1, endChar - removed + indent.length);
        symbolRanges.push([Math.min(start, limit), Math.min(end, limit)]);
    }

    return { text: indent + trimmed, symbolRanges };
}

export function renderDocument(results: ReferenceResults): RenderedDocument {
    const prefix = commentPrefix(results.languageId);
    const out: string[] = [];
    const lines: ResultLine[] = [];
    const push = (text: string, meta: ResultLine) => {
        out.push(text);
        lines.push(meta);
    };

    const files = results.files;
    const total = results.total;
    push(`${prefix} ${total} reference${total === 1 ? '' : 's'} to '${results.symbol}' ` +
        `in ${files.length} file${files.length === 1 ? '' : 's'}`, { kind: 'title' });
    push(`${prefix} Enter / F12 / Ctrl+Click: go to reference   F5: refresh`, { kind: 'title' });

    for (const file of files) {
        const projectLabel = file.project ? `[${file.project}]  ` : '';
        push('', { kind: 'blank' });
        push(`${prefix} ${projectLabel}${relativePath(file.uri)}  (${file.count})`,
            { kind: 'header', file: file.uri });

        for (const source of file.lines) {
            const { text, symbolRanges } = displayLine(source, CODE_INDENT);
            push(text, {
                kind: 'code',
                file: file.uri,
                sourceLine: source.line,
                sourceCol: source.ranges[0].start.character,
                symbolRanges,
            });
        }
    }

    return { content: out.join('\n'), lines };
}

// ---------------------------------------------------------------------------
// Containing member (webview panel only — one symbol request per file)
// ---------------------------------------------------------------------------

type AnySymbol = vscode.DocumentSymbol & vscode.SymbolInformation;

const symbolCache = new Map<string, Promise<vscode.DocumentSymbol[]>>();

function timeout<T>(promise: Thenable<T>, ms: number, fallback: T): Promise<T> {
    return Promise.race([
        Promise.resolve(promise),
        new Promise<T>(resolve => setTimeout(() => resolve(fallback), ms)),
    ]);
}

/** Document symbols as a tree, normalising servers that answer with SymbolInformation[]. */
function documentSymbols(uri: vscode.Uri): Promise<vscode.DocumentSymbol[]> {
    const key = uri.toString();
    const cached = symbolCache.get(key);
    if (cached) {
        return cached;
    }
    const promise = (async () => {
        const raw = await timeout(
            vscode.commands.executeCommand<AnySymbol[]>('vscode.executeDocumentSymbolProvider', uri),
            5000, undefined as unknown as AnySymbol[]);
        if (!raw || raw.length === 0) {
            return [];
        }
        if (raw[0].children !== undefined) {
            return raw as vscode.DocumentSymbol[];
        }
        // Flat SymbolInformation[]: rebuild a shallow tree by containment.
        const flat = raw
            .filter(s => s.location)
            .map(s => new vscode.DocumentSymbol(s.name, '', s.kind, s.location.range, s.location.range));
        flat.sort((a, b) => a.range.start.compareTo(b.range.start) ||
            b.range.end.compareTo(a.range.end));
        const roots: vscode.DocumentSymbol[] = [];
        const stack: vscode.DocumentSymbol[] = [];
        for (const symbol of flat) {
            while (stack.length > 0 && !stack[stack.length - 1].range.contains(symbol.range)) {
                stack.pop();
            }
            (stack.length > 0 ? stack[stack.length - 1].children : roots).push(symbol);
            stack.push(symbol);
        }
        return roots;
    })();
    symbolCache.set(key, promise);
    return promise;
}

const SKIPPED_KINDS = new Set([vscode.SymbolKind.Namespace, vscode.SymbolKind.Module, vscode.SymbolKind.Package]);

/** Dotted path of the innermost symbols containing `line`, e.g. `Profiler.BeginSample`. */
function memberPath(symbols: vscode.DocumentSymbol[], line: number): string {
    const names: string[] = [];
    let level = symbols;
    for (;;) {
        const hit = level.find(s => s.range.start.line <= line && line <= s.range.end.line);
        if (!hit) {
            break;
        }
        if (!SKIPPED_KINDS.has(hit.kind)) {
            names.push(hit.name);
        }
        level = hit.children ?? [];
    }
    return names.join('.');
}

/** Resolves the containing member for every referenced line, per file. */
export async function containingMembers(results: ReferenceResults): Promise<Map<string, string>> {
    const members = new Map<string, string>();
    await Promise.all(results.files.map(async file => {
        const symbols = await documentSymbols(file.uri);
        if (symbols.length === 0) {
            return;
        }
        for (const source of file.lines) {
            members.set(`${file.uri.toString()}|${source.line}`, memberPath(symbols, source.line));
        }
    }));
    return members;
}

/** Symbol information is cached per file; drop it when the file changes. */
export function invalidateSymbols(uri: vscode.Uri): void {
    symbolCache.delete(uri.toString());
}