TypeScript Language Manager in memory cache

InMemoryLanguageServiceHost.ts

import ts from "typescript/lib/tsserverlibrary.js";
import fs from "node:fs";
import path from "node:path";
import { DocumentRegistry } from "typescript";

const tsconfig = ts.readConfigFile("./tsconfig.json", ts.sys.readFile);
const options = ts.parseJsonConfigFileContent(tsconfig.config, ts.sys, "./");
const defaultHost = ts.createCompilerHost(options.options);

const expandPath = (fname: string) => {
  if (fname.startsWith("/")) {
    return fname;
  }
  const root = process.cwd();
  return path.join(root, fname);
};

function applyRenameLocations(
  code: string,
  toName: string,
  renameLocations: readonly ts.RenameLocation[],
) {
  let current = code;
  let offset = 0;
  for (const loc of renameLocations) {
    const start = loc.textSpan.start;
    const end = loc.textSpan.start + loc.textSpan.length;
    current = current.slice(0, start + offset) + toName +
      current.slice(end + offset);
    offset += toName.length - (end - start);
  }
  return current;
}

type SnapshotManager = {
  readFileSnapshot(fileName: string): string | undefined;
  writeFileSnapshot(fileName: string, content: string): ts.SourceFile;
};

export interface InMemoryLanguageServiceHost extends ts.LanguageServiceHost {
  getSnapshotManager: (
    registory: DocumentRegistry,
  ) => SnapshotManager;
}

export function createInMemoryLanguageServiceHost(): InMemoryLanguageServiceHost {
  // read once, write on memory
  const fileContents = new Map<string, string>();
  const fileSnapshots = new Map<string, ts.IScriptSnapshot>();
  const fileVersions = new Map<string, number>();

  const getSnapshotManagerInternal: (
    registory: DocumentRegistry,
  ) => SnapshotManager = (registory: ts.DocumentRegistry) => {
    return {
      readFileSnapshot(fileName: string) {
        fileName = expandPath(fileName);
        console.log("[readFileSnapshot]", fileName);
        if (fileContents.has(fileName)) {
          return fileContents.get(fileName) as string;
        }
        return defaultHost.readFile(fileName);
      },
      writeFileSnapshot(fileName: string, content: string) {
        fileName = expandPath(fileName);
        const nextVersion = (fileVersions.get(fileName) || 0) + 1;
        fileVersions.set(fileName, nextVersion);
        fileContents.set(fileName, content);
        console.log(
          "[writeFileSnapshot]",
          fileName,
          nextVersion,
          content.length,
        );
        const newSource = registory.updateDocument(
          fileName,
          serviceHost,
          ts.ScriptSnapshot.fromString(content),
          String(nextVersion),
        );
        return newSource;
      },
    };
  };

  const serviceHost: InMemoryLanguageServiceHost = {
    getDefaultLibFileName: defaultHost.getDefaultLibFileName,
    fileExists: ts.sys.fileExists,
    readDirectory: ts.sys.readDirectory,
    directoryExists: ts.sys.directoryExists,
    getDirectories: ts.sys.getDirectories,
    getCurrentDirectory: defaultHost.getCurrentDirectory,
    getScriptFileNames: () => options.fileNames,
    getCompilationSettings: () => options.options,
    readFile: (fname, encode) => {
      fname = expandPath(fname);
      // console.log("[readFile]", fname);
      if (fileContents.has(fname)) {
        return fileContents.get(fname) as string;
      }
      const rawFileResult = ts.sys.readFile(fname, encode);
      if (rawFileResult) {
        fileContents.set(fname, rawFileResult);
        fileVersions.set(
          fname,
          (fileVersions.get(fname) || 0) + 1,
        );
      }
      return rawFileResult;
    },
    writeFile: (fileName, content) => {
      fileName = expandPath(fileName);
      console.log("[writeFile]", fileName);
      fileContents.set(fileName, content);
      const version = fileVersions.get(fileName) || 0;
      fileVersions.set(fileName, version + 1);
    },
    getScriptSnapshot: (fileName) => {
      fileName = expandPath(fileName);
      if (fileName.includes("src/index.ts")) {
        console.log("[getScriptSnapshot]", fileName);
      }
      if (fileSnapshots.has(fileName)) {
        return fileSnapshots.get(fileName)!;
      }
      const contentCache = fileContents.get(fileName);
      if (contentCache) {
        const newSnapshot = ts.ScriptSnapshot.fromString(contentCache);
        fileSnapshots.set(fileName, newSnapshot);
        return newSnapshot;
      }
      if (!fs.existsSync(fileName)) return;
      const raw = ts.sys.readFile(fileName, "utf8")!;
      const snopshot = ts.ScriptSnapshot.fromString(raw);
      fileSnapshots.set(fileName, snopshot);
      return snopshot;
    },
    getScriptVersion: (fileName) => {
      fileName = expandPath(fileName);
      return (fileVersions.get(fileName) || 0).toString();
    },
    getSnapshotManager: getSnapshotManagerInternal,
  };
  return serviceHost;
}

{
  // usage
  const prefs: ts.UserPreferences = {};
  const registory = ts.createDocumentRegistry();
  const serviceHost = createInMemoryLanguageServiceHost();
  const languageService = ts.createLanguageService(
    serviceHost,
    registory,
  );
  const snapshotManager = serviceHost.getSnapshotManager(registory);

  // write src/index.ts and check types
  const raw = snapshotManager.readFileSnapshot("src/index.ts");
  const newSource = snapshotManager.writeFileSnapshot(
    "src/index.ts",
    raw + "\nconst y: number = x;",
  );

  // find scoped variables
  const checker = languageService.getProgram()!.getTypeChecker();
  const localVariables = checker.getSymbolsInScope(
    newSource,
    ts.SymbolFlags.BlockScopedVariable,
  );

  // rename x to x_?
  const symbol = localVariables.find((s) => s.name === "x")!;
  const renameLocations = languageService.findRenameLocations(
    "src/index.ts",
    symbol.valueDeclaration!.getStart(),
    false,
    false,
    prefs,
  );
  const targets = new Set(renameLocations!.map((loc) => loc.fileName));

  let current = snapshotManager.readFileSnapshot("src/index.ts")!;
  for (const target of targets) {
    const renameLocationsToTarget = renameLocations!.filter(
      (loc) => expandPath(target) === expandPath(loc.fileName),
    );
    const newSymbolName = `${symbol.name}_${
      Math.random().toString(36).slice(2)
    }`;
    current = applyRenameLocations(
      current,
      newSymbolName,
      renameLocationsToTarget,
    );
  }
  snapshotManager.writeFileSnapshot("src/index.ts", current);
  const result = languageService.getSemanticDiagnostics("src/index.ts");
  console.log("post error", result.length);
  console.log(snapshotManager.readFileSnapshot("src/index.ts"));
}