Option to not memoize non-aliased function parameters

This has been nagging at me for a _long_ time: we unnecessarily memoize function 
callbacks passed to things like Array.prototype.map, even though we know these 
functions can't escape. This PR fixes this as follows: 

* Adds a `noAlias?: boolean` flag to builtin function signatures, defaulting to 
false if not specified. 

* Adds a feature flag, `enableNoAliasOptimizations`, to gate optimizations based 
on the value of that new flag. 

* When the feature is enabled, `PruneNonEscapingScopes` now looks up the 
signature of method calls, and avoids memoizing the arguments if the signatures 
specifies `noAlias: true`. 

* Annotates Array.prototype.map and Array.prototype.filter as `noAlias`. 

This does not mean we'll never memoize arguments to Array.prototype.map, it just 
means that the argument itself won't be considered as escaping. If the function 
still escapes by some other means it will get memoized: 

``` 

function Component(props) { 

const f = () => {}; // memoized! 

const x = [].map(f); // not from here.. 

return [x, f]; // but bc it escapes here 

} 

``` 

Note: this delivers some of the wins from #1640. That PR tried to do a bunch of 
things, part of which I already landed w the introduction of 
ConditionallyMutate, which allowed us to type Array.prototype.map. This PR 
further gives us the ability to understand functions that don't alias their 
params at all. The remaining bit from #1640 is the idea of understanding that 
key hooks such as `useFragment()` return transitively readonly, transitively 
array/object/primitive values, and any `.map()` or `.filter()` calls must be on 
arrays, allowing us to optimize them. Without that extra step, we'll still have 
to memoize a lot of `array.map()` lambdas just because we aren't sure that the 
receiver is an Array. But this PR helps with some cases, and lays the groundwork 
for the rest of that PR.

compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts

@@ -109,6 +109,15 @@ export type EnvironmentConfig = Partial<{
    */
   enableFunctionCallSignatureOptimizations: boolean;
 
+  /**
+   * Enable optimizations based on the `noAlias` flag of method signatures. When enabled,
+   * function signatures can declare that they do not alias their arguments, allowing
+   * Forget to (in some cases) avoid memoizing arguments if they do not otherwise escape.
+   *
+   * Defaults to false
+   */
+  enableNoAliasOptimizations: boolean;
+
   /**
    * When enabled, the compiler assumes that hooks follow the Rules of React:
    * - Hooks may memoize computation based on any of their parameters, thus
@@ -182,6 +191,7 @@ export class Environment {
   enableFunctionCallSignatureOptimizations: boolean;
   enableAssumeHooksFollowRulesOfReact: boolean;
   enableTreatHooksAsFunctions: boolean;
+  enableNoAliasOptimizations: boolean;
   disableAllMemoization: boolean;
   enableEmitFreeze: ExternalFunction | null;
   assertValidMutableRanges: boolean;
@@ -225,6 +235,8 @@ export class Environment {
     this.validateFrozenLambdas = config?.validateFrozenLambdas ?? false;
     this.enableFunctionCallSignatureOptimizations =
       config?.enableFunctionCallSignatureOptimizations ?? false;
+    this.enableNoAliasOptimizations =
+      config?.enableNoAliasOptimizations ?? false;
     this.enableAssumeHooksFollowRulesOfReact =
       config?.enableAssumeHooksFollowRulesOfReact ?? false;
     this.enableTreatHooksAsFunctions =

compiler/packages/babel-plugin-react-forget/src/HIR/ObjectShape.ts

@@ -138,6 +138,12 @@ export type FunctionSignature = {
   returnValueKind: ValueKind;
   calleeEffect: Effect;
   hookKind: HookKind | null;
+  /**
+   * Whether any of the parameters may be aliased by each other or the return
+   * value. Defaults to false (parameters may alias). When true, the compiler
+   * may choose not to memoize arguments if they do not otherwise escape.
+   */
+  noAlias?: boolean;
 };
 
 /**
@@ -217,6 +223,7 @@ addObject(BUILTIN_SHAPES, BuiltInArrayId, [
       // the array object itself is not modified
       calleeEffect: Effect.ConditionallyMutate,
       returnValueKind: ValueKind.Mutable,
+      noAlias: true,
     }),
   ],
   [
@@ -230,6 +237,7 @@ addObject(BUILTIN_SHAPES, BuiltInArrayId, [
       // the array object itself is not modified
       calleeEffect: Effect.ConditionallyMutate,
       returnValueKind: ValueKind.Mutable,
+      noAlias: true,
     }),
   ],
   [

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PruneNonEscapingScopes.ts

@@ -24,6 +24,7 @@ import {
   getHookKind,
   isMutableEffect,
 } from "../HIR";
+import { getFunctionCallSignature } from "../Inference/InferReferenceEffects";
 import { log } from "../Utils/logger";
 import { assertExhaustive } from "../Utils/utils";
 import { getPlaceScope } from "./BuildReactiveBlocks";
@@ -119,7 +120,11 @@ export function pruneNonEscapingScopes(
   for (const param of fn.params) {
     state.declare(param.identifier.id);
   }
-  visitReactiveFunction(fn, new CollectDependenciesVisitor(options), state);
+  visitReactiveFunction(
+    fn,
+    new CollectDependenciesVisitor(fn.env, options),
+    state
+  );
 
   log(() => prettyFormat(state));
 
@@ -343,6 +348,7 @@ type LValueMemoization = {
  * - level: the level of memoization to apply to this value
  */
 function computeMemoizationInputs(
+  env: Environment,
   value: ReactiveValue,
   lvalue: Place | null,
   options: MemoizationOptions
@@ -361,8 +367,10 @@ function computeMemoizationInputs(
             : [],
         rvalues: [
           // Conditionals do not alias their test value.
-          ...computeMemoizationInputs(value.consequent, null, options).rvalues,
-          ...computeMemoizationInputs(value.alternate, null, options).rvalues,
+          ...computeMemoizationInputs(env, value.consequent, null, options)
+            .rvalues,
+          ...computeMemoizationInputs(env, value.alternate, null, options)
+            .rvalues,
         ],
       };
     }
