[compiler][optim] infer mixedReadOnly for numeric and computed properties (#32593)

Expand type inference to infer mixedReadOnly types for numeric and
computed property accesses.
```js
function Component({idx})
  const data = useFragment(...)
  // we want to type `posts` correctly as Array
  const posts = data.viewers[idx].posts.slice(0, 5);
  // ...
}
```
---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with [ReviewStack](https://reviewstack.dev/facebook/react/pull/32593).
* #32596
* #32595
* #32594
* __->__ #32593
* #32522
* #32521

Author: mofeiZ

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/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},
+  ],
+};