[compiler][poc] Improve impurity/ref tracking

We currently model impure functions like `Date.now()` with an `Impure` effect, but the way this works is that the effect is considered an error if it's called at all. But we really want to reflect the fact that the _value_ is impure, and is okay to be used outside of render. For example, `useRef(Date.now())` should be fine.

So this PR changes the Impure effect to describe that an impure value flows `into: Place`. We flow impurity through the data flow graph during InferMutationAliasingRanges, including propagating this to the function's effects. So if a function expression returns a transitively impure value, that's expressed in our inferred signature for that function. Calling it propagates the impurity via our effect system.

We stop this propagation when reaching a ref, allowing `useRef(Date.now())` or `useRef(localFunctionThatReturnsDateNow())`.

This lets us also model accessing a ref as an impure value - we just emit an `Impure` event for PropertyLoad/ComputedLoad off of a ref, and the above tracking kicks in.

A final piece is that we can also use our existing machinery to disallow writing to global values during render to handle writing to refs - except that we need to allow the lazy init pattern. So it probably makes sense to keep the ref validation pass, but scope it back to just handling writes and not reads.

This definitely needs more polish, the error messages would ideally be something like:

```
Error: cannot call impure function in render

Blah blah description of why this is bad...

foo.js:0:0
  <Foo x={x} />
          ^ This value reads from an impure function

foo.js:0:0
  x = Date.now();
      ^^^^^^^^ The value derives from Date.now which is impure
```

Ie show you both the local site (where impurity flows into render) and the ultiamte source of the impure value.

Author: josephsavona

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

@@ -132,6 +132,12 @@ export class CompilerDiagnostic {
     return new CompilerDiagnostic({...options, details: []});
   }
 
+  clone(): CompilerDiagnostic {
+    const cloned = CompilerDiagnostic.create({...this.options});
+    cloned.options.details = [...this.options.details];
+    return cloned;
+  }
+
   get reason(): CompilerDiagnosticOptions['reason'] {
     return this.options.reason;
   }

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts

@@ -96,7 +96,6 @@ import {propagateScopeDependenciesHIR} from '../HIR/PropagateScopeDependenciesHI
 import {outlineJSX} from '../Optimization/OutlineJsx';
 import {optimizePropsMethodCalls} from '../Optimization/OptimizePropsMethodCalls';
 import {transformFire} from '../Transform';
-import {validateNoImpureFunctionsInRender} from '../Validation/ValidateNoImpureFunctionsInRender';
 import {validateStaticComponents} from '../Validation/ValidateStaticComponents';
 import {validateNoFreezingKnownMutableFunctions} from '../Validation/ValidateNoFreezingKnownMutableFunctions';
 import {inferMutationAliasingEffects} from '../Inference/InferMutationAliasingEffects';
@@ -293,10 +292,6 @@ function runWithEnvironment(
       env.logErrors(validateNoJSXInTryStatement(hir));
     }
 
-    if (env.config.validateNoImpureFunctionsInRender) {
-      validateNoImpureFunctionsInRender(hir).unwrap();
-    }
-
     validateNoFreezingKnownMutableFunctions(hir).unwrap();
   }
 

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

@@ -38,7 +38,7 @@ import {
   addObject,
 } from './ObjectShape';
 import {BuiltInType, ObjectType, PolyType} from './Types';
-import {TypeConfig} from './TypeSchema';
+import {AliasingSignatureConfig, TypeConfig} from './TypeSchema';
 import {assertExhaustive} from '../Utils/utils';
 import {isHookName} from './Environment';
 import {CompilerError, SourceLocation} from '..';
@@ -626,6 +626,78 @@ const TYPED_GLOBALS: Array<[string, BuiltInType]> = [
   // TODO: rest of Global objects
 ];
 
