Fix clone prop merging of ref and other props

atomiks · Apr 12, 2024 · 613e444 · 613e444
1 parent cf0f789
commit 613e444
Show file tree

Hide file tree

Showing 15 changed files with 98 additions and 15 deletions.
diff --git a/packages/mui-base/src/Checkbox/CheckboxIndicator.tsx b/packages/mui-base/src/Checkbox/CheckboxIndicator.tsx
@@ -5,6 +5,7 @@ import { CheckboxContext } from './CheckboxContext';
 import { resolveClassName } from '../utils/resolveClassName';
 import { useCheckboxStyleHooks } from './utils';
 import { evaluateRenderProp } from '../utils/evaluateRenderProp';
+import { useExtendedForkRef } from '../utils/useExtendedForkRef';
 
 function defaultRender(props: React.ComponentPropsWithRef<'span'>) {
   return <span {...props} />;
@@ -34,14 +35,15 @@ const CheckboxIndicator = React.forwardRef(function CheckboxIndicator(
   }
 
   const styleHooks = useCheckboxStyleHooks(ownerState);
+  const mergedRef = useExtendedForkRef(render, forwardedRef);
 
   if (!keepMounted && !ownerState.checked && !ownerState.indeterminate) {
     return null;
   }
 
   const elementProps = {
     className: resolveClassName(className, ownerState),
-    ref: forwardedRef,
+    ref: mergedRef,
     ...styleHooks,
     ...otherProps,
   };

diff --git a/packages/mui-base/src/Checkbox/CheckboxRoot.tsx b/packages/mui-base/src/Checkbox/CheckboxRoot.tsx
@@ -6,6 +6,7 @@ import { CheckboxContext } from './CheckboxContext';
 import { useCheckbox } from '../useCheckbox';
 import { useCheckboxStyleHooks } from './utils';
 import { evaluateRenderProp } from '../utils/evaluateRenderProp';
+import { useExtendedForkRef } from '../utils/useExtendedForkRef';
 
 function defaultRender(props: React.ComponentPropsWithRef<'button'>) {
   return <button type="button" {...props} />;
@@ -55,10 +56,11 @@ const CheckboxRoot = React.forwardRef(function CheckboxRoot(
   );
 
   const styleHooks = useCheckboxStyleHooks(ownerState);
+  const mergedRef = useExtendedForkRef(render, forwardedRef);
 
   const buttonProps = getButtonProps({
     className: resolveClassName(className, ownerState),
-    ref: forwardedRef,
+    ref: mergedRef,
     ...styleHooks,
     ...otherProps,
   });

diff --git a/packages/mui-base/src/NumberField/NumberFieldDecrement.tsx b/packages/mui-base/src/NumberField/NumberFieldDecrement.tsx
@@ -4,6 +4,7 @@ import type { NumberFieldDecrementProps } from './NumberField.types';
 import { useNumberFieldContext } from './NumberFieldContext';
 import { resolveClassName } from '../utils/resolveClassName';
 import { evaluateRenderProp } from '../utils/evaluateRenderProp';
+import { useExtendedForkRef } from '../utils/useExtendedForkRef';
 
 function defaultRender(props: React.ComponentPropsWithRef<'button'>) {
   return <button type="button" {...props} />;
@@ -29,8 +30,10 @@ const NumberFieldDecrement = React.forwardRef(function NumberFieldDecrement(
 
   const { getDecrementButtonProps, ownerState } = useNumberFieldContext('Decrement');
 
+  const mergedRef = useExtendedForkRef(render, forwardedRef);
+
   const buttonProps = getDecrementButtonProps({
-    ref: forwardedRef,
+    ref: mergedRef,
     className: resolveClassName(className, ownerState),
     ...otherProps,
   });

diff --git a/packages/mui-base/src/NumberField/NumberFieldGroup.tsx b/packages/mui-base/src/NumberField/NumberFieldGroup.tsx
@@ -4,6 +4,7 @@ import { useNumberFieldContext } from './NumberFieldContext';
 import type { NumberFieldGroupProps } from './NumberField.types';
 import { resolveClassName } from '../utils/resolveClassName';
 import { evaluateRenderProp } from '../utils/evaluateRenderProp';
+import { useExtendedForkRef } from '../utils/useExtendedForkRef';
 
 function defaultRender(props: React.ComponentPropsWithRef<'div'>) {
   return <div {...props} />;
@@ -29,8 +30,10 @@ const NumberFieldGroup = React.forwardRef(function NumberFieldGroup(
 
   const { getGroupProps, ownerState } = useNumberFieldContext('Group');
 
+  const mergedRef = useExtendedForkRef(render, forwardedRef);
+
   const groupProps = getGroupProps({
-    ref: forwardedRef,
+    ref: mergedRef,
     className: resolveClassName(className, ownerState),
     ...otherProps,
   });

diff --git a/packages/mui-base/src/NumberField/NumberFieldIncrement.tsx b/packages/mui-base/src/NumberField/NumberFieldIncrement.tsx
@@ -4,6 +4,7 @@ import type { NumberFieldIncrementProps } from './NumberField.types';
 import { useNumberFieldContext } from './NumberFieldContext';
 import { resolveClassName } from '../utils/resolveClassName';
 import { evaluateRenderProp } from '../utils/evaluateRenderProp';
+import { useExtendedForkRef } from '../utils/useExtendedForkRef';
 
 function defaultRender(props: React.ComponentPropsWithRef<'button'>) {
   return <button type="button" {...props} />;
@@ -29,8 +30,10 @@ const NumberFieldIncrement = React.forwardRef(function NumberFieldIncrement(
 
   const { getIncrementButtonProps, ownerState } = useNumberFieldContext('Increment');
 
+  const mergedRef = useExtendedForkRef(render, forwardedRef);
+
   const buttonProps = getIncrementButtonProps({
-    ref: forwardedRef,
+    ref: mergedRef,
     className: resolveClassName(className, ownerState),
     ...otherProps,
   });

diff --git a/packages/mui-base/src/NumberField/NumberFieldInput.tsx b/packages/mui-base/src/NumberField/NumberFieldInput.tsx
@@ -1,10 +1,10 @@
 import * as React from 'react';
 import PropTypes from 'prop-types';
-import { useForkRef } from '../utils/useForkRef';
 import { useNumberFieldContext } from './NumberFieldContext';
 import type { NumberFieldInputProps } from './NumberField.types';
 import { resolveClassName } from '../utils/resolveClassName';
 import { evaluateRenderProp } from '../utils/evaluateRenderProp';
+import { useExtendedForkRef } from '../utils/useExtendedForkRef';
 
 function defaultRender(props: React.ComponentPropsWithRef<'input'>) {
   return <input {...props} />;
@@ -30,7 +30,7 @@ const NumberFieldInput = React.forwardRef(function NumberFieldInput(
 
   const { getInputProps, inputRef, inputValue, ownerState } = useNumberFieldContext('Input');
 
-  const mergedInputRef = useForkRef(forwardedRef, inputRef);
+  const mergedInputRef = useExtendedForkRef(render, forwardedRef, inputRef);
 
   const inputProps = getInputProps({
     ref: mergedInputRef,

diff --git a/packages/mui-base/src/NumberField/NumberFieldRoot.tsx b/packages/mui-base/src/NumberField/NumberFieldRoot.tsx
@@ -5,6 +5,7 @@ import type { NumberFieldOwnerState, NumberFieldRootProps } from './NumberField.
 import { resolveClassName } from '../utils/resolveClassName';
 import { useNumberField } from '../useNumberField/useNumberField';
 import { evaluateRenderProp } from '../utils/evaluateRenderProp';
+import { useExtendedForkRef } from '../utils/useExtendedForkRef';
 
 function defaultRender(props: React.ComponentPropsWithRef<'div'>) {
   return <div {...props} />;
@@ -77,8 +78,10 @@ const NumberFieldRoot = React.forwardRef(function NumberFieldRoot(
     [numberField, ownerState],
   );
 
+  const mergedRef = useExtendedForkRef(render, forwardedRef);
+
   const rootProps = {
-    ref: forwardedRef,
+    ref: mergedRef,
     className: resolveClassName(className, ownerState),
     ...otherProps,
   };

diff --git a/packages/mui-base/src/NumberField/NumberFieldScrubArea.tsx b/packages/mui-base/src/NumberField/NumberFieldScrubArea.tsx
@@ -1,10 +1,10 @@
 import * as React from 'react';
 import PropTypes from 'prop-types';
-import { useForkRef } from '../utils/useForkRef';
 import { useNumberFieldContext } from './NumberFieldContext';
 import type { NumberFieldScrubAreaProps } from './NumberField.types';
 import { resolveClassName } from '../utils/resolveClassName';
 import { evaluateRenderProp } from '../utils/evaluateRenderProp';
+import { useExtendedForkRef } from '../utils/useExtendedForkRef';
 
 function defaultRender(props: React.ComponentPropsWithRef<'span'>) {
   return <span {...props} />;
@@ -45,7 +45,7 @@ const NumberFieldScrubArea = React.forwardRef(function NumberFieldScrubArea(
     teleportDistance,
   }));
 
-  const mergedRef = useForkRef(scrubAreaRef, forwardedRef);
+  const mergedRef = useExtendedForkRef(render, scrubAreaRef, forwardedRef);
 
   const scrubAreaProps = getScrubAreaProps({
     ref: mergedRef,

diff --git a/packages/mui-base/src/NumberField/NumberFieldScrubAreaCursor.tsx b/packages/mui-base/src/NumberField/NumberFieldScrubAreaCursor.tsx
@@ -1,12 +1,12 @@
 import * as React from 'react';
 import * as ReactDOM from 'react-dom';
 import PropTypes from 'prop-types';
-import { useForkRef } from '../utils/useForkRef';
 import type { NumberFieldScrubAreaCursorProps } from './NumberField.types';
 import { isWebKit } from '../utils/detectBrowser';
 import { resolveClassName } from '../utils/resolveClassName';
 import { useNumberFieldContext } from './NumberFieldContext';
 import { evaluateRenderProp } from '../utils/evaluateRenderProp';
+import { useExtendedForkRef } from '../utils/useExtendedForkRef';
 
 function defaultRender(props: React.ComponentPropsWithRef<'span'>) {
   return <span {...props} />;
@@ -34,7 +34,7 @@ const NumberFieldScrubAreaCursor = React.forwardRef(function NumberFieldScrubAre
   const { isScrubbing, scrubAreaCursorRef, ownerState, getScrubAreaCursorProps } =
     useNumberFieldContext('ScrubAreaCursor');
 
-  const mergedRef = useForkRef(forwardedRef, scrubAreaCursorRef);
+  const mergedRef = useExtendedForkRef(render, forwardedRef, scrubAreaCursorRef);
 
   if (!isScrubbing || isWebKit()) {
     return null;

diff --git a/packages/mui-base/src/Switch/SwitchRoot.tsx b/packages/mui-base/src/Switch/SwitchRoot.tsx
@@ -8,6 +8,7 @@ import { resolveClassName } from '../utils/resolveClassName';
 import { SwitchContext } from './SwitchContext';
 import { useSwitchStyleHooks } from './useSwitchStyleHooks';
 import { evaluateRenderProp } from '../utils/evaluateRenderProp';
+import { useExtendedForkRef } from '../utils/useExtendedForkRef';
 
 function defaultRender(props: React.ComponentPropsWithRef<'button'>) {
   return <button type="button" {...props} />;
@@ -56,10 +57,11 @@ const SwitchRoot = React.forwardRef(function SwitchRoot(
 
   const className = resolveClassName(classNameProp, ownerState);
   const styleHooks = useSwitchStyleHooks(ownerState);
+  const mergedRef = useExtendedForkRef(render, forwardedRef);
 
   const buttonProps = {
     className,
-    ref: forwardedRef,
+    ref: mergedRef,
     ...styleHooks,
     ...other,
   };

diff --git a/packages/mui-base/src/Switch/SwitchThumb.tsx b/packages/mui-base/src/Switch/SwitchThumb.tsx
@@ -5,6 +5,7 @@ import { SwitchContext } from './SwitchContext';
 import { resolveClassName } from '../utils/resolveClassName';
 import { useSwitchStyleHooks } from './useSwitchStyleHooks';
 import { evaluateRenderProp } from '../utils/evaluateRenderProp';
+import { useExtendedForkRef } from '../utils/useExtendedForkRef';
 
 function defaultRender(props: React.ComponentPropsWithRef<'span'>) {
   return <span {...props} />;
@@ -24,10 +25,11 @@ const SwitchThumb = React.forwardRef(function SwitchThumb(
 
   const className = resolveClassName(classNameProp, ownerState);
   const styleHooks = useSwitchStyleHooks(ownerState);
+  const mergedRef = useExtendedForkRef(render, forwardedRef);
 
   const elementProps = {
     className,
-    ref: forwardedRef,
+    ref: mergedRef,
     ...styleHooks,
     ...other,
   };

diff --git a/packages/mui-base/src/utils/evaluateRenderProp.test.tsx b/packages/mui-base/src/utils/evaluateRenderProp.test.tsx
@@ -0,0 +1,22 @@
+import * as React from 'react';
+import { expect } from 'chai';
+import { evaluateRenderProp } from './evaluateRenderProp';
+
+describe('evaluateRenderProp', () => {
+  describe('function', () => {
+    it('should call render function with props and ownerState', () => {
+      const render = (props: any, ownerState: any) => <span {...props} {...ownerState} />;
+      expect(evaluateRenderProp(render, { id: 'a' }, { 'data-test': 'b' })).to.deep.equal(
+        <span id="a" data-test="b" />,
+      );
+    });
+  });
+
+  describe('element', () => {
+    it("merge props, preferring the rendering element's props", () => {
+      expect(evaluateRenderProp(<span id="a" />, { id: 'b', inputMode: 'text' }, {})).to.deep.equal(
+        <span id="a" inputMode="text" />,
+      );
+    });
+  });
+});
diff --git a/packages/mui-base/src/utils/evaluateRenderProp.ts b/packages/mui-base/src/utils/evaluateRenderProp.ts
@@ -8,5 +8,5 @@ export function evaluateRenderProp<ElementType extends React.ElementType, OwnerS
 ) {
   return typeof render === 'function'
     ? render(props, ownerState)
-    : React.cloneElement(render, props);
+    : React.cloneElement(render, { ...props, ...render.props });
 }
diff --git a/packages/mui-base/src/utils/useExtendedForkRef.ts b/packages/mui-base/src/utils/useExtendedForkRef.ts
@@ -0,0 +1,14 @@
+import type { BaseUIComponentProps } from './BaseUI.types';
+import { useForkRef } from './useForkRef';
+
+/**
+ * Merges the rendering element's `ref` in addition to the other `ref`s.
+ * @ignore - internal hook.
+ */
+export function useExtendedForkRef<ElementType extends React.ElementType, OwnerState>(
+  render: BaseUIComponentProps<ElementType, OwnerState>['render'],
+  ...refs: Array<React.Ref<any>>
+) {
+  const childRef = typeof render !== 'function' ? render.ref : null;
+  return useForkRef(childRef, ...refs);
+}
diff --git a/packages/mui-base/test/conformanceTests/renderProp.tsx b/packages/mui-base/test/conformanceTests/renderProp.tsx
@@ -62,5 +62,32 @@ export function testRenderProp(
       expect(instanceFromRef!.tagName).to.equal(Element.toUpperCase());
       expect(instanceFromRef!).to.have.attribute('data-testid', 'wrapped');
     });
+
+    it('should merge the rendering element ref with the custom component ref', () => {
+      let refA = null;
+      let refB = null;
+
+      function Test() {
+        return React.cloneElement(element, {
+          ref: (el: HTMLElement | null) => {
+            refA = el;
+          },
+          render: (
+            <Wrapper
+              ref={(el: HTMLElement | null) => {
+                refB = el;
+              }}
+            />
+          ),
+          'data-testid': 'wrapped',
+        });
+      }
+
+      render(<Test />);
+      expect(refA!.tagName).to.equal(Element.toUpperCase());
+      expect(refA!).to.have.attribute('data-testid', 'wrapped');
+      expect(refB!.tagName).to.equal(Element.toUpperCase());
+      expect(refB!).to.have.attribute('data-testid', 'wrapped');
+    });
   });
 }