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

@@ -1,12 +1,11 @@
 import type { ElementType } from 'react'
 import type { AfterJSXCallback, BeforeJSXCallback, InsteadJSXCallback } from '.'
 
-export const beforeCallbacks = new Map<
-    ElementType,
-    Set<BeforeJSXCallback<any, any>>
+export const jPatches = new Map<
+    ElementType | string,
+    [
+        before?: Set<BeforeJSXCallback>,
+        after?: Set<AfterJSXCallback>,
+        instead?: Set<InsteadJSXCallback>,
+    ]
 >()
-export const insteadCallbacks = new Map<
-    ElementType,
-    Set<InsteadJSXCallback<any, any>>
->()
-export const afterCallbacks = new Map<ElementType, Set<AfterJSXCallback>>()

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<E>,
 > = (
     args: [element: E, props: P, key?: Key | undefined],
+    jsx: AnyJSXFactoryFunction,
 ) => ReturnType<AnyJSXFactoryFunction> | null
 
 export type AfterJSXCallback = (
-    tree: ReactElement,
+    fiber: ReactElement,
 ) => ReturnType<AnyJSXFactoryFunction> | 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<E>,
 >(type: E, patch: BeforeJSXCallback<E, P>) {
-    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<E>,
 >(type: E, patch: InsteadJSXCallback<E, P>) {
-    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)
 }

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<AnyJSXFactoryFunction>,
     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)