From 613e4444f7f90abe4ce250d5000a398832436ae3 Mon Sep 17 00:00:00 2001
From: atomiks <cc.glows@gmail.com>
Date: Fri, 12 Apr 2024 12:27:24 +1000
Subject: [PATCH] Fix clone prop merging of ref and other props

---
 .../src/Checkbox/CheckboxIndicator.tsx        |  4 ++-
 .../mui-base/src/Checkbox/CheckboxRoot.tsx    |  4 ++-
 .../src/NumberField/NumberFieldDecrement.tsx  |  5 +++-
 .../src/NumberField/NumberFieldGroup.tsx      |  5 +++-
 .../src/NumberField/NumberFieldIncrement.tsx  |  5 +++-
 .../src/NumberField/NumberFieldInput.tsx      |  4 +--
 .../src/NumberField/NumberFieldRoot.tsx       |  5 +++-
 .../src/NumberField/NumberFieldScrubArea.tsx  |  4 +--
 .../NumberFieldScrubAreaCursor.tsx            |  4 +--
 packages/mui-base/src/Switch/SwitchRoot.tsx   |  4 ++-
 packages/mui-base/src/Switch/SwitchThumb.tsx  |  4 ++-
 .../src/utils/evaluateRenderProp.test.tsx     | 22 +++++++++++++++
 .../mui-base/src/utils/evaluateRenderProp.ts  |  2 +-
 .../mui-base/src/utils/useExtendedForkRef.ts  | 14 ++++++++++
 .../test/conformanceTests/renderProp.tsx      | 27 +++++++++++++++++++
 15 files changed, 98 insertions(+), 15 deletions(-)
 create mode 100644 packages/mui-base/src/utils/evaluateRenderProp.test.tsx
 create mode 100644 packages/mui-base/src/utils/useExtendedForkRef.ts

diff --git a/packages/mui-base/src/Checkbox/CheckboxIndicator.tsx b/packages/mui-base/src/Checkbox/CheckboxIndicator.tsx
index 6e409c6620..41d06bbb5e 100644
--- 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,6 +35,7 @@ const CheckboxIndicator = React.forwardRef(function CheckboxIndicator(
   }
 
   const styleHooks = useCheckboxStyleHooks(ownerState);
+  const mergedRef = useExtendedForkRef(render, forwardedRef);
 
   if (!keepMounted && !ownerState.checked && !ownerState.indeterminate) {
     return null;
@@ -41,7 +43,7 @@ const CheckboxIndicator = React.forwardRef(function CheckboxIndicator(
 
   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
index 1662433376..6f58c30fa9 100644
--- 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
index 1a5d6c60bc..6b660a6faa 100644
--- 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
index 3880d1102b..99bb20f9b1 100644
--- 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
index e3f55f9a50..8f3e00e801 100644
--- 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
index 7398e57881..3ac4d8d1f9 100644
--- 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
index 2cebbf69cf..6fa924e38f 100644
--- 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
index df326b7e1b..17645b373b 100644
--- 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
index ef74fa4ac7..4de39c22a0 100644
--- 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
index 414f139893..03e222a1cb 100644
--- 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
index b81d66774b..d07e2f8fab 100644
--- 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
new file mode 100644
index 0000000000..c1fdbdb69a
--- /dev/null
+++ 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
index aacd4d81e6..bba250e17e 100644
--- 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
new file mode 100644
index 0000000000..0516f2fb91
--- /dev/null
+++ 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
index a4c57c9f64..bc53830692 100644
--- 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');
+    });
   });
 }