Theia-Modula-2/extensions/modula2-language/src/server.ts

import {
  createConnection, ProposedFeatures, InitializeParams, InitializeResult,
  TextDocuments, TextDocumentSyncKind, Diagnostic, DiagnosticSeverity,
  DidChangeConfigurationParams, CompletionItem, CompletionItemKind,
  DocumentSymbol, SymbolKind, Hover, Location, TextEdit
} from 'vscode-languageserver/node';
import { TextDocument } from 'vscode-languageserver-textdocument';
import * as path from 'path';
import * as os from 'os';
import { promises as fs, readFileSync } from 'fs';
import { spawn } from 'child_process';
import { pathToFileURL } from 'url';
import { lex, Token } from './m2/lexer';
import { M2Symbol, M2Unit, parseUnitText } from './m2/parser';
import { analyseUnit } from './m2/analyse';
import { computeSemanticTokens, semanticTokensLegend } from './m2/semantic';
import { DEFAULT_FORMAT_OPTIONS, FormatOptions, formatDocument, normalizeFormatOptions } from './m2/format';
import {
  moduleExports, recordFieldsOf, resolveName, toLocation, visibleSymbols,
  findReferencesInText, LocatedOccurrence,
} from './m2/resolve';

const connection = createConnection(ProposedFeatures.all);
const documents = new TextDocuments(TextDocument);
let validatorCommand = '';
let validatorArguments: string[] = ['{file}'];
let formatOptions: FormatOptions = { ...DEFAULT_FORMAT_OPTIONS };

connection.onInitialize((params: InitializeParams): InitializeResult => {
  const options = (params.initializationOptions || {}) as {
    validatorCommand?: string; validatorArguments?: string[];
    format?: Partial<FormatOptions>;
  };
  validatorCommand = options.validatorCommand || '';
  validatorArguments = options.validatorArguments || ['{file}'];
  formatOptions = normalizeFormatOptions(options.format ?? {});
  return {
    capabilities: {
      textDocumentSync: TextDocumentSyncKind.Full,
      completionProvider: { triggerCharacters: ['.'] },
      hoverProvider: true,
      definitionProvider: true,
      documentSymbolProvider: true,
      referencesProvider: true,
      renameProvider: { prepareProvider: true },
      documentFormattingProvider: true,
      semanticTokensProvider: {
        legend: semanticTokensLegend(),
        full: true,
      },
    },
  };
});

connection.onDidChangeConfiguration((params: DidChangeConfigurationParams) => {
  const settings = (params.settings || {}) as {
    modula2?: {
      cocoR?: { validatorCommand?: string; validatorArguments?: string[] };
      format?: Partial<FormatOptions>;
    };
  };
  validatorCommand = settings.modula2?.cocoR?.validatorCommand || validatorCommand;
  validatorArguments = settings.modula2?.cocoR?.validatorArguments || validatorArguments;
  if (settings.modula2?.format) {
    formatOptions = normalizeFormatOptions(settings.modula2.format);
  }
});

documents.onDidOpen(e => validate(e.document));
documents.onDidChangeContent(e => validate(e.document));
documents.onDidClose(e => connection.sendDiagnostics({ uri: e.document.uri, diagnostics: [] }));

async function validate(document: TextDocument): Promise<void> {
  const local = analyseDocument(document);
  if (!validatorCommand) {
    connection.sendDiagnostics({ uri: document.uri, diagnostics: local });
    return;
  }

  const filePath = uriToFilePath(document.uri);
  const dir = await fs.mkdtemp(path.join(os.tmpdir(), 'modula2-lsp-'));
  const ext = path.extname(filePath) || '.mod';
  const tempFile = path.join(dir, path.basename(filePath, path.extname(filePath)) + ext);

  try {
    await fs.writeFile(tempFile, document.getText(), 'utf8');
    const args = validatorArguments.map(a => a.replace(/\{file\}/g, tempFile));
    const output = await execute(validatorCommand, args, path.dirname(filePath));
    connection.sendDiagnostics({ uri: document.uri, diagnostics: [...local, ...parseDiagnostics(output)] });
  } catch (error) {
    connection.sendDiagnostics({ uri: document.uri, diagnostics: [...local, {
      severity: DiagnosticSeverity.Warning,
      range: { start: { line: 0, character: 0 }, end: { line: 0, character: 1 } },
      message: `Coco/R validator: ${error instanceof Error ? error.message : String(error)}`,
      source: 'Coco/R'
    }] });
  } finally {
    await fs.rm(dir, { recursive: true, force: true });
  }
}

