From 80155a95ad2b14a6a75cd1cf88dfc5cf2c4de800 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Thu, 1 May 2025 17:43:49 +0900 Subject: [PATCH] [compiler] Fix for uncalled functions that are known-mutable 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: ```js 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: ```js 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. NOTE: The alternative is to reject these cases. If we do that we'd also want to similarly disallow cases like passing a mutable function to a hook. [ghstack-poisoned] --- .../InerAliasForUncalledFunctions.ts | 134 ++++++++++++++++++ .../src/Inference/InferMutableRanges.ts | 2 + ...le-values-memoizes-with-captures-values.js | 38 +++++ ...utable-method-not-memod-together.expect.md | 58 ++++++++ 4 files changed, 232 insertions(+) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/Inference/InerAliasForUncalledFunctions.ts create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-uncalled-function-capturing-mutable-values-memoizes-with-captures-values.js create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.object-captured-by-returned-mutable-method-not-memod-together.expect.md 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.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]); +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.object-captured-by-returned-mutable-method-not-memod-together.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.object-captured-by-returned-mutable-method-not-memod-together.expect.md new file mode 100644 index 0000000000000..11f26962a2833 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.object-captured-by-returned-mutable-method-not-memod-together.expect.md @@ -0,0 +1,58 @@ + +## Input + +```javascript +// @flow +/** + * Creates a map function that lazily caches the results for future invocations. + */ +hook useMemoMap( + map: TInput => TOutput +): TInput => TOutput { + return useMemo(() => { + 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