Fix ReactFiberReconciler annotation to include PI (#8751)

* Fix ReactFiberReconciler annotation to include `PI`

#8628 (comment)

* Make getPublicInstance type safe

* attempting to trigger the issue @sebmarkbage described

#8751 (comment)

* Unify renderer-specific getPublicInstance with getRootInstance

* Switch on fiber type HostComponent for getPublicRootInstance

* Fix test that was too dynamic (and failing)

* Use PI in reconciler public API

* Prettier

Author: iamdustan

scripts/fiber/tests-passing.txt

@@ -610,6 +610,7 @@ src/renderers/__tests__/refs-test.js
 * coerces numbers to strings
 * attaches, detaches from fiber component with stack layer
 * attaches, detaches from stack component with fiber layer
+* attaches and detaches root refs
 
 src/renderers/art/__tests__/ReactART-test.js
 * should have the correct lifecycle state
@@ -1972,6 +1973,7 @@ src/renderers/testing/__tests__/ReactTestRenderer-test.js
 * can update text nodes
 * toTree() renders simple components returning host components
 * toTree() handles null rendering components
+* root instance and createNodeMock ref return the same value
 * toTree() renders complicated trees of composites and hosts
 * can update text nodes when rendered as root
 * can render and update root fragments

src/renderers/__tests__/refs-test.js

@@ -409,3 +409,87 @@ describe('string refs between fiber and stack', () => {
     }
   });
 });
+
+describe('root level refs', () => {
+  beforeEach(() => {
+    var ReactFeatureFlags = require('ReactFeatureFlags');
+    ReactFeatureFlags.disableNewFiberFeatures = false;
+  });
+
+  it('attaches and detaches root refs', () => {
+    var ReactDOM = require('react-dom');
+    var inst = null;
+
+    // host node
+    var ref = jest.fn(value => (inst = value));
+    var container = document.createElement('div');
+    var result = ReactDOM.render(<div ref={ref} />, container);
+    expect(ref).toHaveBeenCalledTimes(1);
+    expect(ref.mock.calls[0][0]).toBeInstanceOf(HTMLDivElement);
+    expect(result).toBe(ref.mock.calls[0][0]);
+    ReactDOM.unmountComponentAtNode(container);
+    expect(ref).toHaveBeenCalledTimes(2);
+    expect(ref.mock.calls[1][0]).toBe(null);
+
+    // composite
+    class Comp extends React.Component {
+      method() {
+        return true;
+      }
+      render() {
+        return <div>Comp</div>;
+      }
+    }
+
+    inst = null;
+    ref = jest.fn(value => (inst = value));
+    result = ReactDOM.render(<Comp ref={ref} />, container);
+
+    expect(ref).toHaveBeenCalledTimes(1);
+    expect(inst).toBeInstanceOf(Comp);
+    expect(result).toBe(inst);
+
+    // ensure we have the correct instance
+    expect(result.method()).toBe(true);
+    expect(inst.method()).toBe(true);
+
+    ReactDOM.unmountComponentAtNode(container);
+    expect(ref).toHaveBeenCalledTimes(2);
+    expect(ref.mock.calls[1][0]).toBe(null);
+
+    if (ReactDOMFeatureFlags.useFiber) {
+      // fragment
+      inst = null;
+      ref = jest.fn(value => (inst = value));
+      var divInst = null;
+      var ref2 = jest.fn(value => (divInst = value));
+      result = ReactDOM.render(
+        [<Comp ref={ref} key="a" />, 5, <div ref={ref2} key="b">Hello</div>],
+        container,
+      );
+
+      // first call should be `Comp`
+      expect(ref).toHaveBeenCalledTimes(1);
+      expect(ref.mock.calls[0][0]).toBeInstanceOf(Comp);
+      expect(result).toBe(ref.mock.calls[0][0]);
+
+      expect(ref2).toHaveBeenCalledTimes(1);
+      expect(divInst).toBeInstanceOf(HTMLDivElement);
+      expect(result).not.toBe(divInst);
+
+      ReactDOM.unmountComponentAtNode(container);
+      expect(ref).toHaveBeenCalledTimes(2);
+      expect(ref.mock.calls[1][0]).toBe(null);
+      expect(ref2).toHaveBeenCalledTimes(2);
+      expect(ref2.mock.calls[1][0]).toBe(null);
+
+      // null
+      result = ReactDOM.render(null, container);
+      expect(result).toBe(null);
+
+      // primitives
+      result = ReactDOM.render(5, container);
+      expect(result).toBeInstanceOf(Text);
+    }
+  });
+});

