Move all markRef calls into begin phase

Certain fiber types may have a ref attached to them. The main ones are
HostComponent and ClassComponent. During the render phase, we check if
a ref was passed to it, and if so, we schedule a Ref effect: `markRef`.

Currently, we're not consistent about whether we call `markRef` in the
begin phase or the complete phase. For some fiber types, I found that
`markRef` was called in both phases, causing redundant work.

After some investigation, I don't believe it's necessary to call
`markRef` in both the begin phase and the complete phase, as long as you
don't bail out before calling `markRef`.

I though that maybe it had to do with the
`attemptEarlyBailoutIfNoScheduledUpdates` branch, which is a fast path
that skips the regular begin phase if no new props, state, or context
were passed. But if the props haven't changed (referentially —
the `memo` and `shouldComponentUpdate` checks happen later), then it
follows that the ref couldn't have changed either. This is true even in
the old `createElement` runtime where `ref` is stored on the element
instead of as a prop, because there's no way to pass a new ref to an
element without also passing new props. You might argue this is a
leaky assumption, but since we're shifting ref to be just a regular prop
anyway, I think it's the correct way to think about it going forward.

I think the pattern of calling `markRef` in the complete phase may have
been left over from an earlier iteration of the implementation before
the bailout logic was structured like it is today.

So, I removed all the `markRef` calls from the complete phase. In the
case of ScopeComponent, which had no corresponding call in the begin
phase, I added one.

We already had a test that asserted that a ref is reattached even if
the component bails out, but I added some new ones to be extra safe.

The reason I'm changing this this is because I'm working on a different
change to move the ref handling logic in `coerceRef` to happen in render
phase of the  component that accepts the ref, instead of during the
parent's reconciliation.

