Skip to content

compiler: More precise escape analysis #29782

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
josephsavona wants to merge 1 commit into from
Closed

compiler: More precise escape analysis #29782

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 47 additions & 16 deletions compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import {
Place,
ReactiveFunction,
ReactiveInstruction,
ReactiveScope,
ReactiveScopeBlock,
ReactiveStatement,
ReactiveTerminal,
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -200,7 +195,9 @@ type IdentifierNode = {

// A scope node describing its dependencies
type ScopeNode = {
dependencies: Array<IdentifierId>;
declarations: Array<IdentifierId>;
inputs: Set<IdentifierId>;
innerScopes: Set<ScopeId>;
seen: boolean;
};

Expand All @@ -216,6 +213,7 @@ class State {
identifiers: Map<IdentifierId, IdentifierNode> = new Map();
scopes: Map<ScopeId, ScopeNode> = new Map();
escapingValues: Set<IdentifierId> = new Set();
currentScope: ReactiveScope | null = null;

constructor(env: Environment) {
this.env = env;
Expand Down Expand Up @@ -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",
Expand All @@ -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);
}
}
}

/*
Expand All @@ -273,7 +294,7 @@ function computeMemoizedIdentifiers(state: State): Set<IdentifierId> {
const memoized = new Set<IdentifierId>();

// 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}\``,
Expand Down Expand Up @@ -301,21 +322,19 @@ function computeMemoizedIdentifiers(state: State): Set<IdentifierId> {

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",
Expand All @@ -328,8 +347,14 @@ function computeMemoizedIdentifiers(state: State): Set<IdentifierId> {
}
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;
}
Expand Down Expand Up @@ -814,6 +839,12 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor<State> {
};
}

override visitScope(scopeBlock: ReactiveScopeBlock, state: State): void {
state.enterScope(scopeBlock.scope, () => {
this.traverseScope(scopeBlock, state);
});
}

override visitInstruction(
instruction: ReactiveInstruction,
state: State
Expand Down
12 changes: 10 additions & 2 deletions ...-react-compiler/src/__tests__/fixtures/compiler/capturing-func-mutate.expect.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
28 changes: 10 additions & 18 deletions ...n-react-compiler/src/__tests__/fixtures/compiler/method-call-computed.expect.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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;
Comment on lines +36 to +47
Copy link
Member Author

@josephsavona josephsavona Jun 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a perfect example of how this PR makes escape analysis more precise. The dependency is on y.method, which we only use where a primitive is expected (in x[y.method]). We don't have to memoize the makeObject(a) call if we're only going to extract a primitive from it.

return z;
}

Expand Down
24 changes: 9 additions & 15 deletions ...r/original-reactive-scopes-fork/mutate-outer-scope-within-value-block.expect.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Comment on lines +75 to +78
Copy link
Member Author

@josephsavona josephsavona Jun 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This output is invalid — note that arr gets mutated in the memo block, but declared outside of it — but that's because this is a bug that is fixed in the new HIR-based reactive scopes mode (which is on by default, this fixture opts out with @enableReactiveScopesInHIR:false). The version of this fixture using HIR-based reactive scopes doesn't change as a result of this PR, and both the arr = shallowCopy(input) and identity(false) are correctly inside the memo block for arr

$[0] = arr;
$[1] = t1;
} else {
t2 = $[2];
t1 = $[1];
}
return t2;
return t1;
}

export const FIXTURE_ENTRYPOINT = {
Expand Down