Simplify type argument to CachedOperation

extensions/ql-vscode/src/language-support/contextual/cached-operation.ts

@@ -1,31 +1,67 @@
 import { asError } from "../../common/helpers-pure";
 
 /**
- * A cached mapping from args of type [string, S] to a value of type Promise<U>.
+ * A type that wraps the cached operation function to avoid the use of
+ * any in the args.
  */
-export class CachedOperation<S extends unknown[], U> {
-  private readonly operation: (t: string, ...args: S) => Promise<U>;
-  private readonly cached: Map<string, U>;
+type CachedOperationFunction<F> = F extends (
+  t: string,
+  ...args: infer A
+) => Promise<infer U>
+  ? (t: string, ...args: A) => Promise<U>
+  : never;
+
+/**
+ * The value to store in the cache for a cached operation.
+ */
+type CachedOperationValue<
+  F extends (t: string, ...args: unknown[]) => Promise<unknown>,
+> = Awaited<ReturnType<F>>;
+
+/**
+ * All parameters of a function except the first one, which must be a string.
+ */
+type CachedOperationArgs<
+  F extends (t: string, ...args: unknown[]) => Promise<unknown>,
+> = Parameters<F> extends [string, ...infer T] ? T : never;
+
+/**
+ * A cached mapping from args of type [string, S] to a value of type Promise<U>.
+ *
+ * F1 needs to be supplied as the type of the function that is being cached,
+ * for example by using `typeof myFunction`. F1 always accepts all arguments,
+ * but if it doesn't match the shape of `CachedOperationFunction` then the constructor
+ * argument will be inferred as `never`. This is because this is the only way to prevent
+ * the use of any in the `args` parameter of the `extends` type.
+ */
+export class CachedOperation<
+  F1,
+  F extends CachedOperationFunction<F1> = CachedOperationFunction<F1>,
+> {
+  private readonly cached: Map<string, CachedOperationValue<F>>;
   private readonly lru: string[];
   private readonly inProgressCallbacks: Map<
     string,
-    Array<[(u: U) => void, (reason?: Error) => void]>
+    Array<[(u: CachedOperationValue<F>) => void, (reason?: Error) => void]>
   >;
 
   constructor(
-    operation: (t: string, ...args: S) => Promise<U>,
+    private readonly operation: F,
     private cacheSize = 100,
   ) {
     this.operation = operation;
     this.lru = [];
     this.inProgressCallbacks = new Map<
       string,
-      Array<[(u: U) => void, (reason?: Error) => void]>
+      Array<[(u: CachedOperationValue<F>) => void, (reason?: Error) => void]>
     >();
-    this.cached = new Map<string, U>();
+    this.cached = new Map();
   }
 
-  async get(t: string, ...args: S): Promise<U> {
+  async get(
+    t: string,
+    ...args: CachedOperationArgs<F>
+  ): Promise<CachedOperationValue<F>> {
     // Try and retrieve from the cache
     const fromCache = this.cached.get(t);
     if (fromCache !== undefined) {
@@ -42,16 +78,21 @@ export class CachedOperation<S extends unknown[], U> {
     const inProgressCallback = this.inProgressCallbacks.get(t);
     if (inProgressCallback !== undefined) {
       // If so wait for it to resolve
-      return await new Promise((resolve, reject) => {
+      return await new Promise<CachedOperationValue<F>>((resolve, reject) => {
         inProgressCallback.push([resolve, reject]);
       });
     }
 
     // Otherwise compute the new value, but leave a callback to allow sharing work
-    const callbacks: Array<[(u: U) => void, (reason?: Error) => void]> = [];
+    const callbacks: Array<
+      [(u: CachedOperationValue<F>) => void, (reason?: Error) => void]
+    > = [];
     this.inProgressCallbacks.set(t, callbacks);
     try {
-      const result = await this.operation(t, ...args);
+      const result = (await this.operation(
+        t,
+        ...args,
+      )) as CachedOperationValue<F>;
       callbacks.forEach((f) => f[0](result));
       this.inProgressCallbacks.delete(t);
       if (this.lru.length > this.cacheSize) {

extensions/ql-vscode/src/language-support/contextual/template-provider.ts

@@ -50,7 +50,7 @@ import { MultiCancellationToken } from "../../common/vscode/multi-cancellation-t
  */
 
 export class TemplateQueryDefinitionProvider implements DefinitionProvider {
-  private cache: CachedOperation<[CancellationToken], LocationLink[]>;
+  private cache: CachedOperation<typeof this.getDefinitions>;
 
   constructor(
     private cli: CodeQLCliServer,
@@ -110,7 +110,7 @@ export class TemplateQueryDefinitionProvider implements DefinitionProvider {
  * or from a selected identifier.
  */
 export class TemplateQueryReferenceProvider implements ReferenceProvider {
-  private cache: CachedOperation<[CancellationToken], FullLocationLink[]>;
+  private cache: CachedOperation<typeof this.getReferences>;
 
   constructor(
     private cli: CodeQLCliServer,
@@ -181,10 +181,7 @@ export class TemplateQueryReferenceProvider implements ReferenceProvider {
  * source-language files.
  */
 export class TemplatePrintAstProvider {
-  private cache: CachedOperation<
-    [ProgressCallback, CancellationToken],
-    CoreCompletedQuery
-  >;
+  private cache: CachedOperation<typeof this.getAst>;
 
   constructor(
     private cli: CodeQLCliServer,
@@ -280,10 +277,7 @@ export class TemplatePrintAstProvider {
  * source-language files.
  */
 export class TemplatePrintCfgProvider {
-  private cache: CachedOperation<
-    [number, number],
-    [Uri, Record<string, string>]
-  >;
+  private cache: CachedOperation<typeof this.getCfgUri>;
 
   constructor(
     private cli: CodeQLCliServer,