Author: acdlite

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -1007,6 +1007,8 @@ function updateProfiler(
 }
 
 function markRef(current: Fiber | null, workInProgress: Fiber) {
+  // TODO: This is also where we should check the type of the ref and error if
+  // an invalid one is passed, instead of during child reconcilation.
   const ref = workInProgress.ref;
   if (
     (current === null && ref !== null) ||
@@ -3531,7 +3533,7 @@ function updateScopeComponent(
 ) {
   const nextProps = workInProgress.pendingProps;
   const nextChildren = nextProps.children;
-
+  markRef(current, workInProgress);
   reconcileChildren(current, workInProgress, nextChildren, renderLanes);
   return workInProgress.child;
 }

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -75,8 +75,6 @@ import {
 } from './ReactWorkTags';
 import {NoMode, ConcurrentMode, ProfileMode} from './ReactTypeOfMode';
 import {
-  Ref,
-  RefStatic,
   Placement,
   Update,
   Visibility,
@@ -186,10 +184,6 @@ function markUpdate(workInProgress: Fiber) {
   workInProgress.flags |= Update;
 }
 
-function markRef(workInProgress: Fiber) {
-  workInProgress.flags |= Ref | RefStatic;
-}
-
 /**
  * In persistent mode, return whether this update needs to clone the subtree.
  */
@@ -1083,9 +1077,6 @@ function completeWork(
           // @TODO refactor this block to create the instance here in complete
           // phase if we are not hydrating.
           markUpdate(workInProgress);
-          if (workInProgress.ref !== null) {
-            markRef(workInProgress);
-          }
           if (nextResource !== null) {
             // This is a Hoistable Resource
 
@@ -1120,9 +1111,6 @@ function completeWork(
             // and require an update
             markUpdate(workInProgress);
           }
-          if (current.ref !== workInProgress.ref) {
-            markRef(workInProgress);
-          }
           if (nextResource !== null) {
             // This is a Hoistable Resource
             // This must come at the very end of the complete phase.
@@ -1194,10 +1182,6 @@ function completeWork(
               renderLanes,
             );
           }
-
-          if (current.ref !== workInProgress.ref) {
-            markRef(workInProgress);
-          }
         } else {
           if (!newProps) {
             if (workInProgress.stateNode === null) {
@@ -1232,11 +1216,6 @@ function completeWork(
             workInProgress.stateNode = instance;
             markUpdate(workInProgress);
           }
-
-          if (workInProgress.ref !== null) {
-            // If there is a ref on a host node we need to schedule a callback
-            markRef(workInProgress);
-          }
         }
         bubbleProperties(workInProgress);
         return null;
@@ -1254,10 +1233,6 @@ function completeWork(
           newProps,
           renderLanes,
         );
-
-        if (current.ref !== workInProgress.ref) {
-          markRef(workInProgress);
-        }
       } else {
         if (!newProps) {
           if (workInProgress.stateNode === null) {
@@ -1310,11 +1285,6 @@ function completeWork(
             markUpdate(workInProgress);
           }
         }
-
-        if (workInProgress.ref !== null) {
-          // If there is a ref on a host node we need to schedule a callback
-          markRef(workInProgress);
-        }
       }
       bubbleProperties(workInProgress);
 
@@ -1739,16 +1709,16 @@ function completeWork(
           workInProgress.stateNode = scopeInstance;
           prepareScopeUpdate(scopeInstance, workInProgress);
           if (workInProgress.ref !== null) {
-            markRef(workInProgress);
+            // Scope components always do work in the commit phase if there's a
+            // ref attached.
             markUpdate(workInProgress);
           }
         } else {
           if (workInProgress.ref !== null) {
+            // Scope components always do work in the commit phase if there's a
+            // ref attached.
             markUpdate(workInProgress);
           }
-          if (current.ref !== workInProgress.ref) {
-            markRef(workInProgress);
-          }
         }
         bubbleProperties(workInProgress);
         return null;

packages/react-reconciler/src/__tests__/ReactFiberRefs-test.js

@@ -0,0 +1,83 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @emails react-core
+ */
+
+'use strict';
+
+let React;
+let Scheduler;
+let ReactNoop;
+let act;
+let assertLog;
+
+describe('ReactFiberRefs', () => {
+  beforeEach(() => {
+    jest.resetModules();
+    React = require('react');
+    Scheduler = require('scheduler');
+    ReactNoop = require('react-noop-renderer');
+    act = require('internal-test-utils').act;
+    assertLog = require('internal-test-utils').assertLog;
+  });
+
+  test('ref is attached even if there are no other updates (class)', async () => {
+    let component;
+    class Component extends React.PureComponent {
+      render() {
+        Scheduler.log('Render');
+        component = this;
+        return 'Hi';
+      }
+    }
+
+    const ref1 = React.createRef();
+    const ref2 = React.createRef();
+    const root = ReactNoop.createRoot();
+
+    // Mount with ref1 attached
+    await act(() => root.render(<Component ref={ref1} />));
+    assertLog(['Render']);
+    expect(root).toMatchRenderedOutput('Hi');
+    expect(ref1.current).toBe(component);
+    // ref2 has no value
+    expect(ref2.current).toBe(null);
+
+    // Switch to ref2, but don't update anything else.
+    await act(() => root.render(<Component ref={ref2} />));
+    // The component did not re-render because no props changed.
+    assertLog([]);
+    expect(root).toMatchRenderedOutput('Hi');
+    // But the refs still should have been swapped.
+    expect(ref1.current).toBe(null);
+    expect(ref2.current).toBe(component);
+  });
+
+  test('ref is attached even if there are no other updates (host component)', async () => {
+    // This is kind of ailly test because host components never bail out if they
+    // receive a new element, and there's no way to update a ref without also
+    // updating the props, but adding it here anyway for symmetry with the
+    // class case above.
+    const ref1 = React.createRef();
+    const ref2 = React.createRef();
+    const root = ReactNoop.createRoot();
+
+    // Mount with ref1 attached
+    await act(() => root.render(<div ref={ref1}>Hi</div>));
+    expect(root).toMatchRenderedOutput(<div>Hi</div>);
+    expect(ref1.current).not.toBe(null);
+    // ref2 has no value
+    expect(ref2.current).toBe(null);
+
+    // Switch to ref2, but don't update anything else.
+    await act(() => root.render(<div ref={ref2}>Hi</div>));
+    expect(root).toMatchRenderedOutput(<div>Hi</div>);
+    // But the refs still should have been swapped.
+    expect(ref1.current).toBe(null);
+    expect(ref2.current).not.toBe(null);
+  });
+});