[compiler][optim] more shapes for mixedreadonly

- Add `at`, `indexOf`, and `includes`
- Optimize MixedReadOnly which is currently only used by hook return values. Hook return values are typed as Frozen, this change propagates that to return values of aliasing function calls (such as `at`). One potential issue is that  developers may pass `enableAssumeHooksFollowRulesOfReact:false` and set `transitiveMixedData`, expecting their transitive mixed data to be mutable. This is a bit of an edge case and already doesn't have clear semantics.

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

@@ -1126,9 +1126,32 @@ export class Environment {
     );
   }
 
+  getFallthroughPropertyType(
+    receiver: Type,
+    _property: Type,
+  ): BuiltInType | PolyType | null {
+    let shapeId = null;
+    if (receiver.kind === 'Object' || receiver.kind === 'Function') {
+      shapeId = receiver.shapeId;
+    }
+
+    if (shapeId !== null) {
+      const shape = this.#shapes.get(shapeId);
+
+      CompilerError.invariant(shape !== undefined, {
+        reason: `[HIR] Forget internal error: cannot resolve shape ${shapeId}`,
+        description: null,
+        loc: null,
+        suggestions: null,
+      });
+      return shape.properties.get('*') ?? null;
+    }
+    return null;
+  }
+
   getPropertyType(
     receiver: Type,
-    property: string,
+    property: string | number,
   ): BuiltInType | PolyType | null {
     let shapeId = null;
     if (receiver.kind === 'Object' || receiver.kind === 'Function') {
@@ -1146,17 +1169,19 @@ export class Environment {
         loc: null,
         suggestions: null,
       });
-      let value =
-        shape.properties.get(property) ?? shape.properties.get('*') ?? null;
-      if (value === null && isHookName(property)) {
-        value = this.#getCustomHookType();
+      if (typeof property === 'string') {
+        return (
+          shape.properties.get(property) ??
+          shape.properties.get('*') ??
+          (isHookName(property) ? this.#getCustomHookType() : null)
+        );
+      } else {
+        return shape.properties.get('*') ?? null;
       }
-      return value;
-    } else if (isHookName(property)) {
+    } else if (typeof property === 'string' && isHookName(property)) {
       return this.#getCustomHookType();
-    } else {
-      return null;
     }
+    return null;
   }
 
   getFunctionSignature(type: FunctionType): FunctionSignature | null {

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

@@ -535,6 +535,30 @@ addObject(BUILTIN_SHAPES, BuiltInRefValueId, [
   ['*', {kind: 'Object', shapeId: BuiltInRefValueId}],
 ]);
 
+/**
+ * MixedReadOnly =
+ *   | primitive
+ *   | simple objects (Record<string, MixedReadOnly>)
+ *   | Array<MixedReadOnly>
+ *
+ * APIs such as Relay — but also Flux and other data stores — often return a
+ * union of types with some interesting properties in terms of analysis.
+ *
+ * Given this constraint, if data came from Relay, then we should be able to
+ * infer things like `data.items.map(): Array`. That may seem like a leap at
+ * first but remember, we assume you're not patching builtins. Thus the only way
+ * data.items.map can exist and be a function, given the above set of data types
+ * and builtin JS methods, is if `data.items` was an Array, and `data.items.map`
+ * is therefore calling Array.prototype.map. Then we know that function returns
+ * an Array as well. This relies on the fact that map() is being called, so if
+ * data.items was some other type it would error at runtime - so it's sound.
+ *
+ * Note that this shape is currently only used for hook return values, which
+ * means that it's safe to type aliasing method-call return kinds as `Frozen`.
+ *
+ * Also note that all newly created arrays from method-calls (e.g. `.map`)
+ * have the appropriate mutable `BuiltInArray` shape
+ */
 addObject(BUILTIN_SHAPES, BuiltInMixedReadonlyId, [
   [
     'toString',
@@ -546,6 +570,36 @@ addObject(BUILTIN_SHAPES, BuiltInMixedReadonlyId, [
       returnValueKind: ValueKind.Primitive,
     }),
   ],
+  [
+    'indexOf',
+    addFunction(BUILTIN_SHAPES, [], {
+      positionalParams: [],
+      restParam: Effect.Read,
+      returnType: {kind: 'Primitive'},
+      calleeEffect: Effect.Read,
+      returnValueKind: ValueKind.Primitive,
+    }),
+  ],
+  [
+    'includes',
+    addFunction(BUILTIN_SHAPES, [], {
+      positionalParams: [],
+      restParam: Effect.Read,
+      returnType: {kind: 'Primitive'},
+      calleeEffect: Effect.Read,
+      returnValueKind: ValueKind.Primitive,
+    }),
+  ],
+  [
+    'at',
+    addFunction(BUILTIN_SHAPES, [], {
+      positionalParams: [Effect.Read],
+      restParam: null,
+      returnType: {kind: 'Object', shapeId: BuiltInMixedReadonlyId},
+      calleeEffect: Effect.Capture,
+      returnValueKind: ValueKind.Frozen,
+    }),
+  ],
   [
     'map',
     addFunction(BUILTIN_SHAPES, [], {
@@ -642,9 +696,9 @@ addObject(BUILTIN_SHAPES, BuiltInMixedReadonlyId, [
     addFunction(BUILTIN_SHAPES, [], {
       positionalParams: [],
       restParam: Effect.ConditionallyMutate,
-      returnType: {kind: 'Poly'},
+      returnType: {kind: 'Object', shapeId: BuiltInMixedReadonlyId},
       calleeEffect: Effect.ConditionallyMutate,
-      returnValueKind: ValueKind.Mutable,
+      returnValueKind: ValueKind.Frozen,
       noAlias: true,
       mutableOnlyIfOperandsAreMutable: true,
     }),

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

@@ -60,7 +60,15 @@ export type PropType = {
   kind: 'Property';
   objectType: Type;
   objectName: string;
-  propertyName: PropertyLiteral;
+  propertyName:
+    | {
+        kind: 'literal';
+        value: PropertyLiteral;
+      }
+    | {
+        kind: 'computed';
+        value: Type;
+      };
 };
 
 export type ObjectMethod = {

compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts

@@ -307,11 +307,26 @@ function* generateInstructionTypes(
         kind: 'Property',
         objectType: value.object.identifier.type,
         objectName: getName(names, value.object.identifier.id),
-        propertyName: value.property,
+        propertyName: {
+          kind: 'literal',
+          value: value.property,
+        },
       });
       break;
     }
 
+    case 'ComputedLoad': {
+      yield equation(left, {
+        kind: 'Property',
+        objectType: value.object.identifier.type,
+        objectName: getName(names, value.object.identifier.id),
+        propertyName: {
+          kind: 'computed',
+          value: value.property.identifier.type,
+        },
+      });
+      break;
+    }
     case 'MethodCall': {
       const returnType = makeType();
       yield equation(value.property.identifier.type, {
@@ -336,7 +351,10 @@ function* generateInstructionTypes(
               kind: 'Property',
               objectType: value.value.identifier.type,
               objectName: getName(names, value.value.identifier.id),
-              propertyName: makePropertyLiteral(propertyName),
+              propertyName: {
+                kind: 'literal',
+                value: makePropertyLiteral(propertyName),
+              },
             });
           } else {
             break;
@@ -353,7 +371,10 @@ function* generateInstructionTypes(
                 kind: 'Property',
                 objectType: value.value.identifier.type,
                 objectName: getName(names, value.value.identifier.id),
-                propertyName: makePropertyLiteral(property.key.name),
+                propertyName: {
+                  kind: 'literal',
+                  value: makePropertyLiteral(property.key.name),
+                },
               });
             }
           }
@@ -410,7 +431,6 @@ function* generateInstructionTypes(
     case 'RegExpLiteral':
     case 'MetaProperty':
     case 'ComputedStore':
-    case 'ComputedLoad':
     case 'Await':
     case 'GetIterator':
     case 'IteratorNext':
@@ -454,12 +474,13 @@ class Unifier {
         return;
       }
       const objectType = this.get(tB.objectType);
-      let propertyType;
-      if (typeof tB.propertyName === 'number') {
-        propertyType = null;
-      } else {
-        propertyType = this.env.getPropertyType(objectType, tB.propertyName);
-      }
+      const propertyType =
+        tB.propertyName.kind === 'literal'
+          ? this.env.getPropertyType(objectType, tB.propertyName.value)
+          : this.env.getFallthroughPropertyType(
+              objectType,
+              tB.propertyName.value,
+            );
       if (propertyType !== null) {
         this.unify(tA, propertyType);
       }
@@ -677,7 +698,11 @@ class Unifier {
 const RefLikeNameRE = /^(?:[a-zA-Z$_][a-zA-Z$_0-9]*)Ref$|^ref$/;
 
 function isRefLikeName(t: PropType): boolean {
-  return RefLikeNameRE.test(t.objectName) && t.propertyName === 'current';
+  return (
+    t.propertyName.kind === 'literal' &&
+    RefLikeNameRE.test(t.objectName) &&
+    t.propertyName.value === 'current'
+  );
 }
 
 function tryUnionTypes(ty1: Type, ty2: Type): Type | null {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/relay-transitive-mixeddata.expect.md

@@ -0,0 +1,61 @@
+
+## Input
+
+```javascript
+import {useFragment} from 'shared-runtime';
+
+/**
+ * React compiler should infer that the returned value is a primitive and avoid
+ * memoizing it.
+ */
+function useRelayData({query, idx}) {
+  'use memo';
+  const data = useFragment('', query);
+  return data.a[idx].toString();
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useRelayData,
+  params: [{query: '', idx: 0}],
+  sequentialRenders: [
+    {query: '', idx: 0},
+    {query: '', idx: 0},
+    {query: '', idx: 1},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { useFragment } from "shared-runtime";
+
+/**
+ * React compiler should infer that the returned value is a primitive and avoid
+ * memoizing it.
+ */
+function useRelayData(t0) {
+  "use memo";
+  const { query, idx } = t0;
+
+  const data = useFragment("", query);
+  return data.a[idx].toString();
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useRelayData,
+  params: [{ query: "", idx: 0 }],
+  sequentialRenders: [
+    { query: "", idx: 0 },
+    { query: "", idx: 0 },
+    { query: "", idx: 1 },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) "1"
+"1"
+"2"

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/relay-transitive-mixeddata.js

@@ -0,0 +1,21 @@
+import {useFragment} from 'shared-runtime';
+
+/**
+ * React compiler should infer that the returned value is a primitive and avoid
+ * memoizing it.
+ */
+function useRelayData({query, idx}) {
+  'use memo';
+  const data = useFragment('', query);
+  return data.a[idx].toString();
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useRelayData,
+  params: [{query: '', idx: 0}],
+  sequentialRenders: [
+    {query: '', idx: 0},
+    {query: '', idx: 0},
+    {query: '', idx: 1},
+  ],
+};