From d04fb3765a9e025b5153694a6373aa2b362c69b5 Mon Sep 17 00:00:00 2001
From: PalmDevs <git@palmdevs.me>
Date: Thu, 19 Sep 2024 03:18:17 +0700
Subject: [PATCH] feat: better inferred typings

---
 src/getPatchFunc.ts | 12 ++++++------
 src/hook.ts         |  6 +++---
 src/index.ts        | 10 +++-------
 src/shared.ts       | 28 +++++++++++++++++++++-------
 4 files changed, 33 insertions(+), 23 deletions(-)

diff --git a/src/getPatchFunc.ts b/src/getPatchFunc.ts
index 6b2dadb..8b80a6e 100644
--- a/src/getPatchFunc.ts
+++ b/src/getPatchFunc.ts
@@ -3,15 +3,15 @@
 // functional programming strikes again! -- sink
 
 import hook from "./hook";
-import { PatchType, patchedFunctions } from "./shared";
+import { AnyObject, KeysWithFunctionValues, PatchType, PatchTypeToCallbackMap, patchedFunctions } from "./shared";
 import { unpatch } from "./unpatch";
 
 // creates a hook if needed, else just adds one to the patches array
-export default <CallbackType extends Function>(patchType: PatchType) =>
-  (
-    funcName: string,
-    funcParent: any,
-    callback: CallbackType,
+export default <T extends PatchType>(patchType: T) =>
+  <Name extends KeysWithFunctionValues<Parent>, Parent extends AnyObject>(
+    funcName: Name,
+    funcParent: Parent,
+    callback: PatchTypeToCallbackMap<Parent[Name]>[T],
     oneTime = false
   ) => {
     let origFunc = funcParent[funcName];
diff --git a/src/hook.ts b/src/hook.ts
index 5bb46dd..bfa4945 100644
--- a/src/hook.ts
+++ b/src/hook.ts
@@ -1,9 +1,9 @@
 // calls relevant patches and returns the final result
-import { patchedFunctions } from "./shared";
+import { AnyFunction, patchedFunctions } from "./shared";
 
 export default function (
-  patchedFunc: Function,
-  origFunc: Function,
+  patchedFunc: AnyFunction,
+  origFunc: AnyFunction,
   funcArgs: unknown[],
   // the value of `this` to apply
   ctxt: any,
diff --git a/src/index.ts b/src/index.ts
index 670cd7b..0e04543 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,12 +1,8 @@
 import getPatchFunc from "./getPatchFunc";
 
-type BeforeCallback = (args: any[]) => void | undefined | any[];
-type InsteadCallback = (args: any[], origFunc: Function) => any;
-type AfterCallback = (args: any[], ret: any) => void | undefined | any;
-
-const before = getPatchFunc<BeforeCallback>("b");
-const instead = getPatchFunc<InsteadCallback>("i");
-const after = getPatchFunc<AfterCallback>("a");
+const before = getPatchFunc("b");
+const instead = getPatchFunc("i");
+const after = getPatchFunc("a");
 
 export { instead, before, after };
 export { resetPatches as unpatchAll } from "./shared";
\ No newline at end of file
diff --git a/src/shared.ts b/src/shared.ts
index 18248d9..bf4fb58 100644
--- a/src/shared.ts
+++ b/src/shared.ts
@@ -1,5 +1,11 @@
 export type PatchType = "a" | "b" | "i";
 
+export type PatchTypeToCallbackMap<F extends AnyFunction> = {
+  a: (args: Parameters<F>, ret: ReturnType<F>) => ReturnType<F>;
+  b: (args: Parameters<F>) => ReturnType<F>;
+  i: (args: Parameters<F>, origFunc: F) => ReturnType<F>;
+}
+
 // we use this array multiple times
 export const patchTypes: PatchType[] = ["a", "b", "i"];
 
@@ -7,22 +13,30 @@ export type Patch = {
   // function name
   n: string;
   // original function
-  o: Function;
+  o: AnyFunction;
   // WeakRef to parent object
   p: WeakRef<any>;
   // cleanups
-  c: Function[];
+  c: AnyFunction[];
   // after hooks
-  a: Map<symbol, Function>;
+  a: Map<symbol, AnyFunction>;
   // before hooks
-  b: Map<symbol, Function>;
+  b: Map<symbol, AnyFunction>;
   // instead hooks
-  i: Map<symbol, Function>;
+  i: Map<symbol, AnyFunction>;
 };
 
-export let patchedFunctions: WeakMap<Function, Patch>;
+export type AnyFunction = (...args: any[]) => any;
+
+export type KeysWithFunctionValues<T extends AnyObject> = {
+  [K in Extract<keyof T, string>]: T[K] extends AnyFunction ? K : never
+}[Extract<keyof T, string>];
+
+export type AnyObject = Record<any, any>;
+
+export let patchedFunctions: WeakMap<AnyFunction, Patch>;
 export let resetPatches = () =>
-  (patchedFunctions = new WeakMap<Function, Patch>());
+  (patchedFunctions = new WeakMap<AnyFunction, Patch>());
 
 // Manual minification is funny
 resetPatches();