@samitouri / QOS-React / commits / a037dabd42

[compiler] Patch ValidatePreserveMemo to bailout correctly for refs

ghstack-source-id: b9c13bf5f858123b68c9e89ca8c7629cf2b90a15 Pull Request resolved: https://github.com/facebook/react/pull/30603

Mofei Zhang committed Aug 7, 2024 at 16:11 UTC a037dabd42c9c0773526cf6b25a7a2264c8251e3
14 files changed +373 -125
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts
+38 -11
@@ -9,6 +9,7 @@ import {CompilerError} from '../CompilerError';
9 import {
10 DeclarationId,
11 Environment,
12 + Identifier,
13 InstructionId,
14 Pattern,
15 Place,
@@ -24,7 +25,7 @@ import {
25 isMutableEffect,
26 } from '../HIR';
27 import {getFunctionCallSignature} from '../Inference/InferReferenceEffects';
27 -import {assertExhaustive} from '../Utils/utils';
28 +import {assertExhaustive, getOrInsertDefault} from '../Utils/utils';
29 import {getPlaceScope} from './BuildReactiveBlocks';
30 import {
31 ReactiveFunctionTransform,
@@ -935,6 +936,11 @@ class PruneScopesTransform extends ReactiveFunctionTransform<
936 Set<DeclarationId>
937 > {
938 prunedScopes: Set<ScopeId> = new Set();
939 + /**
940 + * Track reassignments so we can correctly set `pruned` flags for
941 + * inlined useMemos.
942 + */
943 + reassignments: Map<DeclarationId, Set<Identifier>> = new Map();
944
945 override transformScope(
946 scopeBlock: ReactiveScopeBlock,
@@ -977,24 +983,45 @@ class PruneScopesTransform extends ReactiveFunctionTransform<
983 }
984 }
985
986 + /**
987 + * If we pruned the scope for a non-escaping value, we know it doesn't
988 + * need to be memoized. Remove associated `Memoize` instructions so that
989 + * we don't report false positives on "missing" memoization of these values.
990 + */
991 override transformInstruction(
992 instruction: ReactiveInstruction,
993 state: Set<DeclarationId>,
994 ): Transformed<ReactiveStatement> {
995 this.traverseInstruction(instruction, state);
996
986 - /**
987 - * If we pruned the scope for a non-escaping value, we know it doesn't
988 - * need to be memoized. Remove associated `Memoize` instructions so that
989 - * we don't report false positives on "missing" memoization of these values.
990 - */
991 - if (instruction.value.kind === 'FinishMemoize') {
992 - const identifier = instruction.value.decl.identifier;
997 + const value = instruction.value;
998 + if (value.kind === 'StoreLocal' && value.lvalue.kind === 'Reassign') {
999 + const ids = getOrInsertDefault(
1000 + this.reassignments,
1001 + value.lvalue.place.identifier.declarationId,
1002 + new Set(),
1003 + );
1004 + ids.add(value.value.identifier);
1005 + } else if (value.kind === 'FinishMemoize') {
1006 + let decls;
1007 + if (value.decl.identifier.scope == null) {
1008 + /**
1009 + * If the manual memo was a useMemo that got inlined, iterate through
1010 + * all reassignments to the iife temporary to ensure they're memoized.
1011 + */
1012 + decls = this.reassignments.get(value.decl.identifier.declarationId) ?? [
1013 + value.decl.identifier,
1014 + ];
1015 + } else {
1016 + decls = [value.decl.identifier];
1017 + }
1018 +
1019 if (
994 - identifier.scope !== null &&
995 - this.prunedScopes.has(identifier.scope.id)
1020 + [...decls].every(
1021 + decl => decl.scope == null || this.prunedScopes.has(decl.scope.id),
1022 + )
1023 ) {
997 - instruction.value.pruned = true;
1024 + value.pruned = true;
1025 }
1026 }
1027
compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts
+47 -9
@@ -30,6 +30,7 @@ import {
30 ReactiveFunctionVisitor,
31 visitReactiveFunction,
32 } from '../ReactiveScopes/visitors';
33 +import {getOrInsertDefault} from '../Utils/utils';
34
35 /**
36 * Validates that all explicit manual memoization (useMemo/useCallback) was accurately
@@ -52,6 +53,16 @@ export function validatePreservedManualMemoization(fn: ReactiveFunction): void {
53 const DEBUG = false;
54
55 type ManualMemoBlockState = {
56 + /**
57 + * Tracks reassigned temporaries.
58 + * This is necessary because useMemo calls are usually inlined.
59 + * Inlining produces a `let` declaration, followed by reassignments
60 + * to the newly declared variable (one per return statement).
61 + * Since InferReactiveScopes does not merge scopes across reassigned
62 + * variables (except in the case of a mutate-after-phi), we need to
63 + * track reassignments to validate we're retaining manual memo.
64 + */
65 + reassignments: Map<DeclarationId, Set<Identifier>>;
66 // The source of the original memoization, used when reporting errors
67 loc: SourceLocation;
68
@@ -425,6 +436,18 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
436 */
437 this.recordTemporaries(instruction, state);
438 const value = instruction.value;
439 + if (
440 + value.kind === 'StoreLocal' &&
441 + value.lvalue.kind === 'Reassign' &&
442 + state.manualMemoState != null
443 + ) {
444 + const ids = getOrInsertDefault(
445 + state.manualMemoState.reassignments,
446 + value.lvalue.place.identifier.declarationId,
447 + new Set(),
448 + );
449 + ids.add(value.value.identifier);
450 + }
451 if (value.kind === 'StartMemoize') {
452 let depsFromSource: Array<ManualMemoDependency> | null = null;
453 if (value.deps != null) {
@@ -442,6 +465,7 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
465 decls: new Set(),
466 depsFromSource,
467 manualMemoId: value.manualMemoId,
468 + reassignments: new Map(),
469 };
470
471 /**
@@ -491,20 +515,34 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
515 suggestions: null,
516 },
517 );
518 + const reassignments = state.manualMemoState.reassignments;
519 state.manualMemoState = null;
520 if (!value.pruned) {
521 for (const {identifier, loc} of eachInstructionValueOperand(
522 value as InstructionValue,
523 )) {
499 - if (isUnmemoized(identifier, this.scopes)) {
500 - state.errors.push({
501 - reason:
502 - 'React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output.',
503 - description: null,
504 - severity: ErrorSeverity.CannotPreserveMemoization,
505 - loc,
506 - suggestions: null,
507 - });
524 + let decls;
525 + if (identifier.scope == null) {
526 + /**
527 + * If the manual memo was a useMemo that got inlined, iterate through
528 + * all reassignments to the iife temporary to ensure they're memoized.
529 + */
530 + decls = reassignments.get(identifier.declarationId) ?? [identifier];
531 + } else {
532 + decls = [identifier];
533 + }
534 +
535 + for (const identifier of decls) {
536 + if (isUnmemoized(identifier, this.scopes)) {
537 + state.errors.push({
538 + reason:
539 + 'React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output.',
540 + description: null,
541 + severity: ErrorSeverity.CannotPreserveMemoization,
542 + loc,
543 + suggestions: null,
544 + });
545 + }
546 }
547 }
548 }
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-maybe-mutable-ref-memo-not-preserved.expect.md deleted
-45
@@ -1,45 +0,0 @@
1 -
2 -## Input
3 -
4 -```javascript
5 -// @validatePreserveExistingMemoizationGuarantees:true
6 -
7 -import {useRef, useMemo} from 'react';
8 -import {makeArray} from 'shared-runtime';
9 -
10 -function useFoo() {
11 - const r = useRef();
12 - return useMemo(() => makeArray(r), []);
13 -}
14 -
15 -export const FIXTURE_ENTRYPOINT = {
16 - fn: useFoo,
17 - params: [],
18 -};
19 -
20 -```
21 -
22 -## Code
23 -
24 -```javascript
25 -// @validatePreserveExistingMemoizationGuarantees:true
26 -
27 -import { useRef, useMemo } from "react";
28 -import { makeArray } from "shared-runtime";
29 -
30 -function useFoo() {
31 - const r = useRef();
32 - let t0;
33 - t0 = makeArray(r);
34 - return t0;
35 -}
36 -
37 -export const FIXTURE_ENTRYPOINT = {
38 - fn: useFoo,
39 - params: [],
40 -};
41 -
42 -```
43 -
44 -### Eval output
45 -(kind: ok) [{}]
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-useMemo-dropped-infer-always-invalidating.expect.md new
+42
@@ -0,0 +1,42 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @validatePreserveExistingMemoizationGuarantees
6 +
7 +import {useMemo} from 'react';
8 +import {useHook} from 'shared-runtime';
9 +
10 +// useMemo values may not be memoized in Forget output if we
11 +// infer that their deps always invalidate.
12 +// This is technically a false positive as the useMemo in source
13 +// was effectively a no-op
14 +function useFoo(props) {
15 + const x = [];
16 + useHook();
17 + x.push(props);
18 +
19 + return useMemo(() => [x], [x]);
20 +}
21 +
22 +export const FIXTURE_ENTRYPOINT = {
23 + fn: useFoo,
24 + params: [{}],
25 +};
26 +
27 +```
28 +
29 +
30 +## Error
31 +
32 +```
33 + 13 | x.push(props);
34 + 14 |
35 +> 15 | return useMemo(() => [x], [x]);
36 + | ^^^^^^^^^^^^^^^^^^^^^^^ CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output. (15:15)
37 + 16 | }
38 + 17 |
39 + 18 | export const FIXTURE_ENTRYPOINT = {
40 +```
41 +
42 +
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-useMemo-dropped-infer-always-invalidating.ts renamed
+2 -2
@@ -5,8 +5,8 @@ import {useHook} from 'shared-runtime';
5
6 // useMemo values may not be memoized in Forget output if we
7 // infer that their deps always invalidate.
8 -// This is still correct as the useMemo in source was effectively
9 -// a no-op already.
8 +// This is technically a false positive as the useMemo in source
9 +// was effectively a no-op
10 function useFoo(props) {
11 const x = [];
12 useHook();
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.maybe-mutable-ref-not-preserved.expect.md new
+35
@@ -0,0 +1,35 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @validatePreserveExistingMemoizationGuarantees:true
6 +
7 +import {useRef, useMemo} from 'react';
8 +import {makeArray} from 'shared-runtime';
9 +
10 +function useFoo() {
11 + const r = useRef();
12 + return useMemo(() => makeArray(r), []);
13 +}
14 +
15 +export const FIXTURE_ENTRYPOINT = {
16 + fn: useFoo,
17 + params: [],
18 +};
19 +
20 +```
21 +
22 +
23 +## Error
24 +
25 +```
26 + 6 | function useFoo() {
27 + 7 | const r = useRef();
28 +> 8 | return useMemo(() => makeArray(r), []);
29 + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output. (8:8)
30 + 9 | }
31 + 10 |
32 + 11 | export const FIXTURE_ENTRYPOINT = {
33 +```
34 +
35 +
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.maybe-mutable-ref-not-preserved.ts renamed
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo-mult-returns-primitive.expect.md new
+52
@@ -0,0 +1,52 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @validatePreserveExistingMemoizationGuarantees
6 +
7 +import {useMemo} from 'react';
8 +import {identity} from 'shared-runtime';
9 +
10 +function useFoo(cond) {
11 + useMemo(() => {
12 + if (cond) {
13 + return 2;
14 + } else {
15 + return identity(5);
16 + }
17 + }, [cond]);
18 +}
19 +
20 +export const FIXTURE_ENTRYPOINT = {
21 + fn: useFoo,
22 + params: [true],
23 +};
24 +
25 +```
26 +
27 +## Code
28 +
29 +```javascript
30 +// @validatePreserveExistingMemoizationGuarantees
31 +
32 +import { useMemo } from "react";
33 +import { identity } from "shared-runtime";
34 +
35 +function useFoo(cond) {
36 + let t0;
37 + if (cond) {
38 + t0 = 2;
39 + } else {
40 + t0 = identity(5);
41 + }
42 +}
43 +
44 +export const FIXTURE_ENTRYPOINT = {
45 + fn: useFoo,
46 + params: [true],
47 +};
48 +
49 +```
50 +
51 +### Eval output
52 +(kind: ok)
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo-mult-returns-primitive.ts new
+19
@@ -0,0 +1,19 @@
1 +// @validatePreserveExistingMemoizationGuarantees
2 +
3 +import {useMemo} from 'react';
4 +import {identity} from 'shared-runtime';
5 +
6 +function useFoo(cond) {
7 + useMemo(() => {
8 + if (cond) {
9 + return 2;
10 + } else {
11 + return identity(5);
12 + }
13 + }, [cond]);
14 +}
15 +
16 +export const FIXTURE_ENTRYPOINT = {
17 + fn: useFoo,
18 + params: [true],
19 +};
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo-mult-returns.expect.md new
+52
@@ -0,0 +1,52 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @validatePreserveExistingMemoizationGuarantees
6 +
7 +import {useMemo} from 'react';
8 +import {identity} from 'shared-runtime';
9 +
10 +function useFoo(cond) {
11 + useMemo(() => {
12 + if (cond) {
13 + return identity(10);
14 + } else {
15 + return identity(5);
16 + }
17 + }, [cond]);
18 +}
19 +
20 +export const FIXTURE_ENTRYPOINT = {
21 + fn: useFoo,
22 + params: [true],
23 +};
24 +
25 +```
26 +
27 +## Code
28 +
29 +```javascript
30 +// @validatePreserveExistingMemoizationGuarantees
31 +
32 +import { useMemo } from "react";
33 +import { identity } from "shared-runtime";
34 +
35 +function useFoo(cond) {
36 + let t0;
37 + if (cond) {
38 + t0 = identity(10);
39 + } else {
40 + t0 = identity(5);
41 + }
42 +}
43 +
44 +export const FIXTURE_ENTRYPOINT = {
45 + fn: useFoo,
46 + params: [true],
47 +};
48 +
49 +```
50 +
51 +### Eval output
52 +(kind: ok)
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo-mult-returns.ts new
+19
@@ -0,0 +1,19 @@
1 +// @validatePreserveExistingMemoizationGuarantees
2 +
3 +import {useMemo} from 'react';
4 +import {identity} from 'shared-runtime';
5 +
6 +function useFoo(cond) {
7 + useMemo(() => {
8 + if (cond) {
9 + return identity(10);
10 + } else {
11 + return identity(5);
12 + }
13 + }, [cond]);
14 +}
15 +
16 +export const FIXTURE_ENTRYPOINT = {
17 + fn: useFoo,
18 + params: [true],
19 +};
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo.expect.md new
+50
@@ -0,0 +1,50 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @validatePreserveExistingMemoizationGuarantees
6 +
7 +import {useMemo} from 'react';
8 +import {identity} from 'shared-runtime';
9 +
10 +/**
11 + * This is technically a false positive, although it makes sense
12 + * to bailout as source code might be doing something sketchy.
13 + */
14 +function useFoo(x) {
15 + useMemo(() => identity(x), [x]);
16 +}
17 +
18 +export const FIXTURE_ENTRYPOINT = {
19 + fn: useFoo,
20 + params: [2],
21 +};
22 +
23 +```
24 +
25 +## Code
26 +
27 +```javascript
28 +// @validatePreserveExistingMemoizationGuarantees
29 +
30 +import { useMemo } from "react";
31 +import { identity } from "shared-runtime";
32 +
33 +/**
34 + * This is technically a false positive, although it makes sense
35 + * to bailout as source code might be doing something sketchy.
36 + */
37 +function useFoo(x) {
38 + let t0;
39 + t0 = identity(x);
40 +}
41 +
42 +export const FIXTURE_ENTRYPOINT = {
43 + fn: useFoo,
44 + params: [2],
45 +};
46 +
47 +```
48 +
49 +### Eval output
50 +(kind: ok)
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo.ts new
+17
@@ -0,0 +1,17 @@
1 +// @validatePreserveExistingMemoizationGuarantees
2 +
3 +import {useMemo} from 'react';
4 +import {identity} from 'shared-runtime';
5 +
6 +/**
7 + * This is technically a false positive, although it makes sense
8 + * to bailout as source code might be doing something sketchy.
9 + */
10 +function useFoo(x) {
11 + useMemo(() => identity(x), [x]);
12 +}
13 +
14 +export const FIXTURE_ENTRYPOINT = {
15 + fn: useFoo,
16 + params: [2],
17 +};
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/useMemo-dropped-infer-always-invalidating.expect.md deleted
-58
@@ -1,58 +0,0 @@
1 -
2 -## Input
3 -
4 -```javascript
5 -// @validatePreserveExistingMemoizationGuarantees
6 -
7 -import {useMemo} from 'react';
8 -import {useHook} from 'shared-runtime';
9 -
10 -// useMemo values may not be memoized in Forget output if we
11 -// infer that their deps always invalidate.
12 -// This is still correct as the useMemo in source was effectively
13 -// a no-op already.
14 -function useFoo(props) {
15 - const x = [];
16 - useHook();
17 - x.push(props);
18 -
19 - return useMemo(() => [x], [x]);
20 -}
21 -
22 -export const FIXTURE_ENTRYPOINT = {
23 - fn: useFoo,
24 - params: [{}],
25 -};
26 -
27 -```
28 -
29 -## Code
30 -
31 -```javascript
32 -// @validatePreserveExistingMemoizationGuarantees
33 -
34 -import { useMemo } from "react";
35 -import { useHook } from "shared-runtime";
36 -
37 -// useMemo values may not be memoized in Forget output if we
38 -// infer that their deps always invalidate.
39 -// This is still correct as the useMemo in source was effectively
40 -// a no-op already.
41 -function useFoo(props) {
42 - const x = [];
43 - useHook();
44 - x.push(props);
45 - let t0;
46 - t0 = [x];
47 - return t0;
48 -}
49 -
50 -export const FIXTURE_ENTRYPOINT = {
51 - fn: useFoo,
52 - params: [{}],
53 -};
54 -
55 -```
56 -
57 -### Eval output
58 -(kind: ok) [[{}]]
\ No newline at end of file