[compiler][ez] Fix reanimated custom type defs for imports (#31137)

When we added support for Reanimated, we didn't distinguish between true
globals (i.e. identifiers with no static resolutions), module types, and
imports #29188. For the past 3-4 months, Reanimated imports were not
being matched to the correct hook / function shape we match globals and
module imports against two different registries.

This PR fixes our support for Reanimated library functions imported
under `react-native-reanimated`. See test fixtures for details

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

@@ -16,7 +16,7 @@ import {
   DEFAULT_SHAPES,
   Global,
   GlobalRegistry,
-  installReAnimatedTypes,
+  getReanimatedModuleType,
   installTypeConfig,
 } from './Globals';
 import {
@@ -688,7 +688,8 @@ export class Environment {
     }
 
     if (config.enableCustomTypeDefinitionForReanimated) {
-      installReAnimatedTypes(this.#globals, this.#shapes);
+      const reanimatedModuleType = getReanimatedModuleType(this.#shapes);
+      this.#moduleTypes.set(REANIMATED_MODULE_NAME, reanimatedModuleType);
     }
 
     this.#contextIdentifiers = contextIdentifiers;
@@ -734,11 +735,11 @@ export class Environment {
   }
 
   #resolveModuleType(moduleName: string, loc: SourceLocation): Global | null {
-    if (this.config.moduleTypeProvider == null) {
-      return null;
-    }
     let moduleType = this.#moduleTypes.get(moduleName);
     if (moduleType === undefined) {
+      if (this.config.moduleTypeProvider == null) {
+        return null;
+      }
       const unparsedModuleConfig = this.config.moduleTypeProvider(moduleName);
       if (unparsedModuleConfig != null) {
         const parsedModuleConfig = TypeSchema.safeParse(unparsedModuleConfig);
@@ -957,6 +958,8 @@ export class Environment {
   }
 }
 
+const REANIMATED_MODULE_NAME = 'react-native-reanimated';
+
 // From https://github.com/facebook/react/blob/main/packages/eslint-plugin-react-hooks/src/RulesOfHooks.js#LL18C1-L23C2
 export function isHookName(name: string): boolean {
   return /^use[A-Z0-9]/.test(name);

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

@@ -25,7 +25,7 @@ import {
   addHook,
   addObject,
 } from './ObjectShape';
-import {BuiltInType, PolyType} from './Types';
+import {BuiltInType, ObjectType, PolyType} from './Types';
 import {TypeConfig} from './TypeSchema';
 import {assertExhaustive} from '../Utils/utils';
 import {isHookName} from './Environment';
@@ -652,10 +652,7 @@ export function installTypeConfig(
   }
 }
 
-export function installReAnimatedTypes(
-  globals: GlobalRegistry,
-  registry: ShapeRegistry,
-): void {
+export function getReanimatedModuleType(registry: ShapeRegistry): ObjectType {
   // hooks that freeze args and return frozen value
   const frozenHooks = [
     'useFrameCallback',
@@ -665,8 +662,9 @@ export function installReAnimatedTypes(
     'useAnimatedReaction',
     'useWorkletCallback',
   ];
+  const reanimatedType: Array<[string, BuiltInType]> = [];
   for (const hook of frozenHooks) {
-    globals.set(
+    reanimatedType.push([
       hook,
       addHook(registry, {
         positionalParams: [],
@@ -677,7 +675,7 @@ export function installReAnimatedTypes(
         calleeEffect: Effect.Read,
         hookKind: 'Custom',
       }),
-    );
+    ]);
   }
 
   /**
@@ -686,7 +684,7 @@ export function installReAnimatedTypes(
    */
   const mutableHooks = ['useSharedValue', 'useDerivedValue'];
   for (const hook of mutableHooks) {
-    globals.set(
+    reanimatedType.push([
       hook,
       addHook(registry, {
         positionalParams: [],
@@ -697,7 +695,7 @@ export function installReAnimatedTypes(
         calleeEffect: Effect.Read,
         hookKind: 'Custom',
       }),
-    );
+    ]);
   }
 
   // functions that return mutable value
@@ -711,7 +709,7 @@ export function installReAnimatedTypes(
     'executeOnUIRuntimeSync',
   ];
   for (const fn of funcs) {
-    globals.set(
+    reanimatedType.push([
       fn,
       addFunction(registry, [], {
         positionalParams: [],
@@ -721,6 +719,8 @@ export function installReAnimatedTypes(
         returnValueKind: ValueKind.Mutable,
         noAlias: true,
       }),
-    );
+    ]);
   }
+
+  return addObject(registry, null, reanimatedType);
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-non-imported-reanimated-shared-value-writes.expect.md

@@ -0,0 +1,36 @@
+
+## Input
+
+```javascript
+// @enableCustomTypeDefinitionForReanimated
+
+/**
+ * Test that a global (i.e. non-imported) useSharedValue is treated as an
+ * unknown hook.
+ */
+function SomeComponent() {
+  const sharedVal = useSharedValue(0);
+  return (
+    <Button
+      onPress={() => (sharedVal.value = Math.random())}
+      title="Randomize"
+    />
+  );
+}
+
+```
+
+
+## Error
+
+```
+   9 |   return (
+  10 |     <Button
+> 11 |       onPress={() => (sharedVal.value = Math.random())}
+     |                       ^^^^^^^^^ InvalidReact: Mutating a value returned from a function whose return value should not be mutated. Found mutation of `sharedVal` (11:11)
+  12 |       title="Randomize"
+  13 |     />
+  14 |   );
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-non-imported-reanimated-shared-value-writes.jsx

@@ -0,0 +1,15 @@
+// @enableCustomTypeDefinitionForReanimated
+
+/**
+ * Test that a global (i.e. non-imported) useSharedValue is treated as an
+ * unknown hook.
+ */
+function SomeComponent() {
+  const sharedVal = useSharedValue(0);
+  return (
+    <Button
+      onPress={() => (sharedVal.value = Math.random())}
+      title="Randomize"
+    />
+  );
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reanimated-no-memo-arg.expect.md

@@ -3,6 +3,7 @@
 
 ```javascript
 // @enableCustomTypeDefinitionForReanimated
+import {useAnimatedProps} from 'react-native-reanimated';
 function Component() {
   const radius = useSharedValue(50);
 
@@ -38,6 +39,7 @@ export const FIXTURE_ENTRYPOINT = {
 
 ```javascript
 import { c as _c } from "react/compiler-runtime"; // @enableCustomTypeDefinitionForReanimated
+import { useAnimatedProps } from "react-native-reanimated";
 function Component() {
   const $ = _c(2);
   const radius = useSharedValue(50);

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reanimated-no-memo-arg.js

@@ -1,4 +1,5 @@
 // @enableCustomTypeDefinitionForReanimated
+import {useAnimatedProps} from 'react-native-reanimated';
 function Component() {
   const radius = useSharedValue(50);
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reanimated-shared-value-writes.expect.md

@@ -0,0 +1,57 @@
+
+## Input
+
+```javascript
+// @enableCustomTypeDefinitionForReanimated
+import {useSharedValue} from 'react-native-reanimated';
+
+/**
+ * https://docs.swmansion.com/react-native-reanimated/docs/2.x/api/hooks/useSharedValue/
+ *
+ * Test that shared values are treated as ref-like, i.e. allowing writes outside
+ * of render
+ */
+function SomeComponent() {
+  const sharedVal = useSharedValue(0);
+  return (
+    <Button
+      onPress={() => (sharedVal.value = Math.random())}
+      title="Randomize"
+    />
+  );
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableCustomTypeDefinitionForReanimated
+import { useSharedValue } from "react-native-reanimated";
+
+/**
+ * https://docs.swmansion.com/react-native-reanimated/docs/2.x/api/hooks/useSharedValue/
+ *
+ * Test that shared values are treated as ref-like, i.e. allowing writes outside
+ * of render
+ */
+function SomeComponent() {
+  const $ = _c(3);
+  const sharedVal = useSharedValue(0);
+
+  const T0 = Button;
+  const t0 = () => (sharedVal.value = Math.random());
+  let t1;
+  if ($[0] !== T0 || $[1] !== t0) {
+    t1 = <T0 onPress={t0} title="Randomize" />;
+    $[0] = T0;
+    $[1] = t0;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+```
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reanimated-shared-value-writes.jsx

@@ -0,0 +1,18 @@
+// @enableCustomTypeDefinitionForReanimated
+import {useSharedValue} from 'react-native-reanimated';
+
+/**
+ * https://docs.swmansion.com/react-native-reanimated/docs/2.x/api/hooks/useSharedValue/
+ *
+ * Test that shared values are treated as ref-like, i.e. allowing writes outside
+ * of render
+ */
+function SomeComponent() {
+  const sharedVal = useSharedValue(0);
+  return (
+    <Button
+      onPress={() => (sharedVal.value = Math.random())}
+      title="Randomize"
+    />
+  );
+}

compiler/packages/snap/src/SproutTodoFilter.ts

@@ -434,6 +434,7 @@ const skipFilter = new Set([
   'todo.useContext-mutate-context-in-callback',
   'loop-unused-let',
   'reanimated-no-memo-arg',
+  'reanimated-shared-value-writes',
 
   'userspace-use-memo-cache',
   'transitive-freeze-function-expressions',