feat: preserves overloads for non-async functions (#22)

* feat: preserves overloads for non-async functions

* chore: format types

* chore: use type cast on trace newFn

* chore: move declarations to trace demo file

* chore: use type cast for memoize

* chore: add memoize demo for overloaded functions

* chore: use type cast for als

* test: update als test type for incBy

---------

Co-authored-by: Rexford Essilfie <[email protected]>

demos/memoize.ts

@@ -19,4 +19,29 @@ export function demo() {
   console.log(tracedFib(10));
 }
 
+// Demo with overloaded functions
+
+function myFn(a: boolean): string | number;
+function myFn<T, U>(a: T, b: U): T | U;
+
+function myFn(...args: unknown[]) {
+  if (!Array.isArray(args)) {
+    throw new Error("Invalid arguments");
+  }
+
+  if (typeof args[0] === "boolean") {
+    return "hello";
+  }
+
+  return args[0];
+}
+
+export function demo2() {
+  console.log("demo2");
+  const memoizedMyFn = memoize((...args: any[]) => args.slice().join())(myFn);
+  console.log(memoizedMyFn(5, 6));
+  console.log(memoizedMyFn(false));
+}
+
 demo();
+demo2();

demos/trace.ts

@@ -11,4 +11,28 @@ export function demo() {
   console.log(tracedAdd(3, 5));
 }
 
+// Demo with overloaded functions
+
+function myFn(a: boolean): string | number;
+function myFn<T, U>(a: T, b: U): T | U;
+
+function myFn(...args: unknown[]) {
+  if (!Array.isArray(args)) {
+    throw new Error("Invalid arguments");
+  }
+
+  if (typeof args[0] === "boolean") {
+    return "hello";
+  }
+
+  return args[0];
+}
+
+export function demo2() {
+  const tracedMyFn = trace(myFn);
+  tracedMyFn(5, 6);
+  tracedMyFn(false);
+}
+
 demo();
+demo2();

src/common/als.ts

@@ -1,14 +1,13 @@
+import type { Any } from "src/common/types";
 import { AsyncLocalStorage } from "node:async_hooks";
 
 export function als<T, IArgs extends any[]>(
   storage: AsyncLocalStorage<T>,
   init: (...args: IArgs) => T
 ) {
-  return function <FArgs extends IArgs, FReturn>(
-    fn: (...args: FArgs) => FReturn
-  ) {
+  return function <Fn extends Any.Function>(fn: Fn) {
     function newFn(...args: Parameters<typeof fn>) {
-      // @ts-ignore TS2683
+      // @ts-expect-error TS2683
       const result = storage.run(init(...args), fn.bind(this), ...args);
       return result;
     }
@@ -20,6 +19,6 @@ export function als<T, IArgs extends any[]>(
       }
     );
 
-    return newFn;
+    return newFn as Fn;
   };
 }

src/common/memoize.ts

@@ -1,13 +1,15 @@
-export function memoize<H extends (...args: any[]) => any>(
+import { Any } from "src/common/types";
+
+export function memoize<H extends Any.Function>(
   hash: H,
   cache = Object.create(null)
 ) {
-  return <FArgs extends any[], FReturn>(fn: (...args: FArgs) => FReturn) => {
+  return <Fn extends Any.Function>(fn: Fn) => {
     function newFn(...args: Parameters<typeof fn>) {
-      // @ts-ignore TS2683
+      // @ts-expect-error TS2683
       const key = hash.apply(this, args);
       if (!(key in cache)) {
-        // @ts-ignore TS2683
+        // @ts-expect-error TS2683
         cache[key] = fn.apply(this, args);
       }
       return cache[key] as ReturnType<typeof fn>;
@@ -20,6 +22,6 @@ export function memoize<H extends (...args: any[]) => any>(
       }
     );
 
-    return newFn;
+    return newFn as Fn;
   };
 }

src/common/trace.ts

@@ -1,10 +1,10 @@
-export const trace = <FArgs extends any[], FReturn>(
-  fn: (...args: FArgs) => FReturn
-) => {
-  function newFn(...args: Parameters<typeof fn>) {
+import type { Any } from "src/common/types";
+
+export const trace = <Fn extends Any.Function>(fn: Fn) => {
+  function newFn(...args: Parameters<Fn>) {
     const tag = `${fn.name}(${args}) | duration`;
     console.time(tag);
-    // @ts-ignore TS2683
+    // @ts-expect-error TS2683
     const result = fn.apply(this, args);
     console.timeEnd(tag);
     return result;
@@ -17,5 +17,5 @@ export const trace = <FArgs extends any[], FReturn>(
     }
   );
 
-  return newFn;
+  return newFn as Fn;
 };

src/common/types.ts

@@ -0,0 +1,9 @@
+export type { Any };
+
+declare namespace Any {
+  type Function<I extends Array<unknown> = Array<any>, O = any> = {
+    (...args: I): O;
+  };
+  type Array<T = unknown> = ReadonlyArray<T>;
+}
+

src/wrapper.ts

@@ -45,17 +45,17 @@ type Spreadable<T> = T extends any[] ? T : never;
 
 type Replace<Source, A, B, Sentinel = never> = Source extends Sentinel
   ? Omit<Source, keyof Sentinel> extends infer T
-    ? Replace<LeftIfNotEqual<T, {}>, A, B, Sentinel> extends infer U
-      ? U | B
-      : never // Dummy ternary just so we can get an alias T
-    : never // Dummy ternary just so we can get an alias U
+  ? Replace<LeftIfNotEqual<T, {}>, A, B, Sentinel> extends infer U
+  ? U | B
+  : never // Dummy ternary just so we can get an alias T
+  : never // Dummy ternary just so we can get an alias U
   : Source extends Promise<infer T>
   ? Promise<Replace<T, A, B, Sentinel>>
   : Source extends [infer AItem, ...infer ARest]
   ? [
-      Replace<AItem, A, B, Sentinel>,
-      ...Spreadable<Replace<ARest, A, B, Sentinel>>
-    ]
+    Replace<AItem, A, B, Sentinel>,
+    ...Spreadable<Replace<ARest, A, B, Sentinel>>
+  ]
   : Source extends Record<string | number | symbol, any>
   ? { [K in keyof Source]: Replace<Source[K], A, B, Sentinel> }
   : Source extends A

test/als.test.ts

@@ -17,7 +17,7 @@ test("maintains reference to this", async (t) => {
     incBy: als(
       storage,
       init
-    )(function (this: any) {
+    )(function (this: any, ..._args: any[]) {
       return this.val + x();
     }),
   };