[compiler][optim] Add Effect.ConditionallyMutateIterator

Adds Effect.ConditionallyMutateIterator, which has the following effects:
- capture for known array, map, and sets
- mutate for all other values

An alternative to this approach could be to add polymorphic shape definitions

Author: mofeiZ

compiler/packages/babel-plugin-react-compiler/src/HIR/Globals.ts

@@ -65,8 +65,6 @@ const UNTYPED_GLOBALS: Set<string> = new Set([
   'Int8Array',
   'Int16Array',
   'Int32Array',
-  'Map',
-  'Set',
   'WeakMap',
   'Uint8Array',
   'Uint8ClampedArray',
@@ -140,7 +138,7 @@ const TYPED_GLOBALS: Array<[string, BuiltInType]> = [
         'from',
         addFunction(DEFAULT_SHAPES, [], {
           positionalParams: [
-            Effect.ConditionallyMutate,
+            Effect.ConditionallyMutateIterator,
             Effect.ConditionallyMutate,
             Effect.ConditionallyMutate,
           ],
@@ -466,7 +464,7 @@ const TYPED_GLOBALS: Array<[string, BuiltInType]> = [
       DEFAULT_SHAPES,
       [],
       {
-        positionalParams: [Effect.ConditionallyMutate],
+        positionalParams: [Effect.ConditionallyMutateIterator],
         restParam: null,
         returnType: {kind: 'Object', shapeId: BuiltInMapId},
         calleeEffect: Effect.Read,
@@ -482,7 +480,7 @@ const TYPED_GLOBALS: Array<[string, BuiltInType]> = [
       DEFAULT_SHAPES,
       [],
       {
-        positionalParams: [Effect.ConditionallyMutate],
+        positionalParams: [Effect.ConditionallyMutateIterator],
         restParam: null,
         returnType: {kind: 'Object', shapeId: BuiltInSetId},
         calleeEffect: Effect.Read,

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

@@ -1396,6 +1396,7 @@ export enum Effect {
   Read = 'read',
   // This reference reads and stores the value
   Capture = 'capture',
+  ConditionallyMutateIterator = 'mutate-iterator?',
   /*
    * This reference *may* write to (mutate) the value. This covers two similar cases:
    * - The compiler is being conservative and assuming that a value *may* be mutated
@@ -1414,11 +1415,11 @@ export enum Effect {
   // This reference may alias to (mutate) the value
   Store = 'store',
 }
-
 export const EffectSchema = z.enum([
   Effect.Read,
   Effect.Mutate,
   Effect.ConditionallyMutate,
+  Effect.ConditionallyMutateIterator,
   Effect.Capture,
   Effect.Store,
   Effect.Freeze,
@@ -1432,6 +1433,7 @@ export function isMutableEffect(
     case Effect.Capture:
     case Effect.Store:
     case Effect.ConditionallyMutate:
+    case Effect.ConditionallyMutateIterator:
     case Effect.Mutate: {
       return true;
     }

compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutableLifetimes.ts

@@ -11,7 +11,10 @@ import {
   Identifier,
   InstructionId,
   InstructionKind,
+  isArrayType,
+  isMapType,
   isRefOrRefValue,
+  isSetType,
   makeInstructionId,
   Place,
 } from '../HIR/HIR';
@@ -90,6 +93,17 @@ function inferPlace(
         infer(place, instrId);
       }
       return;
+    case Effect.ConditionallyMutateIterator: {
+      const identifier = place.identifier;
+      if (
+        !isArrayType(identifier) &&
+        !isSetType(identifier) &&
+        !isMapType(identifier)
+      ) {
+        infer(place, instrId);
+      }
+      return;
+    }
     case Effect.ConditionallyMutate:
     case Effect.Mutate: {
       infer(place, instrId);

compiler/packages/babel-plugin-react-compiler/src/Inference/InferReactivePlaces.ts

@@ -230,6 +230,7 @@ export function inferReactivePlaces(fn: HIRFunction): void {
               case Effect.Capture:
               case Effect.Store:
               case Effect.ConditionallyMutate:
+              case Effect.ConditionallyMutateIterator:
               case Effect.Mutate: {
                 if (isMutable(instruction, operand)) {
                   reactiveIdentifiers.markReactive(operand);

compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts

@@ -29,8 +29,10 @@ import {
   ValueKind,
   ValueReason,
   isArrayType,
+  isMapType,
   isMutableEffect,
   isObjectType,
+  isSetType,
 } from '../HIR/HIR';
 import {FunctionSignature} from '../HIR/ObjectShape';
 import {
@@ -469,6 +471,25 @@ class InferenceState {
         }
         break;
       }
+      case Effect.ConditionallyMutateIterator: {
+        if (
+          valueKind.kind === ValueKind.Mutable ||
+          valueKind.kind === ValueKind.Context
+        ) {
+          if (
+            isArrayType(place.identifier) ||
+            isSetType(place.identifier) ||
+            isMapType(place.identifier)
+          ) {
+            effect = Effect.Capture;
+          } else {
+            effect = Effect.ConditionallyMutate;
+          }
+        } else {
+          effect = Effect.Read;
+        }
+        break;
+      }
       case Effect.Mutate: {
         effect = Effect.Mutate;
         break;
@@ -880,9 +901,7 @@ function inferBlock(
             state.referenceAndRecordEffects(
               freezeActions,
               element.place,
-              isArrayType(element.place.identifier)
-                ? Effect.Capture
-                : Effect.ConditionallyMutate,
+              Effect.ConditionallyMutateIterator,
               ValueReason.Other,
             );
           } else if (element.kind === 'Identifier') {
@@ -1643,7 +1662,13 @@ function inferBlock(
           kind === ValueKind.Mutable || kind === ValueKind.Context;
         let effect;
         let valueKind: AbstractValue;
-        if (!isMutable || isArrayType(instrValue.collection.identifier)) {
+        const iterator = instrValue.collection.identifier;
+        if (
+          !isMutable ||
+          isArrayType(iterator) ||
+          isMapType(iterator) ||
+          isSetType(iterator)
+        ) {
           // Case 1, assume iterator is a separate mutable object
           effect = {
             kind: Effect.Read,
@@ -1684,7 +1709,7 @@ function inferBlock(
         state.referenceAndRecordEffects(
           freezeActions,
           instrValue.iterator,
-          Effect.ConditionallyMutate,
+          Effect.ConditionallyMutateIterator,
           ValueReason.Other,
         );
         /**
@@ -1846,6 +1871,7 @@ export function isKnownMutableEffect(effect: Effect): boolean {
   switch (effect) {
     case Effect.Store:
     case Effect.ConditionallyMutate:
+    case Effect.ConditionallyMutateIterator:
     case Effect.Mutate: {
       return true;
     }
@@ -1949,7 +1975,7 @@ function getArgumentEffect(
         });
       }
       // effects[i] is Effect.Capture | Effect.Read | Effect.Store
-      return Effect.ConditionallyMutate;
+      return Effect.ConditionallyMutateIterator;
     }
   } else {
     return Effect.ConditionallyMutate;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/array-from-arg1-captures-arg0.expect.md

@@ -50,28 +50,55 @@ import { useIdentity, Stringify } from "shared-runtime";
  * (2) the 1st argument might mutate its callee
  */
 function Component(t0) {
-  const $ = _c(4);
+  const $ = _c(10);
   const { value } = t0;
-  const arr = [{ value: "foo" }, { value: "bar" }, { value }];
-  useIdentity();
-  const derived = Array.from(arr, _temp);
   let t1;
-  if ($[0] !== derived) {
-    t1 = derived.at(-1);
-    $[0] = derived;
-    $[1] = t1;
+  let t2;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = { value: "foo" };
+    t2 = { value: "bar" };
+    $[0] = t1;
+    $[1] = t2;
   } else {
-    t1 = $[1];
+    t1 = $[0];
+    t2 = $[1];
   }
-  let t2;
-  if ($[2] !== t1) {
-    t2 = <Stringify>{t1}</Stringify>;
-    $[2] = t1;
-    $[3] = t2;
+  let t3;
+  if ($[2] !== value) {
+    t3 = [t1, t2, { value }];
+    $[2] = value;
+    $[3] = t3;
+  } else {
+    t3 = $[3];
+  }
+  const arr = t3;
+  useIdentity();
+  let t4;
+  if ($[4] !== arr) {
+    t4 = Array.from(arr, _temp);
+    $[4] = arr;
+    $[5] = t4;
+  } else {
+    t4 = $[5];
+  }
+  const derived = t4;
+  let t5;
+  if ($[6] !== derived) {
+    t5 = derived.at(-1);
+    $[6] = derived;
+    $[7] = t5;
+  } else {
+    t5 = $[7];
+  }
+  let t6;
+  if ($[8] !== t5) {
+    t6 = <Stringify>{t5}</Stringify>;
+    $[8] = t5;
+    $[9] = t6;
   } else {
-    t2 = $[3];
+    t6 = $[9];
   }
-  return t2;
+  return t6;
 }
 function _temp(x, idx) {
   return { ...x, id: idx };

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/array-from-captures-arg0.expect.md

@@ -50,28 +50,55 @@ import { useIdentity, Stringify } from "shared-runtime";
  * (2) the 1st argument might mutate its callee
  */
 function Component(t0) {
-  const $ = _c(4);
+  const $ = _c(10);
   const { value } = t0;
-  const arr = [{ value: "foo" }, { value: "bar" }, { value }];
-  useIdentity();
-  const derived = Array.from(arr);
   let t1;
-  if ($[0] !== derived) {
-    t1 = derived.at(-1);
-    $[0] = derived;
-    $[1] = t1;
+  let t2;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = { value: "foo" };
+    t2 = { value: "bar" };
+    $[0] = t1;
+    $[1] = t2;
   } else {
-    t1 = $[1];
+    t1 = $[0];
+    t2 = $[1];
   }
-  let t2;
-  if ($[2] !== t1) {
-    t2 = <Stringify>{t1}</Stringify>;
-    $[2] = t1;
-    $[3] = t2;
+  let t3;
+  if ($[2] !== value) {
+    t3 = [t1, t2, { value }];
+    $[2] = value;
+    $[3] = t3;
+  } else {
+    t3 = $[3];
+  }
+  const arr = t3;
+  useIdentity();
+  let t4;
+  if ($[4] !== arr) {
+    t4 = Array.from(arr);
+    $[4] = arr;
+    $[5] = t4;
+  } else {
+    t4 = $[5];
+  }
+  const derived = t4;
+  let t5;
+  if ($[6] !== derived) {
+    t5 = derived.at(-1);
+    $[6] = derived;
+    $[7] = t5;
+  } else {
+    t5 = $[7];
+  }
+  let t6;
+  if ($[8] !== t5) {
+    t6 = <Stringify>{t5}</Stringify>;
+    $[8] = t5;
+    $[9] = t6;
   } else {
-    t2 = $[3];
+    t6 = $[9];
   }
-  return t2;
+  return t6;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/array-from-maybemutates-arg0.expect.md

@@ -7,7 +7,7 @@ import {mutateAndReturn, Stringify, useIdentity} from 'shared-runtime';
 function Component({value}) {
   const arr = [{value: 'foo'}, {value: 'bar'}, {value}];
   useIdentity();
-  const derived = Array.from(arr, mutateAndReturn);
+  const derived = Array.from(arr).map(mutateAndReturn);
   return (
     <Stringify>
       {derived.at(0)}
@@ -19,7 +19,7 @@ function Component({value}) {
 export const FIXTURE_ENTRYPOINT = {
   fn: Component,
   params: [{value: 5}],
-  sequentialRenders: [{value: 5}, {value: 6}, {value: 6}],
+  sequentialRenders: [{value: 5}, {value: 6}, {value: 6}, {value: 7}],
 };
 
 ```
@@ -35,7 +35,7 @@ function Component(t0) {
   const { value } = t0;
   const arr = [{ value: "foo" }, { value: "bar" }, { value }];
   useIdentity();
-  const derived = Array.from(arr, mutateAndReturn);
+  const derived = Array.from(arr).map(mutateAndReturn);
   let t1;
   if ($[0] !== derived) {
     t1 = derived.at(0);
@@ -72,12 +72,13 @@ function Component(t0) {
 export const FIXTURE_ENTRYPOINT = {
   fn: Component,
   params: [{ value: 5 }],
-  sequentialRenders: [{ value: 5 }, { value: 6 }, { value: 6 }],
+  sequentialRenders: [{ value: 5 }, { value: 6 }, { value: 6 }, { value: 7 }],
 };
 
 ```
       
 ### Eval output
 (kind: ok) <div>{"children":[{"value":"foo","wat0":"joe"},{"value":5,"wat0":"joe"}]}</div>
 <div>{"children":[{"value":"foo","wat0":"joe"},{"value":6,"wat0":"joe"}]}</div>
-<div>{"children":[{"value":"foo","wat0":"joe"},{"value":6,"wat0":"joe"}]}</div>
+<div>{"children":[{"value":"foo","wat0":"joe"},{"value":6,"wat0":"joe"}]}</div>
+<div>{"children":[{"value":"foo","wat0":"joe"},{"value":7,"wat0":"joe"}]}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/array-from-maybemutates-arg0.js

@@ -3,7 +3,7 @@ import {mutateAndReturn, Stringify, useIdentity} from 'shared-runtime';
 function Component({value}) {
   const arr = [{value: 'foo'}, {value: 'bar'}, {value}];
   useIdentity();
-  const derived = Array.from(arr, mutateAndReturn);
+  const derived = Array.from(arr).map(mutateAndReturn);
   return (
     <Stringify>
       {derived.at(0)}
@@ -15,5 +15,5 @@ function Component({value}) {
 export const FIXTURE_ENTRYPOINT = {
   fn: Component,
   params: [{value: 5}],
-  sequentialRenders: [{value: 5}, {value: 6}, {value: 6}],
+  sequentialRenders: [{value: 5}, {value: 6}, {value: 6}, {value: 7}],
 };