A few Suspense related refactors

packages/react-reconciler/src/ReactFiberCompleteWork.new.js

@@ -31,7 +31,6 @@ import type {OffscreenState} from './ReactFiberOffscreenComponent';
 import type {TracingMarkerInstance} from './ReactFiberTracingMarkerComponent.new';
 import type {Cache} from './ReactFiberCacheComponent.new';
 import {
-  enableSuspenseAvoidThisFallback,
   enableLegacyHidden,
   enableHostSingletons,
   enableSuspenseCallback,
@@ -127,11 +126,9 @@ import {
   setShallowSuspenseListContext,
   ForceSuspenseFallback,
   setDefaultShallowSuspenseListContext,
+  isBadSuspenseFallback,
 } from './ReactFiberSuspenseContext.new';
-import {
-  popHiddenContext,
-  isCurrentTreeHidden,
-} from './ReactFiberHiddenContext.new';
+import {popHiddenContext} from './ReactFiberHiddenContext.new';
 import {findFirstSuspended} from './ReactFiberSuspenseComponent.new';
 import {
   isContextProvider as isLegacyContextProvider,
@@ -1272,20 +1269,7 @@ function completeWork(
             // If this render already had a ping or lower pri updates,
             // and this is the first time we know we're going to suspend we
             // should be able to immediately restart from within throwException.
-
-            // Check if this is a "bad" fallback state or a good one. A bad
-            // fallback state is one that we only show as a last resort; if this
-            // is a transition, we'll block it from displaying, and wait for
-            // more data to arrive.
-            const isBadFallback =
-              // It's bad to switch to a fallback if content is already visible
-              (current !== null && !prevDidTimeout && !isCurrentTreeHidden()) ||
-              // Experimental: Some fallbacks are always bad
-              (enableSuspenseAvoidThisFallback &&
-                workInProgress.memoizedProps.unstable_avoidThisFallback ===
-                  true);
-
-            if (isBadFallback) {
+            if (isBadSuspenseFallback(current, newProps)) {
               renderDidSuspendDelayIfPossible();
             } else {
               renderDidSuspend();

packages/react-reconciler/src/ReactFiberCompleteWork.old.js

@@ -31,7 +31,6 @@ import type {OffscreenState} from './ReactFiberOffscreenComponent';
 import type {TracingMarkerInstance} from './ReactFiberTracingMarkerComponent.old';
 import type {Cache} from './ReactFiberCacheComponent.old';
 import {
-  enableSuspenseAvoidThisFallback,
   enableLegacyHidden,
   enableHostSingletons,
   enableSuspenseCallback,
@@ -127,11 +126,9 @@ import {
   setShallowSuspenseListContext,
   ForceSuspenseFallback,
   setDefaultShallowSuspenseListContext,
+  isBadSuspenseFallback,
 } from './ReactFiberSuspenseContext.old';
-import {
-  popHiddenContext,
-  isCurrentTreeHidden,
-} from './ReactFiberHiddenContext.old';
+import {popHiddenContext} from './ReactFiberHiddenContext.old';
 import {findFirstSuspended} from './ReactFiberSuspenseComponent.old';
 import {
   isContextProvider as isLegacyContextProvider,
@@ -1272,20 +1269,7 @@ function completeWork(
             // If this render already had a ping or lower pri updates,
             // and this is the first time we know we're going to suspend we
             // should be able to immediately restart from within throwException.
-
-            // Check if this is a "bad" fallback state or a good one. A bad
-            // fallback state is one that we only show as a last resort; if this
-            // is a transition, we'll block it from displaying, and wait for
-            // more data to arrive.
-            const isBadFallback =
-              // It's bad to switch to a fallback if content is already visible
-              (current !== null && !prevDidTimeout && !isCurrentTreeHidden()) ||
-              // Experimental: Some fallbacks are always bad
-              (enableSuspenseAvoidThisFallback &&
-                workInProgress.memoizedProps.unstable_avoidThisFallback ===
-                  true);
-
-            if (isBadFallback) {
+            if (isBadSuspenseFallback(current, newProps)) {
               renderDidSuspendDelayIfPossible();
             } else {
               renderDidSuspend();

packages/react-reconciler/src/ReactFiberSuspenseComponent.new.js

@@ -27,6 +27,7 @@ export type SuspenseProps = {
   // TODO: Add "unstable_" prefix?
   suspenseCallback?: (Set<Wakeable> | null) => mixed,
 
+  unstable_avoidThisFallback?: boolean,
   unstable_expectedLoadTime?: number,
   unstable_name?: string,
 };

packages/react-reconciler/src/ReactFiberSuspenseComponent.old.js

@@ -27,6 +27,7 @@ export type SuspenseProps = {
   // TODO: Add "unstable_" prefix?
   suspenseCallback?: (Set<Wakeable> | null) => mixed,
 
+  unstable_avoidThisFallback?: boolean,
   unstable_expectedLoadTime?: number,
   unstable_name?: string,
 };

packages/react-reconciler/src/ReactFiberSuspenseContext.new.js

@@ -9,7 +9,10 @@
 
 import type {Fiber} from './ReactInternalTypes';
 import type {StackCursor} from './ReactFiberStack.new';
-import type {SuspenseState} from './ReactFiberSuspenseComponent.new';
+import type {
+  SuspenseState,
+  SuspenseProps,
+} from './ReactFiberSuspenseComponent.new';
 
 import {enableSuspenseAvoidThisFallback} from 'shared/ReactFeatureFlags';
 import {createCursor, push, pop} from './ReactFiberStack.new';
@@ -55,6 +58,33 @@ function shouldAvoidedBoundaryCapture(
   return false;
 }
 
+export function isBadSuspenseFallback(
+  current: Fiber | null,
+  nextProps: SuspenseProps,
+): boolean {
+  // Check if this is a "bad" fallback state or a good one. A bad fallback state
+  // is one that we only show as a last resort; if this is a transition, we'll
+  // block it from displaying, and wait for more data to arrive.
+  if (current !== null) {
+    const prevState: SuspenseState = current.memoizedState;
+    const isShowingFallback = prevState !== null;
+    if (!isShowingFallback && !isCurrentTreeHidden()) {
+      // It's bad to switch to a fallback if content is already visible
+      return true;
+    }
+  }
+
+  if (
+    enableSuspenseAvoidThisFallback &&
+    nextProps.unstable_avoidThisFallback === true
+  ) {
+    // Experimental: Some fallbacks are always bad
+    return true;
+  }
+
+  return false;
+}
+
 export function pushPrimaryTreeSuspenseHandler(handler: Fiber): void {
   const props = handler.pendingProps;
   const handlerOnStack = suspenseHandlerStackCursor.current;

packages/react-reconciler/src/ReactFiberSuspenseContext.old.js

@@ -9,7 +9,10 @@
 
 import type {Fiber} from './ReactInternalTypes';
 import type {StackCursor} from './ReactFiberStack.old';
-import type {SuspenseState} from './ReactFiberSuspenseComponent.old';
+import type {
+  SuspenseState,
+  SuspenseProps,
+} from './ReactFiberSuspenseComponent.old';
 
 import {enableSuspenseAvoidThisFallback} from 'shared/ReactFeatureFlags';
 import {createCursor, push, pop} from './ReactFiberStack.old';
@@ -55,6 +58,33 @@ function shouldAvoidedBoundaryCapture(
   return false;
 }
 
+export function isBadSuspenseFallback(
+  current: Fiber | null,
+  nextProps: SuspenseProps,
+): boolean {
+  // Check if this is a "bad" fallback state or a good one. A bad fallback state
+  // is one that we only show as a last resort; if this is a transition, we'll
+  // block it from displaying, and wait for more data to arrive.
+  if (current !== null) {
+    const prevState: SuspenseState = current.memoizedState;
+    const isShowingFallback = prevState !== null;
+    if (!isShowingFallback && !isCurrentTreeHidden()) {
+      // It's bad to switch to a fallback if content is already visible
+      return true;
+    }
+  }
+
+  if (
+    enableSuspenseAvoidThisFallback &&
+    nextProps.unstable_avoidThisFallback === true
+  ) {
+    // Experimental: Some fallbacks are always bad
+    return true;
+  }
+
+  return false;
+}
+
 export function pushPrimaryTreeSuspenseHandler(handler: Fiber): void {
   const props = handler.pendingProps;
   const handlerOnStack = suspenseHandlerStackCursor.current;