Skip to content

Commit

Permalink
compiler: More precise escape analysis
Browse files Browse the repository at this point in the history 
Our current escape analysis currently works by determing values which are returned, finding the scope those values are produced in, and then walking the dependencies of those scopes. Notably, when we traverse into scopes, we _force_ their values to be memoized — even in cases where we can prove it isn't necessary.

The idea of this PR is to change escape analysis to be purely based on the flow of values. So rather than assume all scope deps need to be force-memoized, we look at how those values are used.

Note: the main motivation for this PR is the follow-up, which allows us to infer dependencies for pruned scopes. Without this change, inferring deps for pruned scopes causes the escape analysis pass to consider values as escaping which shouldn't be.

[ghstack-poisoned]
  • Loading branch information
@josephsavona
josephsavona committed Jun 6, 2024
1 parent 4fb7a93 commit d9b212d
Show file tree
Hide file tree
Showing 4 changed files with 76 additions and 51 deletions.
63 changes: 47 additions & 16 deletions compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts
View file
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
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
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;
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
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);
$[0] = arr;
$[1] = t1;
} else {
t2 = $[2];
t1 = $[1];
}
return t2;
return t1;
}

export const FIXTURE_ENTRYPOINT = {
Expand Down

0 comments on commit d9b212d

Please sign in to comment.