diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InerAliasForUncalledFunctions.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InerAliasForUncalledFunctions.ts new file mode 100644 index 0000000000000..1dc5743b73702 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InerAliasForUncalledFunctions.ts @@ -0,0 +1,134 @@ +/** + * 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 { + Effect, + HIRFunction, + Identifier, + isMutableEffect, + isRefOrRefLikeMutableType, + makeInstructionId, +} from '../HIR/HIR'; +import {eachInstructionValueOperand} from '../HIR/visitors'; +import {isMutable} from '../ReactiveScopes/InferReactiveScopeVariables'; +import DisjointSet from '../Utils/DisjointSet'; + +/** + * If a function captures a mutable value but never gets called, we don't infer a + * mutable range for that function. This means that we also don't alias the function + * with its mutable captures. + * + * This case is tricky, because we don't generally know for sure what is a mutation + * and what may just be a normal function call. For example: + * + * ``` + * hook useFoo() { + * const x = makeObject(); + * return () => { + * return readObject(x); // could be a mutation! + * } + * } + * ``` + * + * If we pessimistically assume that all such cases are mutations, we'd have to group + * lots of memo scopes together unnecessarily. However, if there is definitely a mutation: + * + * ``` + * hook useFoo(createEntryForKey) { + * const cache = new WeakMap(); + * return (key) => { + * let entry = cache.get(key); + * if (entry == null) { + * entry = createEntryForKey(key); + * cache.set(key, entry); // known mutation! + * } + * return entry; + * } + * } + * ``` + * + * Then we have to ensure that the function and its mutable captures alias together and + * end up in the same scope. However, aliasing together isn't enough if the function + * and operands all have empty mutable ranges (end = start + 1). + * + * This pass finds function expressions and object methods that have an empty mutable range + * and known-mutable operands which also don't have a mutable range, and ensures that the + * function and those operands are aliased together *and* that their ranges are updated to + * end after the function expression. This is sufficient to ensure that a reactive scope is + * created for the alias set. + */ +export function inferAliasForUncalledFunctions( + fn: HIRFunction, + aliases: DisjointSet, +): void { + for (const block of fn.body.blocks.values()) { + instrs: for (const instr of block.instructions) { + const {lvalue, value} = instr; + if ( + value.kind !== 'ObjectMethod' && + value.kind !== 'FunctionExpression' + ) { + continue; + } + /* + * If the function is known to be mutated, we will have + * already aliased any mutable operands with it + */ + const range = lvalue.identifier.mutableRange; + if (range.end > range.start + 1) { + continue; + } + /* + * If the function already has operands with an active mutable range, + * then we don't need to do anything — the function will have already + * been visited and included in some mutable alias set. This case can + * also occur due to visiting the same function in an earlier iteration + * of the outer fixpoint loop. + */ + for (const operand of eachInstructionValueOperand(value)) { + if (isMutable(instr, operand)) { + continue instrs; + } + } + const operands: Set = new Set(); + for (const effect of value.loweredFunc.func.effects ?? []) { + if (effect.kind !== 'ContextMutation') { + continue; + } + /* + * We're looking for known-mutations only, so we look at the effects + * rather than function context + */ + if (effect.effect === Effect.Store || effect.effect === Effect.Mutate) { + for (const operand of effect.places) { + /* + * It's possible that function effect analysis thinks there was a context mutation, + * but then InferReferenceEffects figures out some operands are globals and therefore + * creates a non-mutable effect for those operands. + * We should change InferReferenceEffects to swap the ContextMutation for a global + * mutation in that case, but for now we just filter them out here + */ + if ( + isMutableEffect(operand.effect, operand.loc) && + !isRefOrRefLikeMutableType(operand.identifier.type) + ) { + operands.add(operand.identifier); + } + } + } + } + if (operands.size !== 0) { + operands.add(lvalue.identifier); + aliases.union([...operands]); + // Update mutable ranges, if the ranges are empty then a reactive scope isn't created + for (const operand of operands) { + operand.mutableRange.end = makeInstructionId(instr.id + 1); + } + } + } + } +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutableRanges.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutableRanges.ts index a8f33b31d58f3..624c302fbf7ee 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutableRanges.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutableRanges.ts @@ -6,6 +6,7 @@ */ import {HIRFunction, Identifier} from '../HIR/HIR'; +import {inferAliasForUncalledFunctions} from './InerAliasForUncalledFunctions'; import {inferAliases} from './InferAlias'; import {inferAliasForPhis} from './InferAliasForPhis'; import {inferAliasForStores} from './InferAliasForStores'; @@ -76,6 +77,7 @@ export function inferMutableRanges(ir: HIRFunction): void { while (true) { inferMutableRangesForAlias(ir, aliases); inferAliasForPhis(ir, aliases); + inferAliasForUncalledFunctions(ir, aliases); const nextAliases = aliases.canonicalize(); if (areEqualMaps(prevAliases, nextAliases)) { break; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-uncalled-function-capturing-mutable-values-memoizes-with-captures-values.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-uncalled-function-capturing-mutable-values-memoizes-with-captures-values.expect.md new file mode 100644 index 0000000000000..e771bf12bde5c --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-uncalled-function-capturing-mutable-values-memoizes-with-captures-values.expect.md @@ -0,0 +1,77 @@ + +## Input + +```javascript +// @flow +/** + * This hook returns a function that when called with an input object, + * will return the result of mapping that input with the supplied map + * function. Results are cached, so if the same input is passed again, + * the same output object will be returned. + * + * Note that this technically violates the rules of React and is unsafe: + * hooks must return immutable objects and be pure, and a function which + * captures and mutates a value when called is inherently not pure. + * + * However, in this case it is technically safe _if_ the mapping function + * is pure *and* the resulting objects are never modified. This is because + * the function only caches: the result of `returnedFunction(someInput)` + * strictly depends on `returnedFunction` and `someInput`, and cannot + * otherwise change over time. + */ +hook useMemoMap( + map: TInput => TOutput +): TInput => TOutput { + return useMemo(() => { + // The original issue is that `cache` was not memoized together with the returned + // function. This was because neither appears to ever be mutated — the function + // is known to mutate `cache` but the function isn't called. + // + // The fix is to detect cases like this — functions that are mutable but not called - + // and ensure that their mutable captures are aliased together into the same scope. + const cache = new WeakMap(); + return input => { + let output = cache.get(input); + if (output == null) { + output = map(input); + cache.set(input, output); + } + return output; + }; + }, [map]); +} + +``` + +## Code + +```javascript +import { c as _c } from "react/compiler-runtime"; + +function useMemoMap(map) { + const $ = _c(2); + let t0; + let t1; + if ($[0] !== map) { + const cache = new WeakMap(); + t1 = (input) => { + let output = cache.get(input); + if (output == null) { + output = map(input); + cache.set(input, output); + } + return output; + }; + $[0] = map; + $[1] = t1; + } else { + t1 = $[1]; + } + t0 = t1; + return t0; +} + +``` + +### Eval output +(kind: exception) Fixture not implemented \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-uncalled-function-capturing-mutable-values-memoizes-with-captures-values.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-uncalled-function-capturing-mutable-values-memoizes-with-captures-values.js new file mode 100644 index 0000000000000..bce92823e3385 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-uncalled-function-capturing-mutable-values-memoizes-with-captures-values.js @@ -0,0 +1,38 @@ +// @flow +/** + * This hook returns a function that when called with an input object, + * will return the result of mapping that input with the supplied map + * function. Results are cached, so if the same input is passed again, + * the same output object will be returned. + * + * Note that this technically violates the rules of React and is unsafe: + * hooks must return immutable objects and be pure, and a function which + * captures and mutates a value when called is inherently not pure. + * + * However, in this case it is technically safe _if_ the mapping function + * is pure *and* the resulting objects are never modified. This is because + * the function only caches: the result of `returnedFunction(someInput)` + * strictly depends on `returnedFunction` and `someInput`, and cannot + * otherwise change over time. + */ +hook useMemoMap( + map: TInput => TOutput +): TInput => TOutput { + return useMemo(() => { + // The original issue is that `cache` was not memoized together with the returned + // function. This was because neither appears to ever be mutated — the function + // is known to mutate `cache` but the function isn't called. + // + // The fix is to detect cases like this — functions that are mutable but not called - + // and ensure that their mutable captures are aliased together into the same scope. + const cache = new WeakMap(); + return input => { + let output = cache.get(input); + if (output == null) { + output = map(input); + cache.set(input, output); + } + return output; + }; + }, [map]); +}