chore: Introduce middleware support for function patching (#37560)

Author: ggazzo

packages/patch-injection/src/index.ts

@@ -1,2 +1,2 @@
-export type * from './definition';
 export * from './makeFunction';
+export * from './midleware';

packages/patch-injection/src/makeFunction.ts

@@ -1,40 +1,26 @@
-import { addPatch } from './addPatch';
-import { calledFunctions, functions } from './data';
-import type { BaseFunction, PatchData, PatchFunction, PatchedFunction } from './definition';
+import type { BaseFunction, PatchFunction, PatchedFunction } from './definition';
+import { withMiddleware } from './midleware';
 
 export const makeFunction = <T extends BaseFunction>(fn: T): PatchedFunction<T> => {
-	const patches = new Set<PatchData<T>>();
-
-	patches.add({
-		patchFunction: (_next, ...args) => fn(...args),
-	});
-
-	const result = ((...args: Parameters<T>): ReturnType<T> => {
-		let newFn: T = fn;
-
-		for (const patch of patches) {
-			if (patch.condition && !patch.condition()) {
-				continue;
+	const wrapped = withMiddleware(fn);
+	const patch = (fn: PatchFunction<T>, condition?: () => boolean) => {
+		return wrapped.use((ctx, next) => {
+			if (!condition || condition()) {
+				return fn(next as unknown as T, ...ctx);
 			}
-
-			const nextFn = newFn;
-			newFn = ((...args: Parameters<T>) => patch.patchFunction(nextFn, ...args)) as T;
-		}
-
-		calledFunctions.add(result);
-		return newFn(...args);
-	}) as PatchedFunction<T>;
-
-	functions.set(result, patches as Set<PatchData<BaseFunction>>);
-
-	result.patch = (patch: PatchFunction<T>, condition?: () => boolean) => addPatch(result, patch, condition);
-
-	result.originalSignature = (() => {
+			return next(...ctx);
+		});
+	};
+	const originalSignature = (() => {
 		throw new Error('OriginalSignature of patched functions is not meant to be executed directly.');
 	}) as unknown as T;
-	result.patchSignature = (() => {
+	const patchSignature = (() => {
 		throw new Error('PatchSignature of patched functions is not meant to be executed directly.');
 	}) as unknown as PatchFunction<T>;
 
-	return result;
+	return Object.assign(wrapped, {
+		patch,
+		originalSignature,
+		patchSignature,
+	}) as PatchedFunction<T>;
 };

packages/patch-injection/src/midleware.ts

@@ -0,0 +1,36 @@
+type Middleware<F extends (...args: any[]) => any> = (ctx: Parameters<F>, next: NextFunction<F>) => ReturnType<F>;
+type NextFunction<F extends (...args: any[]) => any> = (...args: Parameters<F> | []) => ReturnType<F>;
+
+export function withMiddleware<F extends (...args: any[]) => any>(fn: F) {
+	const middlewares: Middleware<F>[] = [];
+
+	const buildRunner = (): ((ctx: Parameters<F>) => ReturnType<F>) => {
+		return middlewares.reduce(
+			(next, middleware) => {
+				return (ctx) => middleware(ctx, (...args) => next(args.length ? (args as Parameters<F>) : ctx));
+			},
+			(ctx: Parameters<F>) => fn(...ctx),
+		);
+	};
+
+	let runner = buildRunner();
+
+	const use = (middleware: Middleware<F>) => {
+		middlewares.push(middleware);
+
+		runner = buildRunner();
+		return () => {
+			const index = middlewares.indexOf(middleware);
+			if (index > -1) {
+				middlewares.splice(index, 1);
+			}
+			runner = buildRunner();
+		};
+	};
+
+	const run = (...args: Parameters<F>): ReturnType<F> => {
+		return runner(args);
+	};
+
+	return Object.assign(run as F, { use, fn });
+}