[compiler] Add types for WeakMap, WeakSet, and reanimated shared values

Building on mofeiz's recent work to type constructors. Also, types for reanimated values which are useful in the next PR.

ghstack-source-id: 1c81e21
Pull Request resolved: #33077

Author: josephsavona

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

@@ -25,6 +25,9 @@ import {
   BuiltInUseRefId,
   BuiltInUseStateId,
   BuiltInUseTransitionId,
+  BuiltInWeakMapId,
+  BuiltInWeakSetId,
+  ReanimatedSharedValueId,
   ShapeRegistry,
   addFunction,
   addHook,
@@ -491,6 +494,38 @@ const TYPED_GLOBALS: Array<[string, BuiltInType]> = [
       true,
     ),
   ],
+  [
+    'WeakMap',
+    addFunction(
+      DEFAULT_SHAPES,
+      [],
+      {
+        positionalParams: [Effect.ConditionallyMutateIterator],
+        restParam: null,
+        returnType: {kind: 'Object', shapeId: BuiltInWeakMapId},
+        calleeEffect: Effect.Read,
+        returnValueKind: ValueKind.Mutable,
+      },
+      null,
+      true,
+    ),
+  ],
+  [
+    'WeakSet',
+    addFunction(
+      DEFAULT_SHAPES,
+      [],
+      {
+        positionalParams: [Effect.ConditionallyMutateIterator],
+        restParam: null,
+        returnType: {kind: 'Object', shapeId: BuiltInWeakSetId},
+        calleeEffect: Effect.Read,
+        returnValueKind: ValueKind.Mutable,
+      },
+      null,
+      true,
+    ),
+  ],
   // TODO: rest of Global objects
 ];
 
