[core] Restrict setRef usage to ref callback (#16539)

or usage where extra care is taken of cleanup
mui · Jul 10, 2019 · f8fa793 · f8fa793
1 parent f3ff314
commit f8fa793
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 28 deletions.
diff --git a/packages/material-ui/src/GridListTile/GridListTile.test.js b/packages/material-ui/src/GridListTile/GridListTile.test.js
@@ -3,7 +3,6 @@ import { assert } from 'chai';
 import { spy, useFakeTimers } from 'sinon';
 import { createMount, getClasses } from '@material-ui/core/test-utils';
 import describeConformance from '../test-utils/describeConformance';
-import { setRef } from '../utils/reactHelpers';
 import GridListTile from './GridListTile';
 
 describe('<GridListTile />', () => {
@@ -50,15 +49,9 @@ describe('<GridListTile />', () => {
 
   function mountMockImage(imgEl) {
     const Image = React.forwardRef((props, ref) => {
-      return (
-        <img
-          alt="test"
-          ref={() => {
-            setRef(ref, imgEl);
-          }}
-          {...props}
-        />
-      );
+      React.useImperativeHandle(ref, () => imgEl, []);
+
+      return <img alt="test" {...props} />;
     });
     Image.muiName = 'Image';
 

diff --git a/packages/material-ui/src/useScrollTrigger/useScrollTrigger.test.js b/packages/material-ui/src/useScrollTrigger/useScrollTrigger.test.js
@@ -68,7 +68,7 @@ describe('useScrollTrigger', () => {
       assert.strictEqual(values.callCount, 1);
     });
 
-    it('should be false by default when using setRef', () => {
+    it('should be false by default when using ref', () => {
       const TestDefaultWithRef = () => {
         const [container, setContainer] = React.useState();
         const trigger = useScrollTrigger({

diff --git a/packages/material-ui/src/utils/reactHelpers.d.ts b/packages/material-ui/src/utils/reactHelpers.d.ts
@@ -14,6 +14,10 @@ export function isMuiElement(element: any, muiNames: string[]): element is Named
 /**
  * passes {value} to {ref}
  *
+ * WARNING: Be sure to only call this inside a callback that is passed as a ref.
+ * Otherwise make sure to cleanup previous {ref} if it changes. See
+ * https://github.com/mui-org/material-ui/issues/13539
+ *
  * useful if you want to expose the ref of an inner component to the public api
  * while still using it inside the component
  *
@@ -23,3 +27,5 @@ export function setRef<T>(
   ref: React.RefObject<T> | ((instance: T | null) => void) | null | undefined,
   value: T | null,
 ): void;
+
+export function useForkRef<T>(refA: React.Ref<T>, refB: React.Ref<T>): React.Ref<T>;
diff --git a/packages/material-ui/src/utils/reactHelpers.js b/packages/material-ui/src/utils/reactHelpers.js
@@ -4,6 +4,7 @@ export function isMuiElement(element, muiNames) {
   return React.isValidElement(element) && muiNames.indexOf(element.type.muiName) !== -1;
 }
 
+// TODO: Make it private only in v5
 export function setRef(ref, value) {
   if (typeof ref === 'function') {
     ref(value);

diff --git a/packages/material-ui/test/integration/Menu.test.js b/packages/material-ui/test/integration/Menu.test.js
@@ -6,7 +6,7 @@ import Popover from '@material-ui/core/Popover';
 import SimpleMenu from './fixtures/menus/SimpleMenu';
 import Menu from '@material-ui/core/Menu';
 import MenuItem from '@material-ui/core/MenuItem';
-import { setRef } from '../../src/utils/reactHelpers';
+import { useForkRef } from '@material-ui/core/utils';
 import { stub } from 'sinon';
 import PropTypes from 'prop-types';
 import { createMuiTheme } from '@material-ui/core/styles';
@@ -154,27 +154,28 @@ describe('<Menu> integration', () => {
     let menuListFocusTracker = false;
     const tabIndexTracker = [false, false, false];
     const TrackingMenuItem = React.forwardRef(({ itemIndex, ...other }, ref) => {
+      const handleRef = useForkRef(ref, instance => {
+        if (instance && !instance.stubbed) {
+          if (instance.tabIndex === 0) {
+            tabIndexTracker[itemIndex] = true;
+          } else if (instance.tabIndex > 0) {
+            tabIndexTracker[itemIndex] = instance.tabIndex;
+          }
+          const offsetTop = instance.offsetTop;
+          stub(instance, 'offsetTop').get(() => {
+            contentAnchorTracker[itemIndex] = true;
+            return offsetTop;
+          });
+          instance.stubbed = true;
+        }
+      });
+
       return (
         <MenuItem
           onFocus={() => {
             focusTracker[itemIndex] = true;
           }}
-          ref={instance => {
-            setRef(ref, instance);
-            if (instance && !instance.stubbed) {
-              if (instance.tabIndex === 0) {
-                tabIndexTracker[itemIndex] = true;
-              } else if (instance.tabIndex > 0) {
-                tabIndexTracker[itemIndex] = instance.tabIndex;
-              }
-              const offsetTop = instance.offsetTop;
-              stub(instance, 'offsetTop').get(() => {
-                contentAnchorTracker[itemIndex] = true;
-                return offsetTop;
-              });
-              instance.stubbed = true;
-            }
-          }}
+          ref={handleRef}
           {...other}
         />
       );