.

Author: paulshen

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

@@ -650,6 +650,14 @@ export const EnvironmentConfigSchema = z.object({
    *   useMemo(() => { ... }, [...]);
    */
   validateNoVoidUseMemo: z.boolean().default(false),
+
+  /**
+   * A list of function identifier names that should never be skipped by
+   * memoization/flattening optimizations. Calls to any identifier with a name
+   * in this list behave like the `use` operator in terms of never being
+   * skipped, and may be called conditionally.
+   */
+  neverSkipFunctionName: z.array(z.string()).default([]),
 });
 
 export type EnvironmentConfig = z.infer<typeof EnvironmentConfigSchema>;

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

@@ -1915,6 +1915,34 @@ export function isUseOperator(id: Identifier): boolean {
   );
 }
 
+/**
+ * Treat any call with a neverSkip identifier name similar to the `use` operator:
+ * - It may be called conditionally.
+ * - It should never be skipped by memoization/flattening logic.
+ */
+export function isNeverSkipIdentifier(
+  env: Environment,
+  id: Identifier,
+): boolean {
+  const list = env.config.neverSkipFunctionName;
+  if (list == null || list.length === 0) {
+    return false;
+  }
+  if (id.name != null && list.includes(id.name.value)) {
+    return true;
+  }
+  const loc = id.loc as any;
+  if (
+    loc &&
+    typeof loc !== 'symbol' &&
+    typeof loc.identifierName === 'string' &&
+    list.includes(loc.identifierName)
+  ) {
+    return true;
+  }
+  return false;
+}
+
 export function getHookKindForType(
   env: Environment,
   type: Type,

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

@@ -21,6 +21,7 @@ import {
   isStableType,
   isStableTypeContainer,
   isUseOperator,
+  isNeverSkipIdentifier,
 } from '../HIR';
 import {PostDominator} from '../HIR/Dominator';
 import {
@@ -302,13 +303,15 @@ export function inferReactivePlaces(fn: HIRFunction): void {
         if (
           value.kind === 'CallExpression' &&
           (getHookKind(fn.env, value.callee.identifier) != null ||
-            isUseOperator(value.callee.identifier))
+            isUseOperator(value.callee.identifier) ||
+            isNeverSkipIdentifier(fn.env, value.callee.identifier))
         ) {
           hasReactiveInput = true;
         } else if (
           value.kind === 'MethodCall' &&
           (getHookKind(fn.env, value.property.identifier) != null ||
-            isUseOperator(value.property.identifier))
+            isUseOperator(value.property.identifier) ||
+            isNeverSkipIdentifier(fn.env, value.property.identifier))
         ) {
           hasReactiveInput = true;
         }

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenScopesWithHooksOrUseHIR.ts

@@ -13,6 +13,7 @@ import {
   PrunedScopeTerminal,
   getHookKind,
   isUseOperator,
+  isNeverSkipIdentifier,
 } from '../HIR';
 import {retainWhere} from '../Utils/utils';
 
@@ -53,7 +54,8 @@ export function flattenScopesWithHooksOrUseHIR(fn: HIRFunction): void {
             value.kind === 'MethodCall' ? value.property : value.callee;
           if (
             getHookKind(fn.env, callee.identifier) != null ||
-            isUseOperator(callee.identifier)
+            isUseOperator(callee.identifier) ||
+            isNeverSkipIdentifier(fn.env, callee.identifier)
           ) {
             prune.push(...activeScopes.map(entry => entry.block));
             activeScopes.length = 0;