function execute(command: string, args: string[], cwd: string): Promise<string> {
  return new Promise((resolve, reject) => {
    const child = spawn(command, args, { cwd, shell: false });
    let output = '';
    child.stdout.on('data', d => output += d.toString());
    child.stderr.on('data', d => output += d.toString());
    child.on('error', reject);
    child.on('close', () => resolve(output));
  });
}

function parseDiagnostics(output: string): Diagnostic[] {
  const result: Diagnostic[] = [];
  // CRQ/editor-oriented format: file (line,column) message
  const re = /^.*?\s*\((\d+),(\d+)\)\s+(.*)$/gm;
  let m: RegExpExecArray | null;
  while ((m = re.exec(output)) !== null) {
    const message = m[3].trim();
    if (!message || /^(Grammar tests|Undefined tests|Unreachable tests|Circular tests|Underivable tests|LL\(1\) tests)$/.test(message)) continue;
    const line = Math.max(0, Number(m[1]) - 1);
    const character = Math.max(0, Number(m[2]));
    result.push({ severity: DiagnosticSeverity.Error, range: {
      start: { line, character }, end: { line, character: character + 1 }
    }, message, source: 'Coco/R' });
  }
  return result;
}

function uriToFilePath(uri: string): string {
  if (uri.startsWith('file://')) return decodeURIComponent(new URL(uri).pathname);
  return uri;
}

/** Local semantic diagnostics (always on; merged with Coco/R output when configured). */
function analyseDocument(document: TextDocument): Diagnostic[] {
  const text = document.getText();
  const filePath = uriToFilePath(document.uri);
  return analyseUnit(text, filePath, path.dirname(filePath)).map(issue => ({
    severity: issue.severity === 'error' ? DiagnosticSeverity.Error : DiagnosticSeverity.Warning,
    range: {
      start: { line: issue.line, character: issue.ch },
      end: { line: issue.endLine, character: issue.endCh },
    },
    message: issue.message,
    source: 'modula2',
  }));
}

// ---- Phase 2: AST-backed language intelligence ----

const SERVER_KEYWORDS = [
  'AND', 'ARRAY', 'BEGIN', 'BY', 'CASE', 'CONST', 'DEFINITION', 'DIV',
  'DO', 'ELSE', 'ELSIF', 'END', 'EXIT', 'EXPORT', 'FOR', 'FROM', 'IF',
  'IMPLEMENTATION', 'IMPORT', 'IN', 'LOOP', 'MOD', 'MODULE', 'NOT', 'OF',
  'OR', 'POINTER', 'PROCEDURE', 'QUALIFIED', 'RECORD', 'REPEAT', 'RETURN',
  'SET', 'THEN', 'TO', 'TYPE', 'UNTIL', 'VAR', 'WHILE', 'WITH',
];

function parseDocument(document: TextDocument): { unit: M2Unit; tokens: Token[] } {
  const text = document.getText();
  const filePath = uriToFilePath(document.uri);
  return { unit: parseUnitText(text, filePath), tokens: lex(text) };
}

function tokenAt(tokens: Token[], offset: number): Token | null {
  for (const t of tokens) {
    if (t.kind === 'eof') break;
    if (t.offset <= offset && offset <= t.endOffset) return t;
  }
  return null;
}