@@ -908,7 +943,7 @@ export function getReanimatedModuleType(registry: ShapeRegistry): ObjectType {
       addHook(registry, {
         positionalParams: [],
         restParam: Effect.Freeze,
-        returnType: {kind: 'Poly'},
+        returnType: {kind: 'Object', shapeId: ReanimatedSharedValueId},
         returnValueKind: ValueKind.Mutable,
         noAlias: true,
         calleeEffect: Effect.Read,

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

@@ -1725,6 +1725,18 @@ export function isRefOrRefValue(id: Identifier): boolean {
   return isUseRefType(id) || isRefValueType(id);
 }
 
+/*
+ * Returns true if the type is a Ref or a custom user type that acts like a ref when it
+ * shouldn't. For now the only other case of this is Reanimated's shared values.
+ */
+export function isRefOrRefLikeMutableType(type: Type): boolean {
+  return (
+    type.kind === 'Object' &&
+    (type.shapeId === 'BuiltInUseRefId' ||
+      type.shapeId == 'ReanimatedSharedValueId')
+  );
+}
+
 export function isSetStateType(id: Identifier): boolean {
   return id.type.kind === 'Function' && id.type.shapeId === 'BuiltInSetState';
 }

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

@@ -203,6 +203,8 @@ export const BuiltInPropsId = 'BuiltInProps';
 export const BuiltInArrayId = 'BuiltInArray';
 export const BuiltInSetId = 'BuiltInSet';
 export const BuiltInMapId = 'BuiltInMap';
+export const BuiltInWeakSetId = 'BuiltInWeakSet';
+export const BuiltInWeakMapId = 'BuiltInWeakMap';
 export const BuiltInFunctionId = 'BuiltInFunction';
 export const BuiltInJsxId = 'BuiltInJsx';
 export const BuiltInObjectId = 'BuiltInObject';
@@ -225,6 +227,9 @@ export const BuiltInStartTransitionId = 'BuiltInStartTransition';
 export const BuiltInFireId = 'BuiltInFire';
 export const BuiltInFireFunctionId = 'BuiltInFireFunction';
 
+// See getReanimatedModuleType() in Globals.ts — this is part of supporting Reanimated's ref-like types
+export const ReanimatedSharedValueId = 'ReanimatedSharedValueId';
+
 // ShapeRegistry with default definitions for built-ins.
 export const BUILTIN_SHAPES: ShapeRegistry = new Map();
 
@@ -764,6 +769,101 @@ addObject(BUILTIN_SHAPES, BuiltInMapId, [
   ],
 ]);
 
+addObject(BUILTIN_SHAPES, BuiltInWeakSetId, [
+  [
+    /**
+     * add(value)
+     * Parameters
+     *   value: the value of the element to add to the Set object.
+     * Returns the Set object with added value.
+     */
+    'add',
+    addFunction(BUILTIN_SHAPES, [], {
+      positionalParams: [Effect.Capture],
+      restParam: null,
+      returnType: {kind: 'Object', shapeId: BuiltInWeakSetId},
+      calleeEffect: Effect.Store,
+      // returnValueKind is technically dependent on the ValueKind of the set itself
+      returnValueKind: ValueKind.Mutable,
+    }),
+  ],
+  [
+    /**
+     * setInstance.delete(value)
+     * Returns true if value was already in Set; otherwise false.
+     */
+    'delete',
+    addFunction(BUILTIN_SHAPES, [], {
+      positionalParams: [Effect.Read],
+      restParam: null,
+      returnType: PRIMITIVE_TYPE,
+      calleeEffect: Effect.Store,
+      returnValueKind: ValueKind.Primitive,
+    }),
+  ],
+  [
+    'has',
+    addFunction(BUILTIN_SHAPES, [], {
+      positionalParams: [Effect.Read],
+      restParam: null,
+      returnType: PRIMITIVE_TYPE,
+      calleeEffect: Effect.Read,
+      returnValueKind: ValueKind.Primitive,
+    }),
+  ],
+]);
+
+addObject(BUILTIN_SHAPES, BuiltInWeakMapId, [
+  [
+    'delete',
+    addFunction(BUILTIN_SHAPES, [], {
+      positionalParams: [Effect.Read],
+      restParam: null,
+      returnType: PRIMITIVE_TYPE,
+      calleeEffect: Effect.Store,
+      returnValueKind: ValueKind.Primitive,
+    }),
+  ],
+  [
+    'get',
+    addFunction(BUILTIN_SHAPES, [], {
+      positionalParams: [Effect.Read],
+      restParam: null,
+      returnType: {kind: 'Poly'},
+      calleeEffect: Effect.Capture,
+      returnValueKind: ValueKind.Mutable,
+    }),
+  ],
+  [
+    'has',
+    addFunction(BUILTIN_SHAPES, [], {
+      positionalParams: [Effect.Read],
+      restParam: null,
+      returnType: PRIMITIVE_TYPE,
+      calleeEffect: Effect.Read,
+      returnValueKind: ValueKind.Primitive,
+    }),
+  ],
+  [
+    /**
+     * Params
+     *   key: the key of the element to add to the Map object. The key may be
+     *   any JavaScript type (any primitive value or any type of JavaScript
+     *   object).
+     *   value: the value of the element to add to the Map object.
+     * Returns the Map object.
+     */
+    'set',
+    addFunction(BUILTIN_SHAPES, [], {
+      positionalParams: [Effect.Capture, Effect.Capture],
+      restParam: null,
+      returnType: {kind: 'Object', shapeId: BuiltInWeakMapId},
+      calleeEffect: Effect.Store,
+      returnValueKind: ValueKind.Mutable,
+    }),
+  ],
+]);
+
 addObject(BUILTIN_SHAPES, BuiltInUseStateId, [
   ['0', {kind: 'Poly'}],
   [

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakmap-constructor.expect.md

@@ -0,0 +1,176 @@
+
+## Input
+
+```javascript
+import {ValidateMemoization} from 'shared-runtime';
+
+function Component({a, b, c}) {
+  const map = new WeakMap();
+  const mapAlias = map.set(a, 0);
+  mapAlias.set(c, 0);
+
+  const hasB = map.has(b);
+
+  return (
+    <>
+      <ValidateMemoization inputs={[a, c]} output={map} />
+      <ValidateMemoization inputs={[a, c]} output={mapAlias} />
+      <ValidateMemoization inputs={[b]} output={[hasB]} />
+    </>
+  );
+}
+
+const v1 = {value: 1};
+const v2 = {value: 2};
+const v3 = {value: 3};
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: v1, b: v1, c: v1}],
+  sequentialRenders: [
+    {a: v1, b: v1, c: v1},
+    {a: v2, b: v1, c: v1},
+    {a: v1, b: v1, c: v1},
+    {a: v1, b: v2, c: v1},
+    {a: v1, b: v1, c: v1},
+    {a: v3, b: v3, c: v1},
+    {a: v3, b: v3, c: v1},
+    {a: v1, b: v1, c: v1},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { ValidateMemoization } from "shared-runtime";
+
+function Component(t0) {
+  const $ = _c(27);
+  const { a, b, c } = t0;
+  let map;
+  let mapAlias;
+  if ($[0] !== a || $[1] !== c) {
+    map = new WeakMap();
+    mapAlias = map.set(a, 0);
+    mapAlias.set(c, 0);
+    $[0] = a;
+    $[1] = c;
+    $[2] = map;
+    $[3] = mapAlias;
+  } else {
+    map = $[2];
+    mapAlias = $[3];
+  }
+
+  const hasB = map.has(b);
+  let t1;
+  if ($[4] !== a || $[5] !== c) {
+    t1 = [a, c];
+    $[4] = a;
+    $[5] = c;
+    $[6] = t1;
+  } else {
+    t1 = $[6];
+  }
+  let t2;
+  if ($[7] !== map || $[8] !== t1) {
+    t2 = <ValidateMemoization inputs={t1} output={map} />;
+    $[7] = map;
+    $[8] = t1;
+    $[9] = t2;
+  } else {
+    t2 = $[9];
+  }
+  let t3;
+  if ($[10] !== a || $[11] !== c) {
+    t3 = [a, c];
+    $[10] = a;
+    $[11] = c;
+    $[12] = t3;
+  } else {
+    t3 = $[12];
+  }
+  let t4;
+  if ($[13] !== mapAlias || $[14] !== t3) {
+    t4 = <ValidateMemoization inputs={t3} output={mapAlias} />;
+    $[13] = mapAlias;
+    $[14] = t3;
+    $[15] = t4;
+  } else {
+    t4 = $[15];
+  }
+  let t5;
+  if ($[16] !== b) {
+    t5 = [b];
+    $[16] = b;
+    $[17] = t5;
+  } else {
+    t5 = $[17];
+  }
+  let t6;
+  if ($[18] !== hasB) {
+    t6 = [hasB];
+    $[18] = hasB;
+    $[19] = t6;
+  } else {
+    t6 = $[19];
+  }
+  let t7;
+  if ($[20] !== t5 || $[21] !== t6) {
+    t7 = <ValidateMemoization inputs={t5} output={t6} />;
+    $[20] = t5;
+    $[21] = t6;
+    $[22] = t7;
+  } else {
+    t7 = $[22];
+  }
+  let t8;
+  if ($[23] !== t2 || $[24] !== t4 || $[25] !== t7) {
+    t8 = (
+      <>
+        {t2}
+        {t4}
+        {t7}
+      </>
+    );
+    $[23] = t2;
+    $[24] = t4;
+    $[25] = t7;
+    $[26] = t8;
+  } else {
+    t8 = $[26];
+  }
+  return t8;
+}
+
+const v1 = { value: 1 };
+const v2 = { value: 2 };
+const v3 = { value: 3 };
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: v1, b: v1, c: v1 }],
+  sequentialRenders: [
+    { a: v1, b: v1, c: v1 },
+    { a: v2, b: v1, c: v1 },
+    { a: v1, b: v1, c: v1 },
+    { a: v1, b: v2, c: v1 },
+    { a: v1, b: v1, c: v1 },
+    { a: v3, b: v3, c: v1 },
+    { a: v3, b: v3, c: v1 },
+    { a: v1, b: v1, c: v1 },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"inputs":[{"value":1},"[[ cyclic ref *2 ]]"],"output":{}}</div><div>{"inputs":[{"value":1},"[[ cyclic ref *2 ]]"],"output":{}}</div><div>{"inputs":[{"value":1}],"output":[true]}</div>
+<div>{"inputs":[{"value":2},{"value":1}],"output":{}}</div><div>{"inputs":[{"value":2},{"value":1}],"output":{}}</div><div>{"inputs":[{"value":1}],"output":[true]}</div>
+<div>{"inputs":[{"value":1},"[[ cyclic ref *2 ]]"],"output":{}}</div><div>{"inputs":[{"value":1},"[[ cyclic ref *2 ]]"],"output":{}}</div><div>{"inputs":[{"value":1}],"output":[true]}</div>
+<div>{"inputs":[{"value":1},"[[ cyclic ref *2 ]]"],"output":{}}</div><div>{"inputs":[{"value":1},"[[ cyclic ref *2 ]]"],"output":{}}</div><div>{"inputs":[{"value":2}],"output":[false]}</div>
+<div>{"inputs":[{"value":1},"[[ cyclic ref *2 ]]"],"output":{}}</div><div>{"inputs":[{"value":1},"[[ cyclic ref *2 ]]"],"output":{}}</div><div>{"inputs":[{"value":1}],"output":[true]}</div>
+<div>{"inputs":[{"value":3},{"value":1}],"output":{}}</div><div>{"inputs":[{"value":3},{"value":1}],"output":{}}</div><div>{"inputs":[{"value":3}],"output":[true]}</div>
+<div>{"inputs":[{"value":3},{"value":1}],"output":{}}</div><div>{"inputs":[{"value":3},{"value":1}],"output":{}}</div><div>{"inputs":[{"value":3}],"output":[true]}</div>
+<div>{"inputs":[{"value":1},"[[ cyclic ref *2 ]]"],"output":{}}</div><div>{"inputs":[{"value":1},"[[ cyclic ref *2 ]]"],"output":{}}</div><div>{"inputs":[{"value":1}],"output":[true]}</div>