runa/vscode/src/diagnostics.ts

import * as vscode from "vscode";
import { ImageClient } from "./client";
import { SourceDiagnostic } from "./protocol";
import { ACTOR_SCHEME } from "./actorFiles";
import { SOURCE_SCHEME } from "./sourceProvider";

interface Connection {
  session: number;
  endpoint: string;
  client: Pick<ImageClient, "checkSource">;
}

/** Validate editor text without saving, publishing, or installing it in the image. */
export class SourceDiagnostics implements vscode.Disposable {
  private readonly collection = vscode.languages.createDiagnosticCollection("runa");
  private readonly pending = new Map<string, vscode.TextDocument>();
  private readonly subscriptions: vscode.Disposable[];
  private timer?: ReturnType<typeof setTimeout>;
  private running = false;
  private generation = 0;
  private disposed = false;

  constructor(
    private readonly connection: () => Connection | undefined,
    private readonly report: (error: unknown) => void,
  ) {
    this.subscriptions = [
      vscode.workspace.onDidOpenTextDocument((document) => this.schedule(document)),
      vscode.workspace.onDidChangeTextDocument(({ document }) => this.schedule(document)),
      vscode.workspace.onDidCloseTextDocument((document) => {
        this.pending.delete(document.uri.toString());
        this.collection.delete(document.uri);
      }),
    ];
  }

  private matches(document: vscode.TextDocument, connection: Connection): boolean {
    return (
      !document.isClosed &&
      ((document.uri.scheme === SOURCE_SCHEME &&
        document.uri.authority === connection.endpoint &&
        document.uri.path.endsWith(".rn")) ||
        (document.uri.scheme === ACTOR_SCHEME &&
          document.uri.authority === String(connection.session) &&
          document.uri.path.endsWith("/behavior.rn")))
    );
  }

  refresh(): void {
    for (const document of vscode.workspace.textDocuments) this.schedule(document);
  }

  clear(): void {
    this.generation++;
    clearTimeout(this.timer);
    this.timer = undefined;
    this.pending.clear();
    this.collection.clear();
  }

  private schedule(document: vscode.TextDocument): void {
    const connection = this.connection();
    if (this.disposed || !connection || !this.matches(document, connection)) return;
    this.collection.delete(document.uri);
    this.pending.set(document.uri.toString(), document);
    if (this.running) return;
    clearTimeout(this.timer);
    this.timer = setTimeout(() => {
      this.timer = undefined;
      void this.run();
    }, 300);
  }

  private async run(): Promise<void> {
    this.running = true;
    try {
      while (this.pending.size && !this.disposed) {
        const document = this.pending.values().next().value;
        if (!document) break;
        this.pending.delete(document.uri.toString());
        const connection = this.connection();
        if (!connection || !this.matches(document, connection)) continue;
        const generation = this.generation;
        const version = document.version;
        try {
          const diagnostics = await connection.client.checkSource(document.getText());
          if (
            this.disposed ||
            generation !== this.generation ||
            this.connection()?.client !== connection.client ||
            document.isClosed ||
            document.version !== version
          )
            continue;
          this.collection.set(
            document.uri,
            diagnostics.map((diagnostic) => this.toEditor(document, diagnostic)),
          );
        } catch (error) {
          if (
            !this.disposed &&
            generation === this.generation &&
            this.connection()?.client === connection.client &&
            !document.isClosed &&
            document.version === version
          )
            this.report(error);
        }
      }
    } finally {
      this.running = false;
    }
  }

  private toEditor(document: vscode.TextDocument, diagnostic: SourceDiagnostic): vscode.Diagnostic {
    let range = new vscode.Range(0, 0, 0, 0);
    if (diagnostic.range !== null) {
      const { start, end } = diagnostic.range;
      if (
        end.line >= document.lineCount ||
        start.line >= document.lineCount ||
        start.character > document.lineAt(start.line).text.length ||
        end.character > document.lineAt(end.line).text.length
      )
        throw new Error("Rune diagnostic range exceeds the validated document");
      range = new vscode.Range(start.line, start.character, end.line, end.character);
    }
    const result = new vscode.Diagnostic(
      range,
      diagnostic.message,
      diagnostic.severity === "error"
        ? vscode.DiagnosticSeverity.Error
        : vscode.DiagnosticSeverity.Warning,
    );
    result.source = "Rune";
    return result;
  }

  dispose(): void {
    this.disposed = true;
    this.clear();
    for (const subscription of this.subscriptions) subscription.dispose();
    this.collection.dispose();
  }
}