/** Extract `Qualifier.name` (or plain `name`) at a document offset. */
function nameAt(tokens: Token[], offset: number): { qualifier: string | null; name: string } | null {
  for (let k = 0; k < tokens.length; k++) {
    const t = tokens[k];
    if (t.kind === 'eof') break;
    if (!(t.offset <= offset && offset <= t.endOffset)) continue;
    if (t.kind !== 'ident') return null;
    // Cursor on the qualifier part: A|.B
    if (tokens[k + 1]?.text === '.' && tokens[k + 2]?.kind === 'ident') {
      return { qualifier: null, name: t.text };
    }
    // Cursor on the qualified part: A.|B  (or second ident after dot)
    if (t.text !== undefined && tokens[k - 1]?.text === '.' && tokens[k - 2]?.kind === 'ident') {
      return { qualifier: tokens[k - 2].text, name: t.text };
    }
    return { qualifier: null, name: t.text };
  }
  return null;
}

function lspKind(kind: M2Symbol['kind']): SymbolKind {
  switch (kind) {
    case 'module': return SymbolKind.Module;
    case 'procedure': return SymbolKind.Function;
    case 'variable': return SymbolKind.Variable;
    case 'parameter': return SymbolKind.Variable;
    case 'field': return SymbolKind.Field;
    case 'constant': return SymbolKind.Constant;
    case 'type': return SymbolKind.Class;
  }
}

function completionKind(kind: M2Symbol['kind']): CompletionItemKind {
  switch (kind) {
    case 'module': return CompletionItemKind.Module;
    case 'procedure': return CompletionItemKind.Function;
    case 'variable': return CompletionItemKind.Variable;
    case 'parameter': return CompletionItemKind.Variable;
    case 'field': return CompletionItemKind.Field;
    case 'constant': return CompletionItemKind.Constant;
    case 'type': return CompletionItemKind.Class;
  }
}

function toDocumentSymbol(s: M2Symbol): DocumentSymbol {
  return {
    name: s.name,
    kind: lspKind(s.kind),
    range: {
      start: { line: s.extent.startLine, character: s.extent.startCh },
      end: { line: s.extent.endLine, character: s.extent.endCh },
    },
    selectionRange: {
      start: { line: s.nameRange.startLine, character: s.nameRange.startCh },
      end: { line: s.nameRange.endLine, character: s.nameRange.endCh },
    },
    detail: s.detail,
    children: s.children.map(toDocumentSymbol),
  };
}

connection.onDocumentSymbol(params => {
  const document = documents.get(params.textDocument.uri);
  if (!document) return [];
  const { unit } = parseDocument(document);
  const out: DocumentSymbol[] = [];
  for (const s of unit.symbols) {
    for (const c of s.children) out.push(toDocumentSymbol(c));
    if (s.children.length === 0 && s.name) out.push(toDocumentSymbol(s));
  }
  return out;
});

connection.onCompletion(params => {
  const document = documents.get(params.textDocument.uri);
  if (!document) return [];
  const { unit } = parseDocument(document);
  const filePath = uriToFilePath(document.uri);
  const docDir = path.dirname(filePath);
  const offset = document.offsetAt(params.position);
  const prefix = document.getText().slice(0, offset);
  const dotMatch = /([A-Za-z_][A-Za-z0-9_]*)\.\s*[A-Za-z_0-9]*$/.exec(prefix);

  if (dotMatch) {
    const qualifier = dotMatch[1];
    // Imported module: offer its exports.
    const exps = moduleExports(docDir, qualifier);
    if (exps.length > 0) {
      return exps.map(s => ({
        label: s.name, kind: completionKind(s.kind), detail: s.detail,
      } as CompletionItem));
    }
    // Variable of record type: offer fields (inline or via named type).
    const local = unit.symbols.flatMap(s => s.children).find(c => c.name === qualifier);
    if (local) {
      if (local.children.length > 0) {
        return local.children.map(f => ({
          label: f.name, kind: CompletionItemKind.Field, detail: f.detail,
        } as CompletionItem));
      }
      const typeMatch = /:\s*([A-Za-z_][A-Za-z0-9_]*)/.exec(local.detail);
      if (typeMatch) {
        const fields = recordFieldsOf(unit, docDir, typeMatch[1]);
        if (fields.length > 0) {
          return fields.map(f => ({
            label: f.name, kind: CompletionItemKind.Field, detail: f.detail,
          } as CompletionItem));
        }
      }
    }
    // Bare type name with record fields.
    const fields = recordFieldsOf(unit, docDir, qualifier);
    if (fields.length > 0) {
      return fields.map(f => ({
        label: f.name, kind: CompletionItemKind.Field, detail: f.detail,
      } as CompletionItem));
    }
    return [];
  }

  const items: CompletionItem[] = visibleSymbols(unit, params.position.line, params.position.character)
    .map(s => ({ label: s.name, kind: completionKind(s.kind), detail: s.detail } as CompletionItem));
  for (const kw of SERVER_KEYWORDS) {
    items.push({ label: kw, kind: CompletionItemKind.Keyword, detail: 'Modula-2 keyword' });
  }
  return items;
});

