From 14920bc21af1a66d1b50bf807c00e22112d81ae1 Mon Sep 17 00:00:00 2001 From: PalmDevs Date: Fri, 11 Jul 2025 00:59:02 +0700 Subject: [PATCH] fix(lib/react/jsx-runtime/patches): bug and performance fixes - Compute fiber normally if there are no instead patches - Store all patch kinds in one single Map - Attempt cleanup after unpatch --- lib/react/src/jsx-runtime/_internal.ts | 15 +++---- lib/react/src/jsx-runtime/index.ts | 59 ++++++++++++++++++-------- lib/react/src/jsx-runtime/patches.ts | 28 ++++++------ 3 files changed, 65 insertions(+), 37 deletions(-) diff --git a/lib/react/src/jsx-runtime/_internal.ts b/lib/react/src/jsx-runtime/_internal.ts index 2866f62..39b7830 100644 --- a/lib/react/src/jsx-runtime/_internal.ts +++ b/lib/react/src/jsx-runtime/_internal.ts @@ -1,12 +1,11 @@ import type { ElementType } from 'react' import type { AfterJSXCallback, BeforeJSXCallback, InsteadJSXCallback } from '.' -export const beforeCallbacks = new Map< - ElementType, - Set> +export const jPatches = new Map< + ElementType | string, + [ + before?: Set, + after?: Set, + instead?: Set, + ] >() -export const insteadCallbacks = new Map< - ElementType, - Set> ->() -export const afterCallbacks = new Map>() diff --git a/lib/react/src/jsx-runtime/index.ts b/lib/react/src/jsx-runtime/index.ts index d601a9b..518efd0 100644 --- a/lib/react/src/jsx-runtime/index.ts +++ b/lib/react/src/jsx-runtime/index.ts @@ -1,4 +1,4 @@ -import { afterCallbacks, beforeCallbacks, insteadCallbacks } from './_internal' +import { jPatches } from './_internal' import type { ComponentProps, ElementType, Key, ReactElement } from 'react' import type { ReactJSXRuntime } from '..' @@ -16,10 +16,11 @@ export type InsteadJSXCallback< P = ComponentProps, > = ( args: [element: E, props: P, key?: Key | undefined], + jsx: AnyJSXFactoryFunction, ) => ReturnType | null export type AfterJSXCallback = ( - tree: ReactElement, + fiber: ReactElement, ) => ReturnType | null /** @@ -30,14 +31,20 @@ export type AfterJSXCallback = ( * @returns A function to unpatch. */ export function afterJSX(type: ElementType, patch: AfterJSXCallback) { - let set = afterCallbacks.get(type) - if (!set) { - set = new Set() - afterCallbacks.set(type, set) + let patches = jPatches.get(type) + if (!patches) { + patches = [undefined, new Set(), undefined] + jPatches.set(type, patches) } + const set = patches[1]! + set.add(patch as AfterJSXCallback) - return () => set.delete(patch as AfterJSXCallback) + return () => { + const res = set.delete(patch as AfterJSXCallback) + if (res) attemptCleanup(type) + return res + } } /** @@ -51,14 +58,20 @@ export function beforeJSX< E extends ElementType = ElementType, P = ComponentProps, >(type: E, patch: BeforeJSXCallback) { - let set = beforeCallbacks.get(type) - if (!set) { - set = new Set() - beforeCallbacks.set(type, set) + let patches = jPatches.get(type) + if (!patches) { + patches = [new Set(), undefined, undefined] + jPatches.set(type, patches) } + const set = patches[0]! + set.add(patch as BeforeJSXCallback) - return () => set.delete(patch as BeforeJSXCallback) + return () => { + const res = set.delete(patch as BeforeJSXCallback) + if (res) attemptCleanup(type) + return res + } } /** @@ -72,12 +85,24 @@ export function insteadJSX< E extends ElementType = ElementType, P = ComponentProps, >(type: E, patch: InsteadJSXCallback) { - let set = insteadCallbacks.get(type) - if (!set) { - set = new Set() - insteadCallbacks.set(type, set) + let patches = jPatches.get(type) + if (!patches) { + patches = [undefined, undefined, new Set()] + jPatches.set(type, patches) } + const set = patches[2]! + set.add(patch as InsteadJSXCallback) - return () => set.delete(patch as InsteadJSXCallback) + return () => { + const res = set.delete(patch as InsteadJSXCallback) + if (res) attemptCleanup(type) + return res + } +} + +function attemptCleanup(type: ElementType) { + const patches = jPatches.get(type)! + if (patches[0]?.size || patches[1]?.size || patches[2]?.size) return + jPatches.delete(type) } diff --git a/lib/react/src/jsx-runtime/patches.ts b/lib/react/src/jsx-runtime/patches.ts index b8bc9db..9d3f639 100644 --- a/lib/react/src/jsx-runtime/patches.ts +++ b/lib/react/src/jsx-runtime/patches.ts @@ -1,29 +1,33 @@ import { instead } from '@revenge-mod/patcher' import { ReactJSXRuntime } from '..' -import { afterCallbacks, beforeCallbacks, insteadCallbacks } from './_internal' -import type { ElementType, ReactElement } from 'react' +import { jPatches } from './_internal' +import type { ReactElement } from 'react' import type { AnyJSXFactoryFunction } from '.' +const jsx = ReactJSXRuntime.jsx + const patch = ( args: Parameters, orig: AnyJSXFactoryFunction, ) => { - const Comp = args[0] as any - const type: ElementType = Comp.type ?? Comp + const [type] = args + const patches = jPatches.get(type) + if (!patches) return Reflect.apply(orig, ReactJSXRuntime, args) + + const [before, after, instead] = patches - const before = beforeCallbacks.get(type) if (before) for (const cb of before) args = cb(args) - let tree: ReactElement | null | undefined - const instead = insteadCallbacks.get(type) - if (instead) for (const cb of instead) tree = cb(args) + let fiber: ReactElement | null | undefined - const after = afterCallbacks.get(type) - if (after) for (const cb of after) tree = cb(tree!) + // If there are instead patches, fiber will always be set by the instead patches + if (instead) for (const cb of instead) fiber = cb(args, jsx) + // If there aren't any instead patches, we compute the fiber normally, and it will be set + else fiber = Reflect.apply(orig, ReactJSXRuntime, args) - if (tree !== undefined) return tree! // have to cast non-nullish even if returning null works + if (after) for (const cb of after) fiber = cb(fiber!) - return Reflect.apply(orig, ReactJSXRuntime, args) + return fiber! } instead(ReactJSXRuntime, 'jsx', patch)