@@ -374,8 +382,8 @@ function computeMemoizationInputs(
             ? [{ place: lvalue, level: MemoizationLevel.Conditional }]
             : [],
         rvalues: [
-          ...computeMemoizationInputs(value.left, null, options).rvalues,
-          ...computeMemoizationInputs(value.right, null, options).rvalues,
+          ...computeMemoizationInputs(env, value.left, null, options).rvalues,
+          ...computeMemoizationInputs(env, value.right, null, options).rvalues,
         ],
       };
     }
@@ -389,7 +397,8 @@ function computeMemoizationInputs(
         // Only the final value of the sequence is a true rvalue:
         // values from the sequence's instructions are evaluated
         // as separate nodes
-        rvalues: computeMemoizationInputs(value.value, null, options).rvalues,
+        rvalues: computeMemoizationInputs(env, value.value, null, options)
+          .rvalues,
       };
     }
     case "JsxExpression": {
@@ -595,18 +604,42 @@ function computeMemoizationInputs(
       return {
         lvalues: lvalues,
         rvalues: [
-          ...computeMemoizationInputs(value.value, null, options).rvalues,
+          ...computeMemoizationInputs(env, value.value, null, options).rvalues,
         ],
       };
     }
+    case "MethodCall": {
+      const signature = env.enableNoAliasOptimizations
+        ? getFunctionCallSignature(env, value.property.identifier.type)
+        : null;
+      const operands = [...eachReactiveValueOperand(value)];
+      let lvalues = [];
+      if (lvalue !== null) {
+        lvalues.push({ place: lvalue, level: MemoizationLevel.Memoized });
+      }
+      if (signature?.noAlias === true) {
+        return {
+          lvalues,
+          rvalues: [],
+        };
+      }
+      lvalues.push(
+        ...operands
+          .filter((operand) => isMutableEffect(operand.effect, operand.loc))
+          .map((place) => ({ place, level: MemoizationLevel.Memoized }))
+      );
+      return {
+        lvalues,
+        rvalues: operands,
+      };
+    }
     case "RegExpLiteral":
     case "FunctionExpression":
     case "TaggedTemplateExpression":
-    case "CallExpression":
     case "ArrayExpression":
     case "NewExpression":
     case "ObjectExpression":
-    case "MethodCall":
+    case "CallExpression":
     case "PropertyStore": {
       // All of these instructions may produce new values which must be memoized if
       // reachable from a return value. Any mutable rvalue may alias any other rvalue
@@ -680,10 +713,12 @@ function computePatternLValues(pattern: Pattern): Array<LValueMemoization> {
  * identifier's and scope's dependencies.
  */
 class CollectDependenciesVisitor extends ReactiveFunctionVisitor<State> {
+  env: Environment;
   options: MemoizationOptions;
 
-  constructor(options: MemoizationOptions) {
+  constructor(env: Environment, options: MemoizationOptions) {
     super();
+    this.env = env;
     this.options = options;
   }
 
@@ -695,6 +730,7 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor<State> {
 
     // Determe the level of memoization for this value and the lvalues/rvalues
     const aliasing = computeMemoizationInputs(
+      this.env,
       instruction.value,
       instruction.lvalue,
       this.options

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-frozen-array-noAlias.expect.md

@@ -0,0 +1,60 @@
+
+## Input
+
+```javascript
+// @enableNoAliasOptimizations
+function Component(props) {
+  const x = [];
+  <dif>{x}</dif>;
+  const y = x.map((item) => item);
+  return [x, y];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+  isComponent: false,
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react"; // @enableNoAliasOptimizations
+function Component(props) {
+  const $ = useMemoCache(3);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = [];
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const x = t0;
+  let t1;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = x.map((item) => item);
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  const y = t1;
+  let t2;
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t2 = [x, y];
+    $[2] = t2;
+  } else {
+    t2 = $[2];
+  }
+  return t2;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+  isComponent: false,
+};
+
+```
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-frozen-array-noAlias.js

@@ -0,0 +1,13 @@
+// @enableNoAliasOptimizations
+function Component(props) {
+  const x = [];
+  <dif>{x}</dif>;
+  const y = x.map((item) => item);
+  return [x, y];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+  isComponent: false,
+};

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-mutable-array-mutating-lambda-noAlias.expect.md

@@ -0,0 +1,61 @@
+
+## Input
+
+```javascript
+// @enableNoAliasOptimizations
+function Component(props) {
+  const x = [];
+  const y = x.map((item) => {
+    item.updated = true;
+    return item;
+  });
+  return [x, y];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+  isComponent: false,
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react"; // @enableNoAliasOptimizations
+function Component(props) {
+  const $ = useMemoCache(3);
+  let t0;
+  let x;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    x = [];
+    t0 = x.map((item) => {
+      item.updated = true;
+      return item;
+    });
+    $[0] = t0;
+    $[1] = x;
+  } else {
+    t0 = $[0];
+    x = $[1];
+  }
+  const y = t0;
+  let t1;
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = [x, y];
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+  isComponent: false,
+};
+
+```
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-mutable-array-mutating-lambda-noAlias.js

@@ -0,0 +1,15 @@
+// @enableNoAliasOptimizations
+function Component(props) {
+  const x = [];
+  const y = x.map((item) => {
+    item.updated = true;
+    return item;
+  });
+  return [x, y];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+  isComponent: false,
+};

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-noAlias-escaping-function.expect.md

@@ -0,0 +1,62 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  const f = (item) => item;
+  const x = [...props.items].map(f); // `f` doesn't escape here...
+  return [x, f]; // ...but it does here so it's memoized
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ items: [{ id: 1 }] }],
+  isComponent: false,
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component(props) {
+  const $ = useMemoCache(5);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = (item) => item;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const f = t0;
+  const c_1 = $[1] !== props.items;
+  let t1;
+  if (c_1) {
+    t1 = [...props.items].map(f);
+    $[1] = props.items;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  const x = t1;
+  const c_3 = $[3] !== x;
+  let t2;
+  if (c_3) {
+    t2 = [x, f];
+    $[3] = x;
+    $[4] = t2;
+  } else {
+    t2 = $[4];
+  }
+  return t2;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ items: [{ id: 1 }] }],
+  isComponent: false,
+};
+
+```
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-noAlias-escaping-function.js

@@ -0,0 +1,11 @@
+function Component(props) {
+  const f = (item) => item;
+  const x = [...props.items].map(f); // `f` doesn't escape here...
+  return [x, f]; // ...but it does here so it's memoized
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ items: [{ id: 1 }] }],
+  isComponent: false,
+};

compiler/packages/fixture-test-utils/src/compiler-utils.ts

@@ -28,6 +28,7 @@ export function transformFixtureInput(
   let enableEmitFreeze = null;
   let compilationMode: CompilationMode = "all";
   let enableForest = false;
+  let enableNoAliasOptimizations = false;
 
   if (firstLine.indexOf("@compilationMode(annotation)") !== -1) {
     assert(
@@ -83,6 +84,9 @@ export function transformFixtureInput(
   if (firstLine.includes("@enableForest true")) {
     enableForest = true;
   }
+  if (firstLine.includes("@enableNoAliasOptimizations")) {
+    enableNoAliasOptimizations = true;
+  }
 
   return pluginFn(
     input,
@@ -101,6 +105,7 @@ export function transformFixtureInput(
         ]),
         enableAssumeHooksFollowRulesOfReact,
         enableFunctionCallSignatureOptimizations: true,
+        enableNoAliasOptimizations,
         disableAllMemoization,
         enableTreatHooksAsFunctions,
         inlineUseMemo: true,
@@ -117,7 +122,7 @@ export function transformFixtureInput(
       logger: null,
       gating,
       instrumentForget,
-      panicThreshold: 'ALL_ERRORS',
+      panicThreshold: "ALL_ERRORS",
       noEmit: false,
     },
     includeAst