connection.onHover(params => {
  const document = documents.get(params.textDocument.uri);
  if (!document) return null;
  const { unit, tokens } = parseDocument(document);
  const offset = document.offsetAt(params.position);
  const found = nameAt(tokens, offset);
  if (!found) return null;
  const docDir = path.dirname(uriToFilePath(document.uri));
  const resolved = resolveName(unit, docDir, params.position.line, params.position.character, found.qualifier, found.name);
  if (!resolved) return null;
  const hover: Hover = {
    contents: {
      kind: 'markdown',
      value: [
        '```modula2',
        resolved.sym.detail || resolved.sym.name,
        '```',
        `*${resolved.sym.kind}* — \`${path.basename(resolved.filePath)}:${resolved.sym.nameRange.startLine + 1}\``,
      ].join('\n'),
    },
  };
  return hover;
});

connection.onDefinition(params => {
  const document = documents.get(params.textDocument.uri);
  if (!document) return null;
  const { unit, tokens } = parseDocument(document);
  const offset = document.offsetAt(params.position);
  const found = nameAt(tokens, offset);
  if (!found) return null;
  const filePath = uriToFilePath(document.uri);
  const resolved = resolveName(unit, path.dirname(filePath), params.position.line, params.position.character, found.qualifier, found.name);
  if (!resolved) return null;
  const loc = toLocation(resolved.filePath, resolved.sym);
  const targetUri = resolved.filePath === filePath ? document.uri : pathToFileURL(resolved.filePath).toString();
  const location: Location = {
    uri: targetUri,
    range: {
      start: { line: loc.startLine, character: loc.startCh },
      end: { line: loc.endLine, character: loc.endCh },
    },
  };
  return location;
});

/** Resolve the declaration for a rename/references request, or null. */
function resolveRequest(
  document: TextDocument, line: number, character: number,
): { unit: M2Unit; filePath: string; qualifier: string | null; name: string } | null {
  const { unit, tokens } = parseDocument(document);
  const offset = document.offsetAt({ line, character });
  const found = nameAt(tokens, offset);
  if (!found) return null;
  return { unit, filePath: uriToFilePath(document.uri), qualifier: found.qualifier, name: found.name };
}

/** Reference occurrences across the open documents plus the declaration file. */
function collectReferences(
  documents: TextDocument[], currentUri: string,
  unit: M2Unit, filePath: string, line: number, character: number,
  qualifier: string | null, name: string,
  includeDeclaration: boolean,
): { uri: string; occurrences: LocatedOccurrence[] }[] {
  const decl = resolveName(unit, path.dirname(filePath), line, character, qualifier, name);
  if (!decl) return [];
  const seen = new Map<string, TextDocument>();
  for (const doc of documents) {
    if (!seen.has(doc.uri)) seen.set(doc.uri, doc);
  }
  // The declaration may live in a .def that is not open.
  let extra: { uri: string; text: string } | null = null;
  const declOpen = seen.has(pathToFileURL(decl.filePath).toString());
  if (!declOpen) {
    try {
      extra = { uri: pathToFileURL(decl.filePath).toString(), text: readFileSync(decl.filePath, 'utf8') };
    } catch {
      extra = null;
    }
  }
  void currentUri;
  const out: { uri: string; occurrences: LocatedOccurrence[] }[] = [];
  for (const doc of seen.values()) {
    const fp = uriToFilePath(doc.uri);
    const occurrences = findReferencesInText(
      doc.getText(), fp, path.dirname(fp), decl, includeDeclaration,
    );
    if (occurrences.length > 0) out.push({ uri: doc.uri, occurrences });
  }
  if (extra) {
    const occurrences = findReferencesInText(
      extra.text, decl.filePath, path.dirname(decl.filePath), decl, includeDeclaration,
    );
    if (occurrences.length > 0) out.push({ uri: extra.uri, occurrences });
  }
  return out;
}

