From 2c3a9f81de0b63364bb31f4a4dd8e559784ece16 Mon Sep 17 00:00:00 2001 From: AAGaming Date: Tue, 9 Jul 2024 02:37:46 -0400 Subject: [PATCH] feat(utils): react tree patching api --- src/utils/index.ts | 1 + src/utils/patcher.ts | 2 +- src/utils/react-patching.ts | 100 ++++++++++++++++++++++++++++++++++++ 3 files changed, 102 insertions(+), 1 deletion(-) create mode 100644 src/utils/react-patching.ts diff --git a/src/utils/index.ts b/src/utils/index.ts index eb2bff3a..9c7e55ad 100644 --- a/src/utils/index.ts +++ b/src/utils/index.ts @@ -1,5 +1,6 @@ export * from './patcher'; export * from './react'; +export * from './react-patching'; export * from './static-classes'; declare global { diff --git a/src/utils/patcher.ts b/src/utils/patcher.ts index 0e6c41dd..9f5091fb 100644 --- a/src/utils/patcher.ts +++ b/src/utils/patcher.ts @@ -6,7 +6,7 @@ export interface PatchOptions { singleShot?: boolean; } -type GenericPatchHandler = (args: any[], ret?: any) => any; +export type GenericPatchHandler = (args: any[], ret?: any) => any; export interface Patch { original: Function; diff --git a/src/utils/react-patching.ts b/src/utils/react-patching.ts new file mode 100644 index 00000000..8057b5e1 --- /dev/null +++ b/src/utils/react-patching.ts @@ -0,0 +1,100 @@ +import Logger from '../logger'; +import { GenericPatchHandler, afterPatch } from './patcher'; +import { wrapReactClass, wrapReactType } from './react'; + +// TODO max size limit? could implement as a class extending map perhaps +type PatchedComponentCache = Map; +export type GenericNodeStep = (node: any) => any; +// to patch a specific method of a class component other than render. TODO implement +// export type ClassNodeStep = [finder: GenericNodeStep, method: string]; +// export type NodeStep = GenericNodeStep | ClassNodeStep; +export type NodeStep = GenericNodeStep; +export type ReactPatchHandler = GenericPatchHandler; + +// These will get *very* spammy. +let loggingEnabled = false; +let perfLoggingEnabled = false; + +export function setReactPatcherLoggingEnabled(value: boolean = true) { loggingEnabled = value }; +export function setReactPatcherPerformanceLoggingEnabled(value: boolean = true) { perfLoggingEnabled = value }; + +function patchComponent(node: any, handler: GenericPatchHandler, steps: NodeStep[], step: number, caches: PatchedComponentCache[], logger: Logger, prop: string = 'type') { + loggingEnabled && logger.group('Patching node:', node); + // We need to take extra care to not mutate the original node.type + switch (typeof node?.[prop]) { + case 'function': + // Function component + const patch = afterPatch(node, prop, steps[step + 1] ? createStepHandler(handler, steps, step + 1, caches, logger) : handler); + loggingEnabled && logger.debug('Patched a function component', patch); + break; + case 'object': + if (node[prop]?.prototype?.render) { + // Class component + // TODO handle patching custom methods + wrapReactClass(node); + const patch = afterPatch(node[prop].prototype, 'render', steps[step + 1] ? createStepHandler(handler, steps, step + 1, caches, logger) : handler); + loggingEnabled && logger.debug('Patched class component', patch); + } else { + loggingEnabled && logger.debug('Patching forwardref/memo'); + wrapReactType(node, prop); + // Step down the object + patchComponent(node[prop], handler, steps, step, caches, logger, node[prop]?.render ? 'render' : 'type'); + } + break; + default: + logger.error('Unhandled component type', node); + break; + } + loggingEnabled && logger.groupEnd(); +} + +function handleStep(tree: any, handler: GenericPatchHandler, steps: NodeStep[], step: number, caches: PatchedComponentCache[], logger: Logger): any { + const startTime = (loggingEnabled || perfLoggingEnabled) ? performance.now() : 0; + const stepHandler = steps[step]; + const cache = caches[step] || (caches[step] = new Map()); + loggingEnabled && logger.debug(`Patch step ${step} running`, { tree, stepHandler, step, caches }); + + const node = stepHandler(tree); + if (node && node.type) { + loggingEnabled && logger.debug('Found node', node); + } else if (node) { + loggingEnabled && logger.error('Found node without type. Something is probably wrong.', node); + return tree; + } else { + loggingEnabled && logger.warn('Found no node. Depending on your usecase, this might be fine.', node); + return tree; + } + + let cachedType; + if (cachedType = cache.get(node.type)) { + loggingEnabled && logger.debug('Found cached patched component', node); + node.type = cachedType; + (loggingEnabled || perfLoggingEnabled) && logger.debug(`Patch step ${step} took ${performance.now() - startTime}ms with cache`); + return tree; + } + + const originalType = node.type; + + patchComponent(node, handler, steps, step, caches, logger); + + cache.set(originalType, node.type); + + (loggingEnabled || perfLoggingEnabled) && logger.debug(`Patch step ${step} took ${performance.now() - startTime}ms`); + return tree; +}; + +function createStepHandler(handler: GenericPatchHandler, steps: NodeStep[], step: number, caches: PatchedComponentCache[], logger: Logger) { + loggingEnabled && logger.debug(`Creating handler for step ${step}`); + return (_: any, tree: any) => handleStep(tree, handler, steps, step, caches, logger); +} + +// TODO handle createReactTreePatcher inside handler and cache it (or warn) +export function createReactTreePatcher(steps: NodeStep[], handler: GenericPatchHandler, debugName: string = 'ReactPatch'): GenericPatchHandler { + const caches: PatchedComponentCache[] = []; + + const logger = new Logger(`ReactTreePatcher -> ${debugName}`); + + loggingEnabled && logger.debug('Init with options:', steps, debugName); + + return createStepHandler(handler, steps, 0, caches, logger); +} \ No newline at end of file