[RFC] Refine memoization for Array#map with non-mutating callbacks

Improves memoization for cases such as #2409: ```javascript const x = []; useEffect(...); return <div>{x.map(item => <span>{item}</span>)}</div>; ``` We previously thought that the `x.map(...)` call mutated `x` since its kind was Mutable. However, in this case we can determine that the map call cannot mutate `x` (or anything else): the lambda does not mutate any free variables and does not mutate its arguments. This PR adds a new flag to function signatures, used for method calls only, that checks for such cases. The idea is that if the receiver is the only thing that is mutable — including that there are no args which are function expressions which mutate their parameters — then we can infer the effect as a read. See tests which confirm that function expressions which capture or mutate their params bypass the optimization.
2025-11-01 09:12:30 +00:00 · 2023-11-29 10:46:43 -08:00
parent 65c54e2d70
commit f7d16db544
10 changed files with 456 additions and 13 deletions
@@ -971,6 +971,7 @@ export enum Effect {
   * But we do not error if the value is known to be immutable.
   */
  ConditionallyMutate = "mutate?",
+
  /*
   * This reference *does* write to (mutate) the value. It is an error (invalid input)
   * if an immutable value flows into a location with this effect.
@@ -147,6 +147,20 @@ export type FunctionSignature = {
   * may choose not to memoize arguments if they do not otherwise escape.
   */
  noAlias?: boolean;
+
+  /**
+   * Supported only for methods (no-op when used on functions in CallExpression.callee position).
+   *
+   * Indicates that the method can only modify its receiver if any of the arguments
+   * are mutable or are function expressions which mutate their arguments. This is designed
+   * for methods such as Array.prototype.map(), which only mutate the receiver array if they are
+   * passed a callback which has mutable side-effects (including mutating its inputs).
+   *
+   * MethodCalls to such functions will use a different behavior depending on their arguments:
+   * - If arguments are all non-mutable, the arguments get the Read effect and the receiver is Capture.
+   * - Else uses the effects specified by this signature.
+   */
+  mutableOnlyIfOperandsAreMutable?: boolean;
 };

 /*
@@ -231,6 +245,7 @@ addObject(BUILTIN_SHAPES, BuiltInArrayId, [
      calleeEffect: Effect.ConditionallyMutate,
      returnValueKind: ValueKind.Mutable,
      noAlias: true,
+      mutableOnlyIfOperandsAreMutable: true,
    }),
  ],
  [
@@ -247,6 +262,70 @@ addObject(BUILTIN_SHAPES, BuiltInArrayId, [
      calleeEffect: Effect.ConditionallyMutate,
      returnValueKind: ValueKind.Mutable,
      noAlias: true,
+      mutableOnlyIfOperandsAreMutable: true,
+    }),
+  ],
+  [
+    "every",
+    addFunction(BUILTIN_SHAPES, [], {
+      positionalParams: [],
+      restParam: Effect.ConditionallyMutate,
+      returnType: { kind: "Primitive" },
+      /*
+       * callee is ConditionallyMutate because items of the array
+       * flow into the lambda and may be mutated there, even though
+       * the array object itself is not modified
+       */
+      calleeEffect: Effect.ConditionallyMutate,
+      returnValueKind: ValueKind.Immutable,
+      noAlias: true,
+      mutableOnlyIfOperandsAreMutable: true,
+    }),
+  ],
+  [
+    "some",
+    addFunction(BUILTIN_SHAPES, [], {
+      positionalParams: [],
+      restParam: Effect.ConditionallyMutate,
+      returnType: { kind: "Primitive" },
+      /*
+       * callee is ConditionallyMutate because items of the array
+       * flow into the lambda and may be mutated there, even though
+       * the array object itself is not modified
+       */
+      calleeEffect: Effect.ConditionallyMutate,
+      returnValueKind: ValueKind.Immutable,
+      noAlias: true,
+      mutableOnlyIfOperandsAreMutable: true,
+    }),
+  ],
+  [
+    "find",
+    addFunction(BUILTIN_SHAPES, [], {
+      positionalParams: [],
+      restParam: Effect.ConditionallyMutate,
+      returnType: { kind: "Poly" },
+      calleeEffect: Effect.ConditionallyMutate,
+      returnValueKind: ValueKind.Mutable,
+      noAlias: true,
+      mutableOnlyIfOperandsAreMutable: true,
+    }),
+  ],
+  [
+    "findIndex",
+    addFunction(BUILTIN_SHAPES, [], {
+      positionalParams: [],
+      restParam: Effect.ConditionallyMutate,
+      returnType: { kind: "Primitive" },
+      /*
+       * callee is ConditionallyMutate because items of the array
+       * flow into the lambda and may be mutated there, even though
+       * the array object itself is not modified
+       */
+      calleeEffect: Effect.ConditionallyMutate,
+      returnValueKind: ValueKind.Immutable,
+      noAlias: true,
+      mutableOnlyIfOperandsAreMutable: true,
    }),
  ],
  [
@@ -228,6 +228,17 @@ class InferenceState {
    this.#values.set(value, kind);
  }

+  values(place: Place): Array<InstructionValue> {
+    const values = this.#variables.get(place.identifier.id);
+    CompilerError.invariant(values != null, {
+      reason: `[hoisting] Expected value kind to be initialized`,
+      description: `${printPlace(place)}`,
+      loc: place.loc,
+      suggestions: null,
+    });
+    return Array.from(values);
+  }
+
  // Lookup the kind of the given @param value.
  kind(place: Place): ValueKind {
    const values = this.#variables.get(place.identifier.id);
@@ -833,6 +844,26 @@ function inferBlock(
          instrValue.property.identifier.type
        );

+        if (
+          signature !== null &&
+          signature.mutableOnlyIfOperandsAreMutable &&
+          areArgumentsImmutableAndNonMutating(state, instrValue.args)
+        ) {
+          /*
+           * None of the args are mutable or mutate their params, we can downgrade to
+           * treating as all reads
+           */
+          for (const arg of instrValue.args) {
+            const place = arg.kind === "Identifier" ? arg : arg.place;
+            state.reference(place, Effect.Read);
+          }
+          state.reference(instrValue.receiver, Effect.Read);
+          state.initialize(instrValue, signature.returnValueKind);
+          state.define(instr.lvalue, instrValue);
+          instr.lvalue.effect = Effect.ConditionallyMutate;
+          continue;
+        }
+
        const effects =
          signature !== null ? getFunctionEffects(instrValue, signature) : null;
        const returnValueKind =
@@ -1185,3 +1216,48 @@ function getFunctionEffects(
  }
  return results;
 }
+
+/**
+ * Returns true if all of the arguments are both non-mutable (immutable or frozen)
+ * _and_ are not functions which might mutate their arguments. Note that function
+ * expressions count as frozen so long as they do not mutate free variables: this
+ * function checks that such functions also don't mutate their inputs.
+ */
+function areArgumentsImmutableAndNonMutating(
+  state: InferenceState,
+  args: MethodCall["args"]
+): boolean {
+  for (const arg of args) {
+    const place = arg.kind === "Identifier" ? arg : arg.place;
+    const kind = state.kind(place);
+    switch (kind) {
+      case ValueKind.Immutable:
+      case ValueKind.Frozen: {
+        /*
+         * Only immutable values, or frozen lambdas are allowed.
+         * A lambda may appear frozen even if it may mutate its inputs,
+         * so we have a second check even for frozen value types
+         */
+        break;
+      }
+      default: {
+        return false;
+      }
+    }
+    const values = state.values(place);
+    for (const value of values) {
+      if (
+        value.kind === "FunctionExpression" &&
+        value.loweredFunc.func.params.some((param) => {
+          const place = param.kind === "Identifier" ? param : param.place;
+          const range = place.identifier.mutableRange;
+          return range.end > range.start + 1;
+        })
+      ) {
+        // This is a function which may mutate its inputs
+        return false;
+      }
+    }
+  }
+  return true;
+}
@@ -23,18 +23,34 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
 function Component(props) {
-  const $ = useMemoCache(2);
+  const $ = useMemoCache(6);
  let t0;
  if ($[0] !== props.a) {
-    const item = { a: props.a };
-    const items = [item];
-    t0 = items.map((item_0) => item_0);
+    t0 = { a: props.a };
    $[0] = props.a;
    $[1] = t0;
  } else {
    t0 = $[1];
  }
-  const mapped = t0;
+  const item = t0;
+  let t1;
+  if ($[2] !== item) {
+    t1 = [item];
+    $[2] = item;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  const items = t1;
+  let t2;
+  if ($[4] !== items) {
+    t2 = items.map((item_0) => item_0);
+    $[4] = items;
+    $[5] = t2;
+  } else {
+    t2 = $[5];
+  }
+  const mapped = t2;
  return mapped;
 }

@@ -3,25 +3,29 @@

 ```javascript
 import { useEffect, useState } from "react";
+import { mutate } from "shared-runtime";

 function Component(props) {
-  const x = [props.value];
+  const x = [{ ...props.value }];
  useEffect(() => {}, []);
  const onClick = () => {
    console.log(x.length);
  };
+  let y;
  return (
    <div onClick={onClick}>
      {x.map((item) => {
-        return <span key={item}>{item}</span>;
+        y = item;
+        return <span key={item.id}>{item.text}</span>;
      })}
+      {mutate(y)}
    </div>
  );
 }

 export const FIXTURE_ENTRYPOINT = {
  fn: Component,
-  params: [{ value: 42 }],
+  params: [{ value: { id: 0, text: "Hello!" } }],
  isComponent: true,
 };

@@ -35,10 +39,11 @@ import {
  useState,
  unstable_useMemoCache as useMemoCache,
 } from "react";
+import { mutate } from "shared-runtime";

 function Component(props) {
  const $ = useMemoCache(5);
-  const x = [props.value];
+  const x = [{ ...props.value }];
  let t0;
  let t1;
  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
@@ -55,10 +60,20 @@ function Component(props) {
    console.log(x.length);
  };

-  const t2 = x.map((item) => <span key={item}>{item}</span>);
+  let y;
+
+  const t2 = x.map((item) => {
+    y = item;
+    return <span key={item.id}>{item.text}</span>;
+  });
  let t3;
  if ($[2] !== onClick || $[3] !== t2) {
-    t3 = <div onClick={onClick}>{t2}</div>;
+    t3 = (
+      <div onClick={onClick}>
+        {t2}
+        {mutate(y)}
+      </div>
+    );
    $[2] = onClick;
    $[3] = t2;
    $[4] = t3;
@@ -70,11 +85,11 @@ function Component(props) {

 export const FIXTURE_ENTRYPOINT = {
  fn: Component,
-  params: [{ value: 42 }],
+  params: [{ value: { id: 0, text: "Hello!" } }],
  isComponent: true,
 };

 ```
      
 ### Eval output
-(kind: ok) <div><span>42</span></div>
+(kind: ok) <div><span>Hello!</span></div>
@@ -0,0 +1,26 @@
+import { useEffect, useState } from "react";
+import { mutate } from "shared-runtime";
+
+function Component(props) {
+  const x = [{ ...props.value }];
+  useEffect(() => {}, []);
+  const onClick = () => {
+    console.log(x.length);
+  };
+  let y;
+  return (
+    <div onClick={onClick}>
+      {x.map((item) => {
+        y = item;
+        return <span key={item.id}>{item.text}</span>;
+      })}
+      {mutate(y)}
+    </div>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: { id: 0, text: "Hello!" } }],
+  isComponent: true,
+};
@@ -0,0 +1,102 @@
+
+## Input
+
+```javascript
+import { useEffect, useState } from "react";
+import { mutate } from "shared-runtime";
+
+function Component(props) {
+  const x = [{ ...props.value }];
+  useEffect(() => {}, []);
+  const onClick = () => {
+    console.log(x.length);
+  };
+  let y;
+  return (
+    <div onClick={onClick}>
+      {x.map((item) => {
+        item.flag = true;
+        return <span key={item.id}>{item.text}</span>;
+      })}
+      {mutate(y)}
+    </div>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: { id: 0, text: "Hello", flag: false } }],
+  isComponent: true,
+};
+
+```
+
+## Code
+
+```javascript
+import {
+  useEffect,
+  useState,
+  unstable_useMemoCache as useMemoCache,
+} from "react";
+import { mutate } from "shared-runtime";
+
+function Component(props) {
+  const $ = useMemoCache(6);
+  const x = [{ ...props.value }];
+  let t0;
+  let t1;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = () => {};
+    t1 = [];
+    $[0] = t0;
+    $[1] = t1;
+  } else {
+    t0 = $[0];
+    t1 = $[1];
+  }
+  useEffect(t0, t1);
+  const onClick = () => {
+    console.log(x.length);
+  };
+
+  let y;
+
+  const t3 = x.map((item) => {
+    item.flag = true;
+    return <span key={item.id}>{item.text}</span>;
+  });
+  let t2;
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t2 = mutate(y);
+    $[2] = t2;
+  } else {
+    t2 = $[2];
+  }
+  let t4;
+  if ($[3] !== onClick || $[4] !== t3) {
+    t4 = (
+      <div onClick={onClick}>
+        {t3}
+        {t2}
+      </div>
+    );
+    $[3] = onClick;
+    $[4] = t3;
+    $[5] = t4;
+  } else {
+    t4 = $[5];
+  }
+  return t4;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: { id: 0, text: "Hello", flag: false } }],
+  isComponent: true,
+};
+
+```
+      
+### Eval output
+(kind: ok) <div><span>Hello</span></div>
@@ -0,0 +1,26 @@
+import { useEffect, useState } from "react";
+import { mutate } from "shared-runtime";
+
+function Component(props) {
+  const x = [{ ...props.value }];
+  useEffect(() => {}, []);
+  const onClick = () => {
+    console.log(x.length);
+  };
+  let y;
+  return (
+    <div onClick={onClick}>
+      {x.map((item) => {
+        item.flag = true;
+        return <span key={item.id}>{item.text}</span>;
+      })}
+      {mutate(y)}
+    </div>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: { id: 0, text: "Hello", flag: false } }],
+  isComponent: true,
+};
@@ -0,0 +1,102 @@
+
+## Input
+
+```javascript
+import { useEffect, useState } from "react";
+
+function Component(props) {
+  const x = [props.value];
+  useEffect(() => {}, []);
+  const onClick = () => {
+    console.log(x.length);
+  };
+  return (
+    <div onClick={onClick}>
+      {x.map((item) => {
+        return <span key={item}>{item}</span>;
+      })}
+    </div>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+  isComponent: true,
+};
+
+```
+
+## Code
+
+```javascript
+import {
+  useEffect,
+  useState,
+  unstable_useMemoCache as useMemoCache,
+} from "react";
+
+function Component(props) {
+  const $ = useMemoCache(11);
+  let t0;
+  if ($[0] !== props.value) {
+    t0 = [props.value];
+    $[0] = props.value;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const x = t0;
+  let t1;
+  let t2;
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = () => {};
+    t2 = [];
+    $[2] = t1;
+    $[3] = t2;
+  } else {
+    t1 = $[2];
+    t2 = $[3];
+  }
+  useEffect(t1, t2);
+  let t3;
+  if ($[4] !== x.length) {
+    t3 = () => {
+      console.log(x.length);
+    };
+    $[4] = x.length;
+    $[5] = t3;
+  } else {
+    t3 = $[5];
+  }
+  const onClick = t3;
+  let t4;
+  if ($[6] !== x) {
+    t4 = x.map((item) => <span key={item}>{item}</span>);
+    $[6] = x;
+    $[7] = t4;
+  } else {
+    t4 = $[7];
+  }
+  let t5;
+  if ($[8] !== onClick || $[9] !== t4) {
+    t5 = <div onClick={onClick}>{t4}</div>;
+    $[8] = onClick;
+    $[9] = t4;
+    $[10] = t5;
+  } else {
+    t5 = $[10];
+  }
+  return t5;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+  isComponent: true,
+};
+
+```
+      
+### Eval output
+(kind: ok) <div><span>42</span></div>