connection.onReferences(params => {
  const document = documents.get(params.textDocument.uri);
  if (!document) return [];
  const req = resolveRequest(document, params.position.line, params.position.character);
  if (!req) return [];
  const locations: Location[] = [];
  for (const { uri, occurrences } of collectReferences(
    documents.all(), document.uri, req.unit, req.filePath,
    params.position.line, params.position.character,
    req.qualifier, req.name, params.context.includeDeclaration,
  )) {
    for (const occ of occurrences) {
      locations.push({
        uri,
        range: {
          start: { line: occ.line, character: occ.ch },
          end: { line: occ.endLine, character: occ.endCh },
        },
      });
    }
  }
  return locations;
});

connection.onPrepareRename(params => {
  const document = documents.get(params.textDocument.uri);
  if (!document) return null;
  const req = resolveRequest(document, params.position.line, params.position.character);
  if (!req) return null;
  const decl = resolveName(
    req.unit, path.dirname(req.filePath),
    params.position.line, params.position.character, req.qualifier, req.name,
  );
  if (!decl) return null;
  const offset = document.offsetAt(params.position);
  const { tokens } = parseDocument(document);
  const token = tokenAt(tokens, offset);
  if (!token || token.kind !== 'ident') return null;
  return {
    start: { line: token.line, character: token.ch },
    end: { line: token.endLine, character: token.endCh },
  };
});

connection.onRenameRequest(params => {
  const document = documents.get(params.textDocument.uri);
  if (!document) return null;
  const req = resolveRequest(document, params.position.line, params.position.character);
  if (!req) return null;
  const changes: { [uri: string]: { range: Location['range']; newText: string }[] } = {};
  for (const { uri, occurrences } of collectReferences(
    documents.all(), document.uri, req.unit, req.filePath,
    params.position.line, params.position.character,
    req.qualifier, req.name, true,
  )) {
    changes[uri] = occurrences.map(occ => ({
      range: {
        start: { line: occ.line, character: occ.ch },
        end: { line: occ.endLine, character: occ.endCh },
      },
      newText: params.newName,
    }));
  }
  if (Object.keys(changes).length === 0) return null;
  return { changes };
});

connection.onDocumentFormatting(params => {
  const document = documents.get(params.textDocument.uri);
  if (!document) return null;
  const filePath = uriToFilePath(document.uri);
  let formatted: string;
  try {
    formatted = formatDocument(
      document.getText(), formatOptions, filePath, path.dirname(filePath),
    );
  } catch (error) {
    connection.console.error(`formatting failed: ${error instanceof Error ? error.message : String(error)}`);
    return null;
  }
  if (formatted === document.getText()) return [];
  const lines = document.getText().split('\n');
  const last = lines.length - 1;
  const edit: TextEdit = {
    range: {
      start: { line: 0, character: 0 },
      end: { line: last, character: lines[last]?.length ?? 0 },
    },
    newText: formatted,
  };
  return [edit];
});

connection.languages.semanticTokens.on(params => {
  const document = documents.get(params.textDocument.uri);
  if (!document) return { data: [] };
  const filePath = uriToFilePath(document.uri);
  try {
    return {
      data: computeSemanticTokens(document.getText(), filePath, path.dirname(filePath)),
    };
  } catch (error) {
    connection.console.error(`semantic tokens failed: ${error instanceof Error ? error.message : String(error)}`);
    return { data: [] };
  }
});

documents.listen(connection);
connection.listen();