[compiler] Fix for uncalled functions that are known-mutable

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<Identifier>,
+): 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<Identifier> = 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);
+        }
+      }
+    }
+  }
+}

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;

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<TInput: interface {}, TOutput>(
+  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<TInput, TOutput>();
+    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

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<TInput: interface {}, TOutput>(
+  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<TInput, TOutput>();
+    return input => {
+      let output = cache.get(input);
+      if (output == null) {
+        output = map(input);
+        cache.set(input, output);
+      }
+      return output;
+    };
+  }, [map]);
+}