TreeSitterCodeParser

export class TreeSitterCodeParser implements CodeParser {
  readonly id: string;
  readonly languages: readonly string[];
  readonly #engine: IncrementalSyntaxEngine;
  readonly #trees = new WeakMap<CodeParseResult, IncrementalSyntaxTree>();
  #version = 0;

  constructor(
    id: string,
    languages: readonly string[],
    engine: IncrementalSyntaxEngine,
  ) {
    this.id = id;
    this.languages = Object.freeze([...languages]);
    this.#engine = engine;
  }

  async parse(
    source: string,
    options: {
      readonly previous?: CodeParseResult;
      readonly signal: AbortSignal;
    },
  ): Promise<CodeParseResult> {
    if (options.signal.aborted) throw options.signal.reason;
    const tree = await this.#engine.parse(
      source,
      options.previous ? this.#trees.get(options.previous) : undefined,
      options.signal,
    );
    if (options.signal.aborted) throw options.signal.reason;
    const nodes = tree.nodes;
    const result = Object.freeze({
      language: this.languages[0] ?? "text",
      version: ++this.#version,
      spans: Object.freeze(
        nodes.map((node) => ({
          start: node.start,
          end: node.end,
          kind: node.type,
        })),
      ),
      folds: Object.freeze([
        ...new Map(
          nodes.flatMap((node) => {
            if (
              node.startLine === undefined ||
              node.endLine === undefined ||
              node.endLine <= node.startLine
            ) {
              return [];
            }
            return [
              [
                `${node.startLine}:${node.endLine}`,
                {
                  startLine: node.startLine,
                  endLine: node.endLine,
                },
              ] as const,
            ];
          }),
        ).values(),
      ]),
      diagnostics: Object.freeze(
        nodes
          .filter(
            (
              node,
            ): node is IncrementalSyntaxNode & {
              readonly diagnostic: NonNullable<
                IncrementalSyntaxNode["diagnostic"]
              >;
            } => Boolean(node.diagnostic),
          )
          .map((node) => ({
            start: node.start,
            end: node.end,
            kind: node.type,
            ...node.diagnostic,
          })),
      ),
      incremental: tree.incremental,
    });
    this.#trees.set(result, tree);
    return result;
  }

  dispose(): void {
    this.#engine.dispose?.();
  }
}