diff --git a/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts b/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts index f9f1e4a66f..23d368ea42 100644 --- a/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts +++ b/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts @@ -35,6 +35,7 @@ import { } from "../Optimization"; import { CodegenFunction, + alignObjectMethodScopes, alignReactiveScopesToBlockScopes, assertScopeInstructionsWithinScopes, buildReactiveBlocks, @@ -43,7 +44,6 @@ import { extractScopeDeclarationsFromDestructuring, flattenReactiveLoops, flattenScopesWithHooks, - flattenScopesWithObjectMethods, inferReactiveScopeVariables, memoizeFbtOperandsInSameScope, mergeOverlappingReactiveScopes, @@ -198,6 +198,13 @@ function* runWithEnvironment( inferReactiveScopeVariables(hir); yield log({ kind: "hir", name: "InferReactiveScopeVariables", value: hir }); + alignObjectMethodScopes(hir); + yield log({ + kind: "hir", + name: "AlignObjectMethodScopes", + value: hir, + }); + const reactiveFunction = buildReactiveFunction(hir); yield log({ kind: "reactive", @@ -265,13 +272,6 @@ function* runWithEnvironment( value: reactiveFunction, }); - flattenScopesWithObjectMethods(reactiveFunction); - yield log({ - kind: "reactive", - name: "FlattenScopesWithObjectMethods", - value: reactiveFunction, - }); - propagateScopeDependencies(reactiveFunction); yield log({ kind: "reactive", diff --git a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/AlignObjectMethodScopes.ts b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/AlignObjectMethodScopes.ts new file mode 100644 index 0000000000..4d79fbc4b1 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/AlignObjectMethodScopes.ts @@ -0,0 +1,101 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +import { CompilerError } from ".."; +import { + GeneratedSource, + HIRFunction, + Identifier, + ReactiveScope, + makeInstructionId, +} from "../HIR"; +import { eachInstructionValueOperand } from "../HIR/visitors"; +import DisjointSet from "../Utils/DisjointSet"; + +/** + * Align scopes of object method values to that of their enclosing object expressions. + * To produce a well-formed JS program in Codegen, object methods and object expressions + * must be in the same ReactiveBlock as object method definitions must be inlined. + */ + +function findScopesToMerge(fn: HIRFunction): DisjointSet { + const objectMethodDecls: Set = new Set(); + const mergeScopesBuilder = new DisjointSet(); + + for (const [_, block] of fn.body.blocks) { + for (const { lvalue, value } of block.instructions) { + if (value.kind === "ObjectMethod") { + objectMethodDecls.add(lvalue.identifier); + } else if (value.kind === "ObjectExpression") { + for (const operand of eachInstructionValueOperand(value)) { + if (objectMethodDecls.has(operand.identifier)) { + const operandScope = operand.identifier.scope; + const lvalueScope = lvalue.identifier.scope; + + CompilerError.invariant( + operandScope != null && lvalueScope != null, + { + reason: + "Internal error: Expected all ObjectExpressions and ObjectMethods to have non-null scope.", + suggestions: null, + loc: GeneratedSource, + } + ); + mergeScopesBuilder.union([operandScope, lvalueScope]); + } + } + } + } + } + return mergeScopesBuilder; +} + +export function alignObjectMethodScopes(fn: HIRFunction): void { + // Handle inner functions: we assume that Scopes are disjoint across functions + for (const [_, block] of fn.body.blocks) { + for (const { value } of block.instructions) { + if ( + value.kind === "ObjectMethod" || + value.kind === "FunctionExpression" + ) { + alignObjectMethodScopes(value.loweredFunc.func); + } + } + } + + const scopeGroupsMap = findScopesToMerge(fn).canonicalize(); + /** + * Step 1: Merge affected scopes to their canonical root. + */ + for (const [scope, root] of scopeGroupsMap) { + if (scope !== root) { + root.range.start = makeInstructionId( + Math.min(scope.range.start, root.range.start) + ); + root.range.end = makeInstructionId( + Math.max(scope.range.end, root.range.end) + ); + } + } + + /** + * Step 2: Repoint identifiers whose scopes were merged. + */ + for (const [_, block] of fn.body.blocks) { + for (const { + lvalue: { identifier }, + } of block.instructions) { + if (identifier.scope != null) { + const root = scopeGroupsMap.get(identifier.scope); + if (root != null) { + identifier.scope = root; + } + // otherwise, this identifier's scope was not affected by this pass + } + } + } +} diff --git a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/FlattenScopesWithObjectMethods.ts b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/FlattenScopesWithObjectMethods.ts deleted file mode 100644 index 782f3d4cd4..0000000000 --- a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/FlattenScopesWithObjectMethods.ts +++ /dev/null @@ -1,58 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * - * This source code is licensed under the MIT license found in the - * LICENSE file in the root directory of this source tree. - */ - -import { - InstructionId, - ReactiveFunction, - ReactiveScopeBlock, - ReactiveStatement, - ReactiveValue, -} from "../HIR"; -import { - ReactiveFunctionTransform, - Transformed, - visitReactiveFunction, -} from "./visitors"; - -export function flattenScopesWithObjectMethods(fn: ReactiveFunction): void { - visitReactiveFunction(fn, new Transform(), { - hasObjectMethod: false, - }); -} - -type State = { - hasObjectMethod: boolean; -}; - -class Transform extends ReactiveFunctionTransform { - override transformScope( - scope: ReactiveScopeBlock, - outerState: State - ): Transformed { - const innerState: State = { - hasObjectMethod: false, - }; - this.visitScope(scope, innerState); - outerState.hasObjectMethod ||= innerState.hasObjectMethod; - if (innerState.hasObjectMethod) { - return { kind: "replace-many", value: scope.instructions }; - } else { - return { kind: "keep" }; - } - } - - override visitValue( - id: InstructionId, - value: ReactiveValue, - state: State - ): void { - this.traverseValue(id, value, state); - if (value.kind === "ObjectMethod") { - state.hasObjectMethod = true; - } - } -} diff --git a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/index.ts b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/index.ts index 396637fa7f..91fb5c2d37 100644 --- a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/index.ts +++ b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/index.ts @@ -5,6 +5,7 @@ * LICENSE file in the root directory of this source tree. */ +export { alignObjectMethodScopes } from "./AlignObjectMethodScopes"; export { alignReactiveScopesToBlockScopes } from "./AlignReactiveScopesToBlockScopes"; export { assertScopeInstructionsWithinScopes } from "./AssertScopeInstructionsWithinScope"; export { buildReactiveBlocks } from "./BuildReactiveBlocks"; @@ -16,7 +17,6 @@ export { export { extractScopeDeclarationsFromDestructuring } from "./ExtractScopeDeclarationsFromDestructuring"; export { flattenReactiveLoops } from "./FlattenReactiveLoops"; export { flattenScopesWithHooks } from "./FlattenScopesWithHooks"; -export { flattenScopesWithObjectMethods } from "./FlattenScopesWithObjectMethods"; export { inferReactiveScopeVariables } from "./InferReactiveScopeVariables"; export { memoizeFbtOperandsInSameScope } from "./MemoizeFbtOperandsInSameScope"; export { mergeOverlappingReactiveScopes } from "./MergeOverlappingReactiveScopes"; diff --git a/compiler/packages/babel-plugin-react-forget/src/Utils/utils.ts b/compiler/packages/babel-plugin-react-forget/src/Utils/utils.ts index 5675ef476f..8c1bc270b0 100644 --- a/compiler/packages/babel-plugin-react-forget/src/Utils/utils.ts +++ b/compiler/packages/babel-plugin-react-forget/src/Utils/utils.ts @@ -68,6 +68,12 @@ export function Set_union(a: Set, b: Set): Set { return union; } +export function nonNull, U>( + value: T | null | undefined +): value is T { + return value != null; +} + export function hasNode( input: NodePath ): input is NodePath> { diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-literal-method-derived-in-ternary-consequent.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-literal-method-derived-in-ternary-consequent.expect.md index a2eef20655..113a14a35a 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-literal-method-derived-in-ternary-consequent.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-literal-method-derived-in-ternary-consequent.expect.md @@ -24,17 +24,28 @@ export const FIXTURE_ENTRYPOINT = { ## Code ```javascript +import { unstable_useMemoCache as useMemoCache } from "react"; import { identity, createHookWrapper } from "shared-runtime"; function useHook(t17) { + const $ = useMemoCache(3); const { isCond, value } = t17; - return isCond - ? identity({ - getValue() { - return value; - }, - }) - : 42; + let t0; + if ($[0] !== isCond || $[1] !== value) { + t0 = isCond + ? identity({ + getValue() { + return value; + }, + }) + : 42; + $[0] = isCond; + $[1] = value; + $[2] = t0; + } else { + t0 = $[2]; + } + return t0; } export const FIXTURE_ENTRYPOINT = { diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-literal-method-in-ternary-consequent.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-literal-method-in-ternary-consequent.expect.md index 6b942167ae..5fc1344933 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-literal-method-in-ternary-consequent.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-literal-method-in-ternary-consequent.expect.md @@ -24,17 +24,28 @@ export const FIXTURE_ENTRYPOINT = { ## Code ```javascript +import { unstable_useMemoCache as useMemoCache } from "react"; import { createHookWrapper } from "shared-runtime"; function useHook(t15) { + const $ = useMemoCache(3); const { isCond, value } = t15; - return isCond - ? { - getValue() { - return value; - }, - } - : 42; + let t0; + if ($[0] !== isCond || $[1] !== value) { + t0 = isCond + ? { + getValue() { + return value; + }, + } + : 42; + $[0] = isCond; + $[1] = value; + $[2] = t0; + } else { + t0 = $[2]; + } + return t0; } export const FIXTURE_ENTRYPOINT = { diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-maybe-alias.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-maybe-alias.expect.md index 5cf8c0aa9b..2c5074fd35 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-maybe-alias.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-maybe-alias.expect.md @@ -27,20 +27,31 @@ export const FIXTURE_ENTRYPOINT = { ## Code ```javascript +import { unstable_useMemoCache as useMemoCache } from "react"; import { createHookWrapper, setProperty } from "shared-runtime"; function useHook(props) { - const x = { - getX() { - return props; - }, - }; + const $ = useMemoCache(2); + let t0; + if ($[0] !== props) { + const x = { + getX() { + return props; + }, + }; - const y = { - getY() { - return "y"; - }, - }; - return setProperty(x, y); + const y = { + getY() { + return "y"; + }, + }; + + t0 = setProperty(x, y); + $[0] = props; + $[1] = t0; + } else { + t0 = $[1]; + } + return t0; } export const FIXTURE_ENTRYPOINT = { diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-3.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-3.expect.md index 4a8402bd44..f07b493362 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-3.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-3.expect.md @@ -25,17 +25,28 @@ export const FIXTURE_ENTRYPOINT = { ## Code ```javascript +import { unstable_useMemoCache as useMemoCache } from "react"; import { createHookWrapper, mutate } from "shared-runtime"; function useHook(a) { - const x = { a }; - const obj = { - method() { - mutate(x); - return x; - }, - }; - return obj.method(); + const $ = useMemoCache(2); + let t0; + if ($[0] !== a) { + const x = { a }; + const obj = { + method() { + mutate(x); + return x; + }, + }; + + t0 = obj.method(); + $[0] = a; + $[1] = t0; + } else { + t0 = $[1]; + } + return t0; } export const FIXTURE_ENTRYPOINT = { diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-aliased-mutate-after.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-aliased-mutate-after.expect.md index d7807ce219..086f6be456 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-aliased-mutate-after.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-aliased-mutate-after.expect.md @@ -27,23 +27,23 @@ export const FIXTURE_ENTRYPOINT = { import { unstable_useMemoCache as useMemoCache } from "react"; import { createHookWrapper, mutate, mutateAndReturn } from "shared-runtime"; function useHook(t21) { - const $ = useMemoCache(1); + const $ = useMemoCache(2); const { value } = t21; - const x = mutateAndReturn({ value }); - let t0; - if ($[0] === Symbol.for("react.memo_cache_sentinel")) { - t0 = { + let obj; + if ($[0] !== value) { + const x = mutateAndReturn({ value }); + obj = { getValue() { return value; }, }; - $[0] = t0; - } else { - t0 = $[0]; - } - const obj = t0; - mutate(x); + mutate(x); + $[0] = value; + $[1] = obj; + } else { + obj = $[1]; + } return obj; } diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-derived-value.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-derived-value.expect.md index d1ba886cd2..4e02a87eed 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-derived-value.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-derived-value.expect.md @@ -26,7 +26,7 @@ export const FIXTURE_ENTRYPOINT = { import { unstable_useMemoCache as useMemoCache } from "react"; import { createHookWrapper, mutateAndReturn } from "shared-runtime"; function useHook(t18) { - const $ = useMemoCache(3); + const $ = useMemoCache(4); const { value } = t18; let t0; if ($[0] !== value) { @@ -38,15 +38,16 @@ function useHook(t18) { } const x = t0; let t1; - if ($[2] === Symbol.for("react.memo_cache_sentinel")) { + if ($[2] !== x) { t1 = { getValue() { return x; }, }; - $[2] = t1; + $[2] = x; + $[3] = t1; } else { - t1 = $[2]; + t1 = $[3]; } const obj = t1; return obj; diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-hook-dep.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-hook-dep.expect.md index 914c845acc..4798e4a627 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-hook-dep.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-hook-dep.expect.md @@ -26,18 +26,19 @@ export const FIXTURE_ENTRYPOINT = { import { createHookWrapper } from "shared-runtime"; import { useState, unstable_useMemoCache as useMemoCache } from "react"; function useFoo() { - const $ = useMemoCache(1); + const $ = useMemoCache(2); const [state] = useState(false); let t0; - if ($[0] === Symbol.for("react.memo_cache_sentinel")) { + if ($[0] !== state) { t0 = { func() { return state; }, }; - $[0] = t0; + $[0] = state; + $[1] = t0; } else { - t0 = $[0]; + t0 = $[1]; } return t0; } diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-mutated-after.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-mutated-after.expect.md index 8d6883aeac..17a5cacd72 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-mutated-after.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-method-shorthand-mutated-after.expect.md @@ -24,17 +24,26 @@ export const FIXTURE_ENTRYPOINT = { ## Code ```javascript +import { unstable_useMemoCache as useMemoCache } from "react"; import { createHookWrapper, mutate, mutateAndReturn } from "shared-runtime"; function useHook(t21) { + const $ = useMemoCache(2); const { value } = t21; - const x = mutateAndReturn({ value }); - const obj = { - getValue() { - return x; - }, - }; + let obj; + if ($[0] !== value) { + const x = mutateAndReturn({ value }); + obj = { + getValue() { + return x; + }, + }; - mutate(obj); + mutate(obj); + $[0] = value; + $[1] = obj; + } else { + obj = $[1]; + } return obj; } diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-1.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-1.expect.md index c2392aaef7..1a1124f8d6 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-1.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-1.expect.md @@ -27,7 +27,7 @@ export const FIXTURE_ENTRYPOINT = { import { unstable_useMemoCache as useMemoCache } from "react"; import { createHookWrapper } from "shared-runtime"; function useHook(t16) { - const $ = useMemoCache(4); + const $ = useMemoCache(5); const { a, b } = t16; let t0; if ($[0] !== a) { @@ -40,17 +40,18 @@ function useHook(t16) { t0 = $[1]; } let t1; - if ($[2] !== t0) { + if ($[2] !== b || $[3] !== t0) { t1 = { x: t0, y() { return [b]; }, }; - $[2] = t0; - $[3] = t1; + $[2] = b; + $[3] = t0; + $[4] = t1; } else { - t1 = $[3]; + t1 = $[4]; } return t1; } diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-2.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-2.expect.md index aad497ae8d..552522e514 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-2.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-2.expect.md @@ -28,7 +28,7 @@ import { unstable_useMemoCache as useMemoCache } from "react"; import { createHookWrapper } from "shared-runtime"; function useHook(t16) { - const $ = useMemoCache(7); + const $ = useMemoCache(8); const { a, b, c } = t16; let t0; if ($[0] !== a) { @@ -38,16 +38,16 @@ function useHook(t16) { } else { t0 = $[1]; } - let t1; - if ($[2] !== c) { - t1 = { c }; - $[2] = c; - $[3] = t1; - } else { - t1 = $[3]; - } let t2; - if ($[4] !== t0 || $[5] !== t1) { + if ($[2] !== b || $[3] !== c || $[4] !== t0) { + let t1; + if ($[6] !== c) { + t1 = { c }; + $[6] = c; + $[7] = t1; + } else { + t1 = $[7]; + } t2 = { x: t0, y() { @@ -55,11 +55,12 @@ function useHook(t16) { }, z: t1, }; + $[2] = b; + $[3] = c; $[4] = t0; - $[5] = t1; - $[6] = t2; + $[5] = t2; } else { - t2 = $[6]; + t2 = $[5]; } return t2; } diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-nested.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-nested.expect.md index 010b3b5a2d..c84e813415 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-nested.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/object-shorthand-method-nested.expect.md @@ -35,11 +35,11 @@ import { useState, unstable_useMemoCache as useMemoCache } from "react"; import { createHookWrapper } from "shared-runtime"; function useHook(t21) { - const $ = useMemoCache(1); + const $ = useMemoCache(3); const { value } = t21; const [state] = useState(false); let t0; - if ($[0] === Symbol.for("react.memo_cache_sentinel")) { + if ($[0] !== value || $[1] !== state) { t0 = { getX() { return { @@ -51,9 +51,11 @@ function useHook(t21) { }; }, }; - $[0] = t0; + $[0] = value; + $[1] = state; + $[2] = t0; } else { - t0 = $[0]; + t0 = $[2]; } return t0; }