diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts index 205eb07c7dd..48cee368bc6 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts @@ -14,6 +14,7 @@ import { Place, ReactiveFunction, ReactiveInstruction, + ReactiveScope, ReactiveScopeBlock, ReactiveStatement, ReactiveTerminal, @@ -122,18 +123,12 @@ export function pruneNonEscapingScopes(fn: ReactiveFunction): void { } visitReactiveFunction(fn, new CollectDependenciesVisitor(fn.env), state); - // log(() => prettyFormat(state)); - /* * Then walk outward from the returned values and find all captured operands. * This forms the set of identifiers which should be memoized. */ const memoized = computeMemoizedIdentifiers(state); - // log(() => prettyFormat(memoized)); - - // log(() => printReactiveFunction(fn)); - // Prune scopes that do not declare/reassign any escaping values visitReactiveFunction(fn, new PruneScopesTransform(), memoized); } @@ -200,7 +195,9 @@ type IdentifierNode = { // A scope node describing its dependencies type ScopeNode = { - dependencies: Array; + declarations: Array; + inputs: Set; + innerScopes: Set; seen: boolean; }; @@ -216,6 +213,7 @@ class State { identifiers: Map = new Map(); scopes: Map = new Map(); escapingValues: Set = new Set(); + currentScope: ReactiveScope | null = null; constructor(env: Environment) { this.env = env; @@ -247,11 +245,18 @@ class State { let node = this.scopes.get(scope.id); if (node === undefined) { node = { - dependencies: [...scope.dependencies].map((dep) => dep.identifier.id), + declarations: [...scope.declarations] + .map(([id]) => id) + .concat( + [...scope.reassignments].map((identifier) => identifier.id) + ), + inputs: new Set(), + innerScopes: new Set(), seen: false, }; this.scopes.set(scope.id, node); } + node.inputs.add(identifier); const identifierNode = this.identifiers.get(identifier); CompilerError.invariant(identifierNode !== undefined, { reason: "Expected identifier to be initialized", @@ -262,6 +267,22 @@ class State { identifierNode.scopes.add(scope.id); } } + + enterScope(scope: ReactiveScope, fn: () => void): void { + const parentScope = this.currentScope; + this.currentScope = scope; + fn(); + this.currentScope = parentScope; + + if ( + parentScope !== null && + scope.declarations.size === 0 && + scope.reassignments.size === 0 + ) { + const parentNode = this.scopes.get(parentScope.id)!; + parentNode.innerScopes.add(scope.id); + } + } } /* @@ -273,7 +294,7 @@ function computeMemoizedIdentifiers(state: State): Set { const memoized = new Set(); // Visit an identifier, optionally forcing it to be memoized - function visit(id: IdentifierId, forceMemoize: boolean = false): boolean { + function visit(id: IdentifierId): boolean { const node = state.identifiers.get(id); CompilerError.invariant(node !== undefined, { reason: `Expected a node for all identifiers, none found for \`${id}\``, @@ -301,21 +322,19 @@ function computeMemoizedIdentifiers(state: State): Set { if ( node.level === MemoizationLevel.Memoized || - (node.level === MemoizationLevel.Conditional && - (hasMemoizedDependency || forceMemoize)) || - (node.level === MemoizationLevel.Unmemoized && forceMemoize) + (node.level === MemoizationLevel.Conditional && hasMemoizedDependency) ) { node.memoized = true; memoized.add(id); for (const scope of node.scopes) { - forceMemoizeScopeDependencies(scope); + memoizeScopeDependencies(scope); } } return node.memoized; } // Force all the scope's optionally-memoizeable dependencies (not "Never") to be memoized - function forceMemoizeScopeDependencies(id: ScopeId): void { + function memoizeScopeDependencies(id: ScopeId): void { const node = state.scopes.get(id); CompilerError.invariant(node !== undefined, { reason: "Expected a node for all scopes", @@ -328,8 +347,14 @@ function computeMemoizedIdentifiers(state: State): Set { } node.seen = true; - for (const dep of node.dependencies) { - visit(dep, true); + for (const dep of node.declarations) { + visit(dep); + } + for (const dep of node.inputs) { + visit(dep); + } + for (const scope of node.innerScopes) { + memoizeScopeDependencies(scope); } return; } @@ -814,6 +839,12 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor { }; } + override visitScope(scopeBlock: ReactiveScopeBlock, state: State): void { + state.enterScope(scopeBlock.scope, () => { + this.traverseScope(scopeBlock, state); + }); + } + override visitInstruction( instruction: ReactiveInstruction, state: State diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/capturing-func-mutate.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/capturing-func-mutate.expect.md index 5b48ec9070c..6bdcb726fd9 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/capturing-func-mutate.expect.md +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/capturing-func-mutate.expect.md @@ -26,11 +26,19 @@ export const FIXTURE_ENTRYPOINT = { ```javascript import { c as _c } from "react/compiler-runtime"; function component(a, b) { - const $ = _c(3); + const $ = _c(5); let z; if ($[0] !== a || $[1] !== b) { z = { a }; - const y = { b }; + let t0; + if ($[3] !== b) { + t0 = { b }; + $[3] = b; + $[4] = t0; + } else { + t0 = $[4]; + } + const y = t0; const x = function () { z.a = 2; console.log(y.b); diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/method-call-computed.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/method-call-computed.expect.md index 887479c01eb..371f3861559 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/method-call-computed.expect.md +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/method-call-computed.expect.md @@ -23,7 +23,7 @@ function foo(a, b, c) { ```javascript import { c as _c } from "react/compiler-runtime"; function foo(a, b, c) { - const $ = _c(8); + const $ = _c(6); let t0; if ($[0] !== a) { t0 = makeObject(a); @@ -33,26 +33,18 @@ function foo(a, b, c) { t0 = $[1]; } const x = t0; + const y = makeObject(a); let t1; - if ($[2] !== a) { - t1 = makeObject(a); - $[2] = a; - $[3] = t1; - } else { - t1 = $[3]; - } - const y = t1; - let t2; - if ($[4] !== x || $[5] !== y.method || $[6] !== b) { - t2 = x[y.method](b); - $[4] = x; - $[5] = y.method; - $[6] = b; - $[7] = t2; + if ($[2] !== x || $[3] !== y.method || $[4] !== b) { + t1 = x[y.method](b); + $[2] = x; + $[3] = y.method; + $[4] = b; + $[5] = t1; } else { - t2 = $[7]; + t1 = $[5]; } - const z = t2; + const z = t1; return z; } diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/original-reactive-scopes-fork/mutate-outer-scope-within-value-block.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/original-reactive-scopes-fork/mutate-outer-scope-within-value-block.expect.md index bcb9caa4b63..8d7dab2376a 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/original-reactive-scopes-fork/mutate-outer-scope-within-value-block.expect.md +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/original-reactive-scopes-fork/mutate-outer-scope-within-value-block.expect.md @@ -68,26 +68,20 @@ import { CONST_TRUE, identity, shallowCopy } from "shared-runtime"; * should be merged. */ function useFoo(t0) { - const $ = _c(3); + const $ = _c(2); const { input } = t0; const arr = shallowCopy(input); + + const cond = identity(false); let t1; - if ($[0] === Symbol.for("react.memo_cache_sentinel")) { - t1 = identity(false); - $[0] = t1; - } else { - t1 = $[0]; - } - const cond = t1; - let t2; - if ($[1] !== arr) { - t2 = cond ? { val: CONST_TRUE } : mutate(arr); - $[1] = arr; - $[2] = t2; + if ($[0] !== arr) { + t1 = cond ? { val: CONST_TRUE } : mutate(arr); + $[0] = arr; + $[1] = t1; } else { - t2 = $[2]; + t1 = $[1]; } - return t2; + return t1; } export const FIXTURE_ENTRYPOINT = {