Cut evaluation plugin overhead and debounce diagnostics

language-server/src/build-server.ts

@@ -26,7 +26,11 @@ import type { Connection } from "vscode-languageserver";
 export type LanguageServerSettings = {
 };
 
-export const buildServer = (connection: Connection): Server => {
+export type BuildServerOptions = {
+  diagnosticsDelay?: number;
+};
+
+export const buildServer = (connection: Connection, options: BuildServerOptions = {}): Server => {
   const server = new Server(connection);
   const workspace = new Workspace(server);
 
@@ -44,7 +48,7 @@ export const buildServer = (connection: Connection): Server => {
     new SchemaRegistrationDiagnosticsProvider(registry, jsonSchema),
     new SchemaValidationDiagnosticsProvider(jsonSchema),
     new DeprecatedDiagnosticsProvider(server, jsonSchema)
-  ]);
+  ], options.diagnosticsDelay);
 
   new Formatting(server, jsonDocuments);
   new Hover(server, jsonDocuments, jsonSchema);

language-server/src/evaluation/LspEvaluationPlugin.ts

@@ -1,7 +1,6 @@
 import { getKeyword } from "@hyperjump/json-schema/experimental";
 import * as Instance from "@hyperjump/json-schema/instance/experimental";
 import * as JsonPointer from "@hyperjump/json-pointer";
-import * as Pact from "@hyperjump/pact";
 import { JsonValueSet } from "./JsonValueSet.ts";
 import { SubschemaTracker } from "./SubschemaTracker.ts";
 import { Annotation } from "./Annotation.ts";
@@ -52,10 +51,22 @@ export class LspEvaluationPlugin implements EvaluationPlugin<LspContext> {
   private annotations: RecordedAnnotations = Object.create(null);
   private completions: Completions = Object.create(null);
   private incompleteLocations: Set<string>;
+  // Incomplete locations by the pointer of their parent so the keywords that
+  // apply to an object only look at its own incomplete properties
+  private incompleteChildren = new Map<string, [string, string][]>();
   private subschemaTracker = new SubschemaTracker<SchemaResult>();
 
   constructor(incompleteLocations: Set<string> = new Set()) {
     this.incompleteLocations = incompleteLocations;
+    for (const pointer of incompleteLocations) {
+      const [parentPointer, propertyName] = splitPointer(pointer);
+      let children = this.incompleteChildren.get(parentPointer);
+      if (!children) {
+        children = [];
+        this.incompleteChildren.set(parentPointer, children);
+      }
+      children.push([pointer, propertyName]);
+    }
   }
 
   // The plugin is registered once, in build-server.ts
@@ -107,9 +118,8 @@ export class LspEvaluationPlugin implements EvaluationPlugin<LspContext> {
             }
           }
 