src/renderers/shared/fiber/ReactFiberCommitWork.js

@@ -518,8 +518,14 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
   function commitAttachRef(finishedWork: Fiber) {
     const ref = finishedWork.ref;
     if (ref !== null) {
-      const instance = getPublicInstance(finishedWork.stateNode);
-      ref(instance);
+      const instance = finishedWork.stateNode;
+      switch (finishedWork.tag) {
+        case HostComponent:
+          ref(getPublicInstance(instance));
+          break;
+        default:
+          ref(instance);
+      }
     }
   }
 

src/renderers/shared/fiber/ReactFiberReconciler.js

@@ -28,6 +28,7 @@ var {
 } = require('ReactFiberContext');
 var {createFiberRoot} = require('ReactFiberRoot');
 var ReactFiberScheduler = require('ReactFiberScheduler');
+var {HostComponent} = require('ReactTypeOfWork');
 
 if (__DEV__) {
   var warning = require('fbjs/lib/warning');
@@ -167,6 +168,8 @@ getContextForSubtree._injectFiber(function(fiber: Fiber) {
 module.exports = function<T, P, I, TI, PI, C, CX, PL>(
   config: HostConfig<T, P, I, TI, PI, C, CX, PL>,
 ): Reconciler<C, I, TI> {
+  var {getPublicInstance} = config;
+
   var {
     scheduleUpdate,
     getPriorityContext,
@@ -269,15 +272,20 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
 
     getPublicRootInstance(
       container: OpaqueRoot,
-    ): ReactComponent<any, any, any> | I | TI | null {
+    ): ReactComponent<any, any, any> | PI | null {
       const containerFiber = container.current;
       if (!containerFiber.child) {
         return null;
       }
-      return containerFiber.child.stateNode;
+      switch (containerFiber.child.tag) {
+        case HostComponent:
+          return getPublicInstance(containerFiber.child.stateNode);
+        default:
+          return containerFiber.child.stateNode;
+      }
     },
 
-    findHostInstance(fiber: Fiber): I | TI | null {
+    findHostInstance(fiber: Fiber): PI | null {
       const hostFiber = findCurrentHostFiber(fiber);
       if (hostFiber === null) {
         return null;

src/renderers/testing/ReactTestRendererFiber.js

@@ -219,7 +219,7 @@ var TestRenderer = ReactFiberReconciler({
 
   useSyncScheduling: true,
 
-  getPublicInstance(inst) {
+  getPublicInstance(inst: Instance | TextInstance): * {
     switch (inst.tag) {
       case 'INSTANCE':
         const createNodeMock = inst.rootContainerInstance.createNodeMock;

src/renderers/testing/__tests__/ReactTestRenderer-test.js

@@ -568,6 +568,17 @@ describe('ReactTestRenderer', () => {
     });
   });
 
+  it('root instance and createNodeMock ref return the same value', () => {
+    var createNodeMock = ref => ({node: ref});
+    var refInst = null;
+    var renderer = ReactTestRenderer.create(
+      <div ref={ref => (refInst = ref)} />,
+      {createNodeMock},
+    );
+    var root = renderer.getInstance();
+    expect(root).toEqual(refInst);
+  });
+
   it('toTree() renders complicated trees of composites and hosts', () => {
     // SFC returning host. no children props.
     var Qoo = () => <span className="Qoo">Hello World!</span>;