+const RenderHookAliasing: (
+  reason: ValueReason,
+) => AliasingSignatureConfig = reason => ({
+  receiver: '@receiver',
+  params: [],
+  rest: '@rest',
+  returns: '@returns',
+  temporaries: [],
+  effects: [
+    // Freeze the arguments
+    {
+      kind: 'Freeze',
+      value: '@rest',
+      reason: ValueReason.HookCaptured,
+    },
+    // Render the arguments
+    {
+      kind: 'Render',
+      place: '@rest',
+    },
+    // Returns a frozen value
+    {
+      kind: 'Create',
+      into: '@returns',
+      value: ValueKind.Frozen,
+      reason,
+    },
+    // May alias any arguments into the return
+    {
+      kind: 'Alias',
+      from: '@rest',
+      into: '@returns',
+    },
+  ],
+});
+
+const EffectHookAliasing: AliasingSignatureConfig = {
+  receiver: '@receiver',
+  params: [],
+  rest: '@rest',
+  returns: '@returns',
+  temporaries: ['@effect'],
+  effects: [
+    // Freezes the function and deps
+    {
+      kind: 'Freeze',
+      value: '@rest',
+      reason: ValueReason.Effect,
+    },
+    // Internally creates an effect object that captures the function and deps
+    {
+      kind: 'Create',
+      into: '@effect',
+      value: ValueKind.Frozen,
+      reason: ValueReason.KnownReturnSignature,
+    },
+    // The effect stores the function and dependencies
+    {
+      kind: 'Capture',
+      from: '@rest',
+      into: '@effect',
+    },
+    // Returns undefined
+    {
+      kind: 'Create',
+      into: '@returns',
+      value: ValueKind.Primitive,
+      reason: ValueReason.KnownReturnSignature,
+    },
+  ],
+};
+
 /*
  * TODO(mofeiZ): We currently only store rest param effects for hooks.
  * now that FeatureFlag `enableTreatHooksAsFunctions` is removed we can
@@ -644,6 +716,7 @@ const REACT_APIS: Array<[string, BuiltInType]> = [
         hookKind: 'useContext',
         returnValueKind: ValueKind.Frozen,
         returnValueReason: ValueReason.Context,
+        aliasing: RenderHookAliasing(ValueReason.Context),
       },
       BuiltInUseContextHookId,
     ),
@@ -658,6 +731,7 @@ const REACT_APIS: Array<[string, BuiltInType]> = [
       hookKind: 'useState',
       returnValueKind: ValueKind.Frozen,
       returnValueReason: ValueReason.State,
+      aliasing: RenderHookAliasing(ValueReason.State),
     }),
   ],
   [
@@ -670,6 +744,7 @@ const REACT_APIS: Array<[string, BuiltInType]> = [
       hookKind: 'useActionState',
       returnValueKind: ValueKind.Frozen,
       returnValueReason: ValueReason.State,
+      aliasing: RenderHookAliasing(ValueReason.HookCaptured),
     }),
   ],
   [
@@ -682,6 +757,7 @@ const REACT_APIS: Array<[string, BuiltInType]> = [
       hookKind: 'useReducer',
       returnValueKind: ValueKind.Frozen,
       returnValueReason: ValueReason.ReducerState,
+      aliasing: RenderHookAliasing(ValueReason.ReducerState),
     }),
   ],
   [
@@ -715,6 +791,7 @@ const REACT_APIS: Array<[string, BuiltInType]> = [
       calleeEffect: Effect.Read,
       hookKind: 'useMemo',
       returnValueKind: ValueKind.Frozen,
+      aliasing: RenderHookAliasing(ValueReason.HookCaptured),
     }),
   ],
   [
@@ -726,6 +803,7 @@ const REACT_APIS: Array<[string, BuiltInType]> = [
       calleeEffect: Effect.Read,
       hookKind: 'useCallback',
       returnValueKind: ValueKind.Frozen,
+      aliasing: RenderHookAliasing(ValueReason.HookCaptured),
     }),
   ],
   [
@@ -739,41 +817,7 @@ const REACT_APIS: Array<[string, BuiltInType]> = [
         calleeEffect: Effect.Read,
         hookKind: 'useEffect',
         returnValueKind: ValueKind.Frozen,
-        aliasing: {
-          receiver: '@receiver',
-          params: [],
-          rest: '@rest',
-          returns: '@returns',
-          temporaries: ['@effect'],
-          effects: [
-            // Freezes the function and deps
-            {
-              kind: 'Freeze',
-              value: '@rest',
-              reason: ValueReason.Effect,
-            },
-            // Internally creates an effect object that captures the function and deps
-            {
-              kind: 'Create',
-              into: '@effect',
-              value: ValueKind.Frozen,
-              reason: ValueReason.KnownReturnSignature,
-            },
-            // The effect stores the function and dependencies
-            {
-              kind: 'Capture',
-              from: '@rest',
-              into: '@effect',
-            },
-            // Returns undefined
-            {
-              kind: 'Create',
-              into: '@returns',
-              value: ValueKind.Primitive,
-              reason: ValueReason.KnownReturnSignature,
-            },
-          ],
-        },
+        aliasing: EffectHookAliasing,
       },
       BuiltInUseEffectHookId,
     ),
@@ -789,6 +833,7 @@ const REACT_APIS: Array<[string, BuiltInType]> = [
         calleeEffect: Effect.Read,
         hookKind: 'useLayoutEffect',
         returnValueKind: ValueKind.Frozen,
+        aliasing: EffectHookAliasing,
       },
       BuiltInUseLayoutEffectHookId,
     ),
@@ -804,6 +849,7 @@ const REACT_APIS: Array<[string, BuiltInType]> = [
         calleeEffect: Effect.Read,
         hookKind: 'useInsertionEffect',
         returnValueKind: ValueKind.Frozen,
+        aliasing: EffectHookAliasing,
       },
       BuiltInUseInsertionEffectHookId,
     ),
@@ -817,6 +863,7 @@ const REACT_APIS: Array<[string, BuiltInType]> = [
       calleeEffect: Effect.Read,
       hookKind: 'useTransition',
       returnValueKind: ValueKind.Frozen,
+      aliasing: RenderHookAliasing(ValueReason.HookCaptured),
     }),
   ],
   [
@@ -829,6 +876,7 @@ const REACT_APIS: Array<[string, BuiltInType]> = [
       hookKind: 'useOptimistic',
       returnValueKind: ValueKind.Frozen,
       returnValueReason: ValueReason.State,
+      aliasing: RenderHookAliasing(ValueReason.HookCaptured),
     }),
   ],
   [
@@ -842,6 +890,7 @@ const REACT_APIS: Array<[string, BuiltInType]> = [
         returnType: {kind: 'Poly'},
         calleeEffect: Effect.Read,
         returnValueKind: ValueKind.Frozen,
+        aliasing: RenderHookAliasing(ValueReason.HookCaptured),
       },
       BuiltInUseOperatorId,
     ),

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

@@ -5,7 +5,11 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {CompilerError} from '../CompilerError';
+import {
+  CompilerDiagnostic,
+  CompilerError,
+  ErrorCategory,
+} from '../CompilerError';
 import {AliasingEffect, AliasingSignature} from '../Inference/AliasingEffects';
 import {assertExhaustive} from '../Utils/utils';
 import {
@@ -190,16 +194,28 @@ function parseAliasingSignatureConfig(
           };
         }
         case 'Impure': {
-          const place = lookup(effect.place);
+          const into = lookup(effect.into);
           return {
             kind: 'Impure',
-            place,
-            error: CompilerError.throwTodo({
-              reason: 'Support impure effect declarations',
-              loc: GeneratedSource,
+            into,
+            error: CompilerDiagnostic.create({
+              category: ErrorCategory.Purity,
+              reason: 'Cannot call impure function during render',
+              description: effect.reason,
+            }).withDetails({
+              kind: 'error',
+              loc,
+              message: 'Cannot call impure function',
             }),
           };
         }
+        case 'Render': {
+          const place = lookup(effect.place);
+          return {
+            kind: 'Render',
+            place,
+          };
+        }
         case 'Apply': {
           const receiver = lookup(effect.receiver);
           const fn = lookup(effect.function);
@@ -1513,6 +1529,11 @@ export const DefaultNonmutatingHook = addHook(
           value: '@rest',
           reason: ValueReason.HookCaptured,
         },
+        // Render the arguments
+        {
+          kind: 'Render',
+          place: '@rest',
+        },
         // Returns a frozen value
         {
           kind: 'Create',

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

@@ -1009,7 +1009,7 @@ export function printAliasingEffect(effect: AliasingEffect): string {
       return `MutateGlobal ${printPlaceForAliasEffect(effect.place)} reason=${JSON.stringify(effect.error.reason)}`;
     }
     case 'Impure': {
-      return `Impure ${printPlaceForAliasEffect(effect.place)} reason=${JSON.stringify(effect.error.reason)}`;
+      return `Impure ${printPlaceForAliasEffect(effect.into)} reason=${effect.error.reason}`;
     }
     case 'Render': {
       return `Render ${printPlaceForAliasEffect(effect.place)}`;

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

@@ -185,11 +185,23 @@ export const ApplyEffectSchema: z.ZodType<ApplyEffectConfig> = z.object({
 
 export type ImpureEffectConfig = {
   kind: 'Impure';
-  place: string;
+  into: string;
+  reason: string;
 };
 
 export const ImpureEffectSchema: z.ZodType<ImpureEffectConfig> = z.object({
   kind: z.literal('Impure'),
+  into: LifetimeIdSchema,
+  reason: z.string(),
+});
+
+export type RenderEffectConfig = {
+  kind: 'Render';
+  place: string;
+};
+
+export const RenderEffectSchema: z.ZodType<RenderEffectConfig> = z.object({
+  kind: z.literal('Render'),
   place: LifetimeIdSchema,
 });
 
@@ -204,7 +216,8 @@ export type AliasingEffectConfig =
   | ImpureEffectConfig
   | MutateEffectConfig
   | MutateTransitiveConditionallyConfig
-  | ApplyEffectConfig;
+  | ApplyEffectConfig
+  | RenderEffectConfig;
 
 export const AliasingEffectSchema: z.ZodType<AliasingEffectConfig> = z.union([
   FreezeEffectSchema,

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

@@ -162,7 +162,7 @@ export type AliasingEffect =
   /**
    * Indicates a side-effect that is not safe during render
    */
-  | {kind: 'Impure'; place: Place; error: CompilerDiagnostic}
+  | {kind: 'Impure'; into: Place; error: CompilerDiagnostic}
   /**
    * Indicates that a given place is accessed during render. Used to distingush
    * hook arguments that are known to be called immediately vs those used for
@@ -222,6 +222,14 @@ export function hashEffect(effect: AliasingEffect): string {
       return [effect.kind, effect.value.identifier.id, effect.reason].join(':');
     }
     case 'Impure':
+      return [
+        effect.kind,
+        effect.into.identifier.id,
+        effect.error.severity,
+        effect.error.reason,
+        effect.error.description,
+        printSourceLocation(effect.error.primaryLocation() ?? GeneratedSource),
+      ].join(':');
     case 'Render': {
       return [effect.kind, effect.place.identifier.id].join(':');
     }