diff --git a/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts b/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts index ea394c9ea6..18dd8c5f42 100644 --- a/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts +++ b/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts @@ -376,11 +376,12 @@ function shouldVisitNode(fn: BabelFn, pass: CompilerPass): boolean { return false; } case "infer": { + const hookPattern = pass.opts.environment?.hookPattern ?? null; return ( // Component declarations are known components (fn.isFunctionDeclaration() && isComponentDeclaration(fn.node)) || // Otherwise check if this is a component or hook-like function - isComponentOrHookLike(fn) + isComponentOrHookLike(fn, hookPattern) ); } case "all": { @@ -414,7 +415,10 @@ function hasUseMemoCacheCall( return hasUseMemoCache; } -function isHookName(s: string): boolean { +function isHookName(s: string, hookPattern: string | null): boolean { + if (hookPattern !== null) { + return new RegExp(hookPattern).test(s); + } return /^use[A-Z0-9]/.test(s); } @@ -423,13 +427,16 @@ function isHookName(s: string): boolean { * containing a hook name. */ -function isHook(path: NodePath): boolean { +function isHook( + path: NodePath, + hookPattern: string | null +): boolean { if (path.isIdentifier()) { - return isHookName(path.node.name); + return isHookName(path.node.name, hookPattern); } else if ( path.isMemberExpression() && !path.node.computed && - isHook(path.get("property")) + isHook(path.get("property"), hookPattern) ) { const obj = path.get("object").node; const isPascalCaseNameSpace = /^[A-Z].*/; @@ -518,18 +525,20 @@ function isValidComponentParams( function isComponentOrHookLike( node: NodePath< t.FunctionDeclaration | t.ArrowFunctionExpression | t.FunctionExpression - > + >, + hookPattern: string | null ): boolean { const functionName = getFunctionName(node); // Check if the name is component or hook like: if (functionName !== null && isComponentName(functionName)) { return ( // As an added check we also look for hook invocations or JSX - callsHooksOrCreatesJsx(node) && isValidComponentParams(node.get("params")) + callsHooksOrCreatesJsx(node, hookPattern) && + isValidComponentParams(node.get("params")) ); - } else if (functionName !== null && isHook(functionName)) { + } else if (functionName !== null && isHook(functionName, hookPattern)) { // Hooks have hook invocations or JSX, but can take any # of arguments - return callsHooksOrCreatesJsx(node); + return callsHooksOrCreatesJsx(node, hookPattern); } /* @@ -539,7 +548,7 @@ function isComponentOrHookLike( if (node.isFunctionExpression() || node.isArrowFunctionExpression()) { if (isForwardRefCallback(node) || isMemoCallback(node)) { // As an added check we also look for hook invocations or JSX - return callsHooksOrCreatesJsx(node); + return callsHooksOrCreatesJsx(node, hookPattern); } else { return false; } @@ -547,7 +556,10 @@ function isComponentOrHookLike( return false; } -function callsHooksOrCreatesJsx(node: NodePath): boolean { +function callsHooksOrCreatesJsx( + node: NodePath, + hookPattern: string | null +): boolean { let invokesHooks = false; let createsJsx = false; node.traverse({ @@ -556,7 +568,7 @@ function callsHooksOrCreatesJsx(node: NodePath): boolean { }, CallExpression(call) { const callee = call.get("callee"); - if (callee.isExpression() && isHook(callee)) { + if (callee.isExpression() && isHook(callee, hookPattern)) { invokesHooks = true; } }, diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts index 3fac85f8f5..44fe5ffd1d 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts @@ -344,6 +344,19 @@ const EnvironmentConfigSchema = z.object({ * non-ideal. */ enableTreatFunctionDepsAsConditional: z.boolean().default(false), + + /** + * If specified, this value is used as a pattern for determing which global values should be + * treated as hooks. The pattern should have a single capture group, which will be used as + * the hook name for the purposes of resolving hook definitions (for builtin hooks)_. + * + * For example, by default `React$useState` would not be treated as a hook. By specifying + * `hookPattern: 'React$(\w+)'`, the compiler will treat this value equivalently to `useState()`. + * + * This setting is intended for cases where Forget is compiling code that has been prebundled + * and identifiers have been changed. + */ + hookPattern: z.string().nullable().default(null), }); export type EnvironmentConfig = z.infer; @@ -447,7 +460,20 @@ export class Environment { } getGlobalDeclaration(name: string): Global | null { - let resolvedGlobal: Global | null = this.#globals.get(name) ?? null; + let resolvedName = name; + + if (this.config.hookPattern != null) { + const match = new RegExp(this.config.hookPattern).exec(name); + if ( + match != null && + typeof match[1] === "string" && + isHookName(match[1]) + ) { + resolvedName = match[1]; + } + } + + let resolvedGlobal: Global | null = this.#globals.get(resolvedName) ?? null; if (resolvedGlobal === null) { // Hack, since we don't track module level declarations and imports if (isHookName(name)) { diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts index ec4e16bba4..c2dc19eab4 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts @@ -1233,9 +1233,15 @@ export function isUseInsertionEffectHookType(id: Identifier): boolean { } export function getHookKind(env: Environment, id: Identifier): HookKind | null { - const idType = id.type; - if (idType.kind === "Function") { - const signature = env.getFunctionSignature(idType); + return getHookKindForType(env, id.type); +} + +export function getHookKindForType( + env: Environment, + type: Type +): HookKind | null { + if (type.kind === "Function") { + const signature = env.getFunctionSignature(type); return signature?.hookKind ?? null; } return null; diff --git a/compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts b/compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts index 4697f48842..80994aca39 100644 --- a/compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts +++ b/compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts @@ -14,6 +14,7 @@ import { Instruction, Place, SpreadPattern, + getHookKindForType, makeInstructionId, } from "../HIR"; import { createTemporaryPlace, markInstructionIds } from "../HIR/HIRBuilder"; @@ -43,11 +44,11 @@ export function dropManualMemoization(func: HIRFunction): void { break; } case "LoadGlobal": { - if ( - instr.value.name === "useMemo" || - instr.value.name === "useCallback" - ) { - hooks.set(instr.lvalue.identifier.id, instr.value.name); + const global = func.env.getGlobalDeclaration(instr.value.name); + const hookKind = + global !== null ? getHookKindForType(func.env, global) : null; + if (hookKind === "useMemo" || hookKind === "useCallback") { + hooks.set(instr.lvalue.identifier.id, hookKind); } else if (instr.value.name === "React") { react.add(instr.lvalue.identifier.id); } diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/hooks-with-prefix.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/hooks-with-prefix.expect.md new file mode 100644 index 0000000000..86afdad4b2 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/hooks-with-prefix.expect.md @@ -0,0 +1,75 @@ + +## Input + +```javascript +// @hookPattern:"React\$(\w+)" + +import * as React from "react"; +import { makeArray } from "shared-runtime"; + +const React$useState = React.useState; +const React$useMemo = React.useMemo; + +function Component() { + const [state, setState] = React$useState(0); + const doubledArray = React$useMemo(() => { + return makeArray(state); + }, [state]); + return
{doubledArray.join("")}
; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], +}; + +``` + +## Code + +```javascript +import { unstable_useMemoCache as useMemoCache } from "react"; // @hookPattern:"React\$(\w+)" + +import * as React from "react"; +import { makeArray } from "shared-runtime"; + +const React$useState = React.useState; +const React$useMemo = React.useMemo; + +function Component() { + const $ = useMemoCache(5); + const [state] = React$useState(0); + let t15; + let t0; + if ($[0] !== state) { + t15 = makeArray(state); + const doubledArray = t15; + + t0 = doubledArray.join(""); + $[0] = state; + $[1] = t0; + $[2] = t15; + } else { + t0 = $[1]; + t15 = $[2]; + } + let t1; + if ($[3] !== t0) { + t1 =
{t0}
; + $[3] = t0; + $[4] = t1; + } else { + t1 = $[4]; + } + return t1; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], +}; + +``` + +### Eval output +(kind: ok)
0
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/hooks-with-prefix.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/hooks-with-prefix.js new file mode 100644 index 0000000000..02574b49d5 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/hooks-with-prefix.js @@ -0,0 +1,20 @@ +// @hookPattern:"React\$(\w+)" + +import * as React from "react"; +import { makeArray } from "shared-runtime"; + +const React$useState = React.useState; +const React$useMemo = React.useMemo; + +function Component() { + const [state, setState] = React$useState(0); + const doubledArray = React$useMemo(() => { + return makeArray(state); + }, [state]); + return
{doubledArray.join("")}
; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], +}; diff --git a/compiler/packages/fixture-test-utils/src/compiler-utils.ts b/compiler/packages/fixture-test-utils/src/compiler-utils.ts index 599dd07cb8..59d773182b 100644 --- a/compiler/packages/fixture-test-utils/src/compiler-utils.ts +++ b/compiler/packages/fixture-test-utils/src/compiler-utils.ts @@ -37,6 +37,7 @@ export function transformFixtureInput( let compilationMode: CompilationMode = "all"; let enableUseMemoCachePolyfill = false; let panicThreshold: PanicThresholdOptions = "ALL_ERRORS"; + let hookPattern: string | null = null; if (firstLine.indexOf("@compilationMode(annotation)") !== -1) { assert( @@ -85,9 +86,24 @@ export function transformFixtureInput( } let eslintSuppressionRules: Array | null = null; - const match = /@eslintSuppressionRules\(([^)]+)\)/.exec(firstLine); - if (match != null) { - eslintSuppressionRules = match[1].split("|"); + const eslintSuppressionMatch = /@eslintSuppressionRules\(([^)]+)\)/.exec( + firstLine + ); + if (eslintSuppressionMatch != null) { + eslintSuppressionRules = eslintSuppressionMatch[1].split("|"); + } + + const hookPatternMatch = /@hookPattern:"([^"]+)"/.exec(firstLine); + if ( + hookPatternMatch && + hookPatternMatch.length > 1 && + hookPatternMatch[1].trim().length > 0 + ) { + hookPattern = hookPatternMatch[1].trim(); + } else if (firstLine.includes("@hookPattern")) { + throw new Error( + 'Invalid @hookPattern:"..." pragma, must contain the prefix between balanced double quotes eg @hookPattern:"pattern"' + ); } const config = parseConfigPragmaFn(firstLine); @@ -131,6 +147,7 @@ export function transformFixtureInput( enableEmitInstrumentForget, enableEmitHookGuards, assertValidMutableRanges: true, + hookPattern, }, compilationMode, logger: null,