@computekit/core 0.1.3 → 0.2.0

package/src/index.ts

@@ -11,6 +11,14 @@ import type {
   ComputeKitEvents,
 } from './types';
 
+import type {
+  ComputeFunctionRegistry,
+  RegisteredFunctionName,
+  FunctionInput,
+  FunctionOutput,
+  ComputeFn,
+} from './registry';
+
 import { WorkerPool } from './pool';
 import { EventEmitter, isWasmSupported, createLogger } from './utils';
 
@@ -60,48 +68,79 @@ export class ComputeKit extends EventEmitter<ComputeKitEvents> {
   /**
    * Register a compute function
    *
-   * @param name - Unique name for the function
+   * @param name - Unique name for the function (autocompletes if registry is extended)
    * @param fn - The function to execute (will run in a Web Worker)
    *
    * @example
    * ```ts
+   * // Basic usage
    * kit.register('sum', (arr: number[]) => arr.reduce((a, b) => a + b, 0));
+   *
+   * // With typed registry (extend ComputeFunctionRegistry for autocomplete)
+   * // declare module '@computekit/core' {
+   * //   interface ComputeFunctionRegistry {
+   * //     sum: { input: number[]; output: number };
+   * //   }
+   * // }
+   * // kit.register('sum', (arr) => arr.reduce((a, b) => a + b, 0));
    * ```
    */
-  register<TInput, TOutput>(
-    name: string,
-    fn: (input: TInput) => TOutput | Promise<TOutput>
+  register<
+    TName extends RegisteredFunctionName,
+    TInput = FunctionInput<TName extends string ? TName : never>,
+    TOutput = FunctionOutput<TName extends string ? TName : never>,
+  >(
+    name: TName,
+    fn: TName extends keyof ComputeFunctionRegistry
+      ? ComputeFn<
+          ComputeFunctionRegistry[TName]['input'],
+          ComputeFunctionRegistry[TName]['output']
+        >
+      : ComputeFn<TInput, TOutput>
   ): this {
-    this.pool.register(name, fn);
+    this.pool.register(name as string, fn as ComputeFn<unknown, unknown>);
     return this;
   }
 
   /**
    * Execute a registered compute function
    *
-   * @param name - Name of the registered function
-   * @param input - Input data for the function
+   * @param name - Name of the registered function (autocompletes if registry is extended)
+   * @param input - Input data for the function (type-safe if registry is extended)
    * @param options - Execution options
-   * @returns Promise resolving to the function result
+   * @returns Promise resolving to the function result (type-safe if registry is extended)
    *
    * @example
    * ```ts
    * const sum = await kit.run('sum', [1, 2, 3, 4, 5]);
+   *
+   * // With typed registry, input/output types are inferred:
+   * // const result = await kit.run('fibonacci', 50); // result: number
    * ```
    */
-  async run<TInput, TOutput>(
-    name: string,
-    input: TInput,
+  async run<
+    TName extends RegisteredFunctionName,
+    TInput = FunctionInput<TName extends string ? TName : never>,
+    TOutput = FunctionOutput<TName extends string ? TName : never>,
+  >(
+    name: TName,
+    input: TName extends keyof ComputeFunctionRegistry
+      ? ComputeFunctionRegistry[TName]['input']
+      : TInput,
     options?: ComputeOptions
-  ): Promise<TOutput> {
-    return this.pool.execute<TInput, TOutput>(name, input, options);
+  ): Promise<
+    TName extends keyof ComputeFunctionRegistry
+      ? ComputeFunctionRegistry[TName]['output']
+      : TOutput
+  > {
+    return this.pool.execute(name as string, input, options);
   }
 
   /**
    * Execute a registered compute function with full result metadata
    *
-   * @param name - Name of the registered function
-   * @param input - Input data for the function
+   * @param name - Name of the registered function (autocompletes if registry is extended)
+   * @param input - Input data for the function (type-safe if registry is extended)
    * @param options - Execution options
    * @returns Promise resolving to ComputeResult with metadata
    *
@@ -111,13 +150,33 @@ export class ComputeKit extends EventEmitter<ComputeKitEvents> {
    * console.log(`Took ${result.duration}ms`);
    * ```
    */
-  async runWithMetadata<TInput, TOutput>(
-    name: string,
-    input: TInput,
+  async runWithMetadata<
+    TName extends RegisteredFunctionName,
+    TInput = FunctionInput<TName extends string ? TName : never>,
+    TOutput = FunctionOutput<TName extends string ? TName : never>,
+  >(
+    name: TName,
+    input: TName extends keyof ComputeFunctionRegistry
+      ? ComputeFunctionRegistry[TName]['input']
+      : TInput,
     options?: ComputeOptions
-  ): Promise<ComputeResult<TOutput>> {
+  ): Promise<
+    ComputeResult<
+      TName extends keyof ComputeFunctionRegistry
+        ? ComputeFunctionRegistry[TName]['output']
+        : TOutput
+    >
+  > {
+    type ActualOutput = TName extends keyof ComputeFunctionRegistry
+      ? ComputeFunctionRegistry[TName]['output']
+      : TOutput;
+
     const startTime = performance.now();
-    const data = await this.pool.execute<TInput, TOutput>(name, input, options);
+    const data = (await this.pool.execute(
+      name as string,
+      input,
+      options
+    )) as ActualOutput;
     const duration = performance.now() - startTime;
 
     return {
@@ -175,23 +234,64 @@ export function getDefaultInstance(): ComputeKit {
 
 /**
  * Register a function on the default instance
+ *
+ * @example
+ * ```ts
+ * import { register } from '@computekit/core';
+ *
+ * register('fibonacci', (n: number) => {
+ *   if (n <= 1) return n;
+ *   let a = 0, b = 1;
+ *   for (let i = 2; i <= n; i++) {
+ *     [a, b] = [b, a + b];
+ *   }
+ *   return b;
+ * });
+ * ```
  */
-export function register<TInput, TOutput>(
-  name: string,
-  fn: (input: TInput) => TOutput | Promise<TOutput>
+export function register<
+  TName extends RegisteredFunctionName,
+  TInput = FunctionInput<TName extends string ? TName : never>,
+  TOutput = FunctionOutput<TName extends string ? TName : never>,
+>(
+  name: TName,
+  fn: TName extends keyof ComputeFunctionRegistry
+    ? ComputeFn<
+        ComputeFunctionRegistry[TName]['input'],
+        ComputeFunctionRegistry[TName]['output']
+      >
+    : ComputeFn<TInput, TOutput>
 ): void {
   getDefaultInstance().register(name, fn);
 }
 
 /**
  * Run a function on the default instance
+ *
+ * @example
+ * ```ts
+ * import { run } from '@computekit/core';
+ *
+ * const result = await run('fibonacci', 50);
+ * console.log(result); // Type is inferred if registry is extended
+ * ```
  */
-export async function run<TInput, TOutput>(
-  name: string,
-  input: TInput,
+export async function run<
+  TName extends RegisteredFunctionName,
+  TInput = FunctionInput<TName extends string ? TName : never>,
+  TOutput = FunctionOutput<TName extends string ? TName : never>,
+>(
+  name: TName,
+  input: TName extends keyof ComputeFunctionRegistry
+    ? ComputeFunctionRegistry[TName]['input']
+    : TInput,
   options?: ComputeOptions
-): Promise<TOutput> {
-  return getDefaultInstance().run<TInput, TOutput>(name, input, options);
+): Promise<
+  TName extends keyof ComputeFunctionRegistry
+    ? ComputeFunctionRegistry[TName]['output']
+    : TOutput
+> {
+  return getDefaultInstance().run(name, input, options);
 }
 
 // Re-export types
@@ -221,6 +321,18 @@ export type {
   ParallelBatchResult,
 } from './types';
 
+// Re-export typed registry types
+export type {
+  ComputeFunctionRegistry,
+  RegisteredFunctionName,
+  FunctionInput,
+  FunctionOutput,
+  ComputeFn,
+  InferComputeFn,
+  DefineFunction,
+  HasRegisteredFunctions,
+} from './registry';
+
 // Re-export utilities
 export {
   isWasmSupported,

package/src/registry.ts

@@ -0,0 +1,132 @@
+/**
+ * ComputeKit Typed Registry
+ *
+ * This module provides type-safe function registration and execution.
+ * Users can extend the ComputeFunctionRegistry interface to get autocomplete
+ * and type safety for their registered functions.
+ *
+ * @example
+ * ```ts
+ * // Extend the registry interface (in a .d.ts file or at the top of your file)
+ * declare module '@computekit/core' {
+ *   interface ComputeFunctionRegistry {
+ *     fibonacci: { input: number; output: number };
+ *     sum: { input: number[]; output: number };
+ *     processData: { input: { items: string[] }; output: { count: number } };
+ *   }
+ * }
+ *
+ * // Now you get autocomplete and type safety!
+ * const kit = new ComputeKit();
+ * kit.register('fibonacci', (n) => ...);  // n is inferred as number
+ * const result = await kit.run('fibonacci', 42);  // result is number
+ * ```
+ */
+
+/**
+ * Registry interface for compute functions.
+ * Extend this interface using module augmentation to add your own functions.
+ *
+ * Each entry should be in the format:
+ * ```ts
+ * functionName: { input: InputType; output: OutputType }
+ * ```
+ *
+ * @example
+ * ```ts
+ * declare module '@computekit/core' {
+ *   interface ComputeFunctionRegistry {
+ *     myFunction: { input: string; output: number };
+ *   }
+ * }
+ * ```
+ */
+// eslint-disable-next-line @typescript-eslint/no-empty-object-type
+export interface ComputeFunctionRegistry {}
+
+/**
+ * Helper type to get all registered function names.
+ * If no functions are registered, falls back to string.
+ */
+export type RegisteredFunctionName = keyof ComputeFunctionRegistry extends never
+  ? string
+  : keyof ComputeFunctionRegistry;
+
+/**
+ * Helper type to check if the registry has any entries.
+ */
+export type HasRegisteredFunctions = keyof ComputeFunctionRegistry extends never
+  ? false
+  : true;
+
+/**
+ * Get the input type for a registered function.
+ * Falls back to TFallback if the function is not registered.
+ */
+export type FunctionInput<
+  TName extends string,
+  TFallback = unknown,
+> = TName extends keyof ComputeFunctionRegistry
+  ? ComputeFunctionRegistry[TName]['input']
+  : TFallback;
+
+/**
+ * Get the output type for a registered function.
+ * Falls back to TFallback if the function is not registered.
+ */
+export type FunctionOutput<
+  TName extends string,
+  TFallback = unknown,
+> = TName extends keyof ComputeFunctionRegistry
+  ? ComputeFunctionRegistry[TName]['output']
+  : TFallback;
+
+/**
+ * Type for a compute function based on registry or explicit types.
+ */
+export type ComputeFn<TInput, TOutput> = (input: TInput) => TOutput | Promise<TOutput>;
+
+/**
+ * Infer the compute function type for a registered function name.
+ */
+export type InferComputeFn<TName extends string> =
+  TName extends keyof ComputeFunctionRegistry
+    ? ComputeFn<
+        ComputeFunctionRegistry[TName]['input'],
+        ComputeFunctionRegistry[TName]['output']
+      >
+    : ComputeFn<unknown, unknown>;
+
+/**
+ * Type helper for creating registry entries.
+ * Use this to define your function types more easily.
+ *
+ * @example
+ * ```ts
+ * declare module '@computekit/core' {
+ *   interface ComputeFunctionRegistry {
+ *     fibonacci: DefineFunction<number, number>;
+ *     sum: DefineFunction<number[], number>;
+ *   }
+ * }
+ * ```
+ */
+export type DefineFunction<TInput, TOutput> = {
+  input: TInput;
+  output: TOutput;
+};
+
+/**
+ * Helper type for typing the register method.
+ * Provides proper type inference based on whether the function is in the registry.
+ */
+export type RegisterFn<
+  TName extends string,
+  TInput,
+  TOutput,
+> = TName extends keyof ComputeFunctionRegistry
+  ? ComputeFn<
+      ComputeFunctionRegistry[TName]['input'],
+      ComputeFunctionRegistry[TName]['output']
+    >
+  : ComputeFn<TInput, TOutput>;