-          for (const pointer of this.incompleteLocations) {
-            const [parentPointer, propertyName] = splitPointer(pointer);
-            if (parentPointer === instance.pointer && pattern.test(propertyName)) {
+          for (const [pointer, propertyName] of this.incompleteChildren.get(instance.pointer) ?? []) {
+            if (pattern.test(propertyName)) {
               addCompletions(schemaContext.completions!, pointer, completions);
               context.evaluatedProperties?.add(propertyName);
             }
@@ -129,9 +139,8 @@ export class LspEvaluationPlugin implements EvaluationPlugin<LspContext> {
           }
         }
 
-        for (const pointer of this.incompleteLocations) {
-          const [parentPointer, propertyName] = splitPointer(pointer);
-          if (parentPointer === instance.pointer && !isDeclaredProperty.test(propertyName)) {
+        for (const [pointer, propertyName] of this.incompleteChildren.get(instance.pointer) ?? []) {
+          if (!isDeclaredProperty.test(propertyName)) {
             addCompletions(schemaContext.completions!, pointer, completions);
             context.evaluatedProperties?.add(propertyName);
           }
@@ -150,9 +159,8 @@ export class LspEvaluationPlugin implements EvaluationPlugin<LspContext> {
           }
         }
 
-        for (const pointer of this.incompleteLocations) {
-          const [parentPointer, propertyName] = splitPointer(pointer);
-          if (parentPointer === instance.pointer && !context.schemaEvaluatedProperties!.has(propertyName)) {
+        for (const [pointer, propertyName] of this.incompleteChildren.get(instance.pointer) ?? []) {
+          if (!context.schemaEvaluatedProperties!.has(propertyName)) {
             addCompletions(schemaContext.completions!, pointer, completions);
             context.evaluatedProperties?.add(propertyName);
           }
@@ -275,7 +283,7 @@ export class LspEvaluationPlugin implements EvaluationPlugin<LspContext> {
 
     const isConditional = conditionalKeywords.has(keywordId);
     for (const { pointer, annotations } of annotationValues) {
-      appendAnnotations(schemaContext.schemaAnnotations!, isConditional ? asValueAnnotations(annotations, pointer) : annotations);
+      appendAnnotations(schemaContext.schemaAnnotations!, annotations, isConditional ? pointer : undefined);
     }
 
     // The schema's completions are only shared once the schema is done, so
@@ -363,12 +371,13 @@ export class LspEvaluationPlugin implements EvaluationPlugin<LspContext> {
   }
 
   private getAnnotationsOfKind(instanceLocation: string, kind: RecordedAnnotation["kind"]): AnnotationRecord[] {
-    return Pact.pipe(
-      this.annotations[instanceLocation],
-      Pact.filter((recorded) => recorded.kind === kind),
-      Pact.map(({ annotation }) => annotation),
-      Pact.collectArray
-    );
+    const annotations: AnnotationRecord[] = [];
+    for (const recorded of this.annotations[instanceLocation]) {
+      if (recorded.kind === kind) {
+        annotations.push(recorded.annotation);
+      }
+    }
+    return annotations;
   }
 
   * getCompletions(pointer: string): Generator<Completion> {
@@ -564,38 +573,23 @@ function negateInto(target: Completions, source: Completions) {
 }
 
 function exclusiveUnion(records: Completions[]): Completions {
-  const pointers = new Set<string>();
+  const setsByPointer: Record<string, LazyValueSet[]> = Object.create(null);
   for (const record of records) {
     for (const pointer in record) {
-      pointers.add(pointer);
+      setsByPointer[pointer] ??= [];
+      setsByPointer[pointer].push(record[pointer]);
     }
   }
 
   const result: Completions = Object.create(null);
-  for (const pointer of pointers) {
-    const sets = Pact.pipe(
-      records,
-      Pact.map((record) => record[pointer]),
-      Pact.filter((set) => set !== undefined),
-      Pact.collectArray
-    );
+  for (const pointer in setsByPointer) {
+    const sets = setsByPointer[pointer];
     result[pointer] = lazy(() => JsonValueSet.exclusiveUnion(sets.map((set) => set())));
   }
 
   return result;
 }
 
-function asValueAnnotations(annotations: RecordedAnnotations, pointer: string): RecordedAnnotations {
-  if (!(pointer in annotations)) {
-    return annotations;
-  }
-
-  return {
-    ...annotations,
-    [pointer]: annotations[pointer].map(({ annotation }) => ({ kind: "value", annotation }))
-  };
-}
-
 // Adds a keyword's annotation without an instance, for locations that
 // evaluation doesn't reach.
 // Returns whether the keyword contributed an annotation
@@ -616,10 +610,17 @@ function toAnnotations(records: AnnotationRecord[]): Annotation[] {
   return records.map((record) => new Annotation(record));
 }
 
-function appendAnnotations(target: RecordedAnnotations, source: RecordedAnnotations) {
+// Annotations at `valuePointer` are appended as value annotations
+function appendAnnotations(target: RecordedAnnotations, source: RecordedAnnotations, valuePointer?: string) {
   for (const pointer in source) {
     target[pointer] ??= [];
-    target[pointer].push(...source[pointer]);
+    if (pointer === valuePointer) {
+      for (const { annotation } of source[pointer]) {
+        target[pointer].push({ kind: "value", annotation });
+      }
+    } else {
+      target[pointer].push(...source[pointer]);
+    }
   }
 }
 

language-server/src/evaluation/SubschemaTracker.ts

@@ -14,6 +14,12 @@ type SchemaFailures = {
   failedType: boolean;
 };
 
+type TrackerState<T> = {
+  failedLocations?: Set<string>;
+  schemaFailures?: SchemaFailures;
+  subschemaResults?: SubschemaResult<T>[];
+};
+
 /**
  * Tracks the result of each subschema evaluated by a keyword and the instance
  * locations where it failed. This allows anyOf/oneOf alternatives to be
@@ -23,30 +29,32 @@ type SchemaFailures = {
  * separately because an object or array is discriminated by its type, not by
  * any other failure.
  *
- * State is keyed by context object rather than stored on the context so that
- * plugins using separate trackers don't share state.
+ * State is stored on the context under a key that's unique to the tracker so
+ * that plugins using separate trackers don't share state. A property is much
+ * faster than a WeakMap for something that happens for every keyword.
  */
 export class SubschemaTracker<T> {
-  private failedLocations = new WeakMap<object, Set<string>>();
-  private schemaFailures = new WeakMap<object, SchemaFailures>();
-  private subschemaResults = new WeakMap<object, SubschemaResult<T>[]>();
+  private key = Symbol("SubschemaTracker");
+
+  private state(context: object): TrackerState<T> {
+    return (context as Record<symbol, TrackerState<T>>)[this.key] ??= {};
+  }
 
   beforeSchema(context: object): void {
-    if (!this.failedLocations.has(context)) {
-      this.failedLocations.set(context, new Set());
-    }
-    this.schemaFailures.set(context, { failedLocations: new Set(), failedType: false });
+    const state = this.state(context);
+    state.failedLocations ??= new Set();
+    state.schemaFailures = { failedLocations: new Set(), failedType: false };
   }
 
   beforeKeyword(keywordContext: object): void {
-    this.subschemaResults.set(keywordContext, []);
+    this.state(keywordContext).subschemaResults = [];
   }
 
   afterKeyword(keywordId: string, keywordValue: unknown, instance: JsonNode, keywordContext: object, valid: boolean, schemaContext: object): void {
     if (!valid) {
-      const schemaFailures = this.schemaFailures.get(schemaContext)!;
+      const schemaFailures = this.state(schemaContext).schemaFailures!;
       schemaFailures.failedLocations.add(instance.pointer);
-      for (const location of this.failedLocations.get(keywordContext) ?? []) {
+      for (const location of this.state(keywordContext).failedLocations ?? []) {
         schemaFailures.failedLocations.add(location);
       }
 
@@ -57,7 +65,8 @@ export class SubschemaTracker<T> {
   }
 
   afterSchema(instance: JsonNode, context: object, valid: boolean, value: T): void {
-    const schemaFailures = this.schemaFailures.get(context)!;
+    const state = this.state(context);
+    const schemaFailures = state.schemaFailures!;
     if (!valid) {
       // A schema that fails without any of its keywords failing is a `false`
       // schema, which rules out every type
@@ -67,12 +76,12 @@ export class SubschemaTracker<T> {
       schemaFailures.failedLocations.add(instance.pointer);
     }
 
-    const failedLocations = this.failedLocations.get(context)!;
+    const failedLocations = state.failedLocations!;
     for (const location of schemaFailures.failedLocations) {
       failedLocations.add(location);
     }
 
-    this.subschemaResults.get(context)?.push({
+    state.subschemaResults?.push({
       value,
       pointer: instance.pointer,
       failedLocations: schemaFailures.failedLocations,
@@ -81,7 +90,7 @@ export class SubschemaTracker<T> {
   }
 
   getSubschemaResults(keywordContext: object): SubschemaResult<T>[] {
-    return this.subschemaResults.get(keywordContext) ?? [];
+    return this.state(keywordContext).subschemaResults ?? [];
   }
 
   /**

language-server/src/features/DocumentColors.ts

@@ -38,21 +38,21 @@ export class DocumentColors {
 
         const plugin = LspEvaluationPlugin.from(result);
 
-        const stringNodes: Node[] = [];
-        jsonDocument.walkNodes(ast, (node) => {
+        const stringNodes: [Node, string][] = [];
+        jsonDocument.walkNodes(ast, (node, pointer) => {
           if (node.type === "string") {
-            stringNodes.push(node);
+            stringNodes.push([node, pointer]);
           }
         });
 
         const colors: ColorInformation[] = [];
-        for (const node of stringNodes) {
+        for (const [node, pointer] of stringNodes) {
           const color = colorFromHex(node.value as string);
           if (!color) {
             continue;
           }
 
-          const annotations = plugin.getAnnotations(jsonDocument.getPointerForNode(node));
+          const annotations = plugin.getAnnotations(pointer);
           if (annotations.some(isColorHex)) {
             colors.push({ color, range: jsonDocument.rangeAt(node.offset, node.offset + node.length) });
           }

language-server/src/features/diagnostics/Diagnostics.test.ts

@@ -3,6 +3,7 @@ import { readFile } from "node:fs/promises";
 import { fileURLToPath } from "node:url";
 import { TestClient } from "../../test/TestClient.ts";
 import { ReadFileRequest } from "../../protocol/hyperjump-readFile.ts";
+import { wait } from "../../test/test-utils.ts";
 
 import type { Diagnostic } from "vscode-languageserver";
 
@@ -102,3 +103,70 @@ describe("Diagnostics", () => {
     await expect(reopenDiagnostics).resolves.toHaveLength(1);
   });
 });
+
+describe("Diagnostics delay", () => {
+  const diagnosticsDelay = 200;
+  // Long enough to validate, but well within the delay
+  const pause = async () => wait(diagnosticsDelay / 5);
+  let client: TestClient;
+
+  beforeEach(async () => {
+    client = new TestClient(undefined, { diagnosticsDelay });
+    await client.start();
+  });
+
+  afterEach(async () => {
+    await client.stop();
+  });
+
+  test("opening a document publishes diagnostics without waiting for the delay", async () => {
+    await client.writeDocument("test.json", `{ "name": }`);
+    const initialDiagnostics = client.getDiagnostics("test.json");
+
+    const start = performance.now();
+    await client.openDocument("test.json");
+    await expect(initialDiagnostics).resolves.toHaveLength(1);
+
+    expect(performance.now() - start).toBeLessThan(diagnosticsDelay);
+  });
+
+  test("changes in quick succession publish diagnostics once, for the last change", async () => {
+    await client.writeDocument("test.json", `{ "name": "foo" }`);
+    const initialDiagnostics = client.getDiagnostics("test.json");
+    await client.openDocument("test.json");
+    await expect(initialDiagnostics).resolves.toEqual([]);
+
+    const receivedDiagnostics: Diagnostic[][] = [];
+    await client.onDiagnostics("test.json", (diagnostics) => {
+      receivedDiagnostics.push(diagnostics);
+    });
+
+    await client.changeDocument("test.json", `{ "name": }`);
+    await pause();
+    await client.changeDocument("test.json", `{ "name": "f }`);
+    await pause();
+    await client.changeDocument("test.json", `{ "name": "foo" }`);
+    await wait(diagnosticsDelay * 4);
+
+    expect(receivedDiagnostics).toEqual([[]]);
+  });
+
+  test("closing a document before the delay is over publishes only the empty list", async () => {
+    await client.writeDocument("test.json", `{ "name": "foo" }`);
+    const initialDiagnostics = client.getDiagnostics("test.json");
+    await client.openDocument("test.json");
+    await expect(initialDiagnostics).resolves.toEqual([]);
+
+    const receivedDiagnostics: Diagnostic[][] = [];
+    await client.onDiagnostics("test.json", (diagnostics) => {
+      receivedDiagnostics.push(diagnostics);
+    });
+
+    await client.changeDocument("test.json", `{ "name": }`);
+    await pause();
+    await client.closeDocument("test.json");
+    await wait(diagnosticsDelay * 4);
+
+    expect(receivedDiagnostics).toEqual([[]]);
+  });
+});