[Base] Improve the return type of useSlotProps

packages/mui-base/src/FormControlUnstyled/FormControlUnstyled.tsx

@@ -131,29 +131,30 @@ const FormControlUnstyled = React.forwardRef(function FormControlUnstyled<
 
   const classes = useUtilityClasses(ownerState);
 
+  const renderChildren = () => {
+    if (typeof children === 'function') {
+      return children(childContext);
+    }
+
+    return children;
+  };
+
   const Root = component ?? components.Root ?? 'div';
   const rootProps: WithOptionalOwnerState<FormControlUnstyledRootSlotProps> = useSlotProps({
     elementType: Root,
     externalSlotProps: componentsProps.root,
     externalForwardedProps: other,
     additionalProps: {
       ref,
+      children: renderChildren(),
     },
     ownerState,
     className: classes.root,
   });
 
-  const renderChildren = () => {
-    if (typeof children === 'function') {
-      return children(childContext);
-    }
-
-    return children;
-  };
-
   return (
     <FormControlUnstyledContext.Provider value={childContext}>
-      <Root {...rootProps}>{renderChildren()}</Root>
+      <Root {...rootProps} />
     </FormControlUnstyledContext.Provider>
   );
 }) as OverridableComponent<FormControlUnstyledTypeMap>;

packages/mui-base/src/InputUnstyled/InputUnstyled.types.ts

@@ -145,7 +145,7 @@ export type InputUnstyledInputSlotProps = Simplify<
     'aria-labelledby': React.AriaAttributes['aria-labelledby'];
     autoComplete: string | undefined;
     autoFocus: boolean | undefined;
-    className: string;
+    className?: string;
     id: string | undefined;
     name: string | undefined;
     onKeyDown: React.KeyboardEventHandler<HTMLInputElement> | undefined;

packages/mui-base/src/SwitchUnstyled/SwitchUnstyled.types.ts

@@ -66,26 +66,26 @@ export type SwitchUnstyledOwnerState = Simplify<
 
 export type SwitchUnstyledRootSlotProps = {
   ownerState: SwitchUnstyledOwnerState;
-  className: string;
+  className?: string;
   children?: React.ReactNode;
 };
 
 export type SwitchUnstyledThumbSlotProps = {
   ownerState: SwitchUnstyledOwnerState;
-  className: string;
+  className?: string;
   children?: React.ReactNode;
 };
 
 export type SwitchUnstyledTrackSlotProps = {
   ownerState: SwitchUnstyledOwnerState;
-  className: string;
+  className?: string;
   children?: React.ReactNode;
 };
 
 export type SwitchUnstyledInputSlotProps = Simplify<
   UseSwitchInputSlotProps & {
     ownerState: SwitchUnstyledOwnerState;
-    className: string;
+    className?: string;
     children?: React.ReactNode;
   }
 >;

packages/mui-base/src/utils/appendOwnerState.spec.tsx

@@ -0,0 +1,28 @@
+import * as React from 'react';
+import appendOwnerState from './appendOwnerState';
+
+const divProps = appendOwnerState('div', { otherProp: true }, { ownerStateProps: true });
+
+// ownerState is not available on a host component
+// @ts-expect-error
+const test1 = divProps.ownerState.ownerStateProps;
+// @ts-expect-error
+const test2 = divProps.ownerState?.ownerStateProps;
+
+const componentProps = appendOwnerState(
+  () => <div />,
+  { otherProp: true },
+  { ownerStateProps: true },
+);
+
+// ownerState is present on a custom component
+const test3: boolean = componentProps.ownerState.ownerStateProps;
+
+function test(element: React.ElementType) {
+  const props = appendOwnerState(element, { otherProp: true }, { ownerStateProps: true });
+
+  // ownerState may be present on a provided element type (it depends on its exact type)
+  // @ts-expect-error
+  const test4 = props.ownerState.ownerStateProps;
+  const test5: boolean | undefined = props.ownerState?.ownerStateProps;
+}

packages/mui-base/src/utils/appendOwnerState.ts

@@ -1,26 +1,53 @@
+import { Simplify } from '@mui/types';
+import React from 'react';
 import isHostComponent from './isHostComponent';
 
+/**
+ * Type of the ownerState based on the type of an element it applies to.
+ * This resolves to the provided OwnerState for React components and `undefined` for host components.
+ * Falls back to `OwnerState | undefined` when the exact type can't be determined in development time.
+ */
+type OwnerStateWhenApplicable<
+  ElementType extends React.ElementType,
+  OwnerState,
+> = ElementType extends React.ComponentType<any>
+  ? OwnerState
+  : ElementType extends keyof JSX.IntrinsicElements
+  ? undefined
+  : OwnerState | undefined;
+
+export type AppendOwnerStateReturnType<
+  ElementType extends React.ElementType,
+  OtherProps,
+  OwnerState,
+> = Simplify<
+  OtherProps & {
+    ownerState: OwnerStateWhenApplicable<ElementType, OwnerState>;
+  }
+>;
+
 /**
  * Appends the ownerState object to the props, merging with the existing one if necessary.
  *
- * @param elementType Type of the element that owns the `existingProps`. If the element is a DOM node, `ownerState` are not applied.
- * @param existingProps Props of the element.
+ * @param elementType Type of the element that owns the `existingProps`. If the element is a DOM node, `ownerState` is not applied.
+ * @param otherProps Props of the element.
  * @param ownerState
  */
 export default function appendOwnerState<
-  TExistingProps extends Record<string, any>,
-  TOwnerState extends {},
+  ElementType extends React.ElementType,
+  OtherProps extends Record<string, any>,
+  OwnerState,
 >(
-  elementType: React.ElementType,
-  existingProps: TExistingProps = {} as TExistingProps,
-  ownerState: TOwnerState,
-): TExistingProps & { ownerState?: TOwnerState } {
+  elementType: ElementType,
+  otherProps: OtherProps = {} as OtherProps,
+  ownerState: OwnerState,
+): AppendOwnerStateReturnType<ElementType, OtherProps, OwnerState> {
   if (isHostComponent(elementType)) {
-    return existingProps;
+    return otherProps as AppendOwnerStateReturnType<ElementType, OtherProps, OwnerState>;
   }
 
   return {
-    ...existingProps,
-    ownerState: { ...existingProps.ownerState, ...ownerState },
-  };
+    ...otherProps,
+    ownerState: { ...otherProps.ownerState, ...ownerState },
+  } as AppendOwnerStateReturnType<ElementType, OtherProps, OwnerState>;
 }

packages/mui-base/src/utils/useSlotProps.test.tsx

@@ -8,13 +8,15 @@ import useSlotProps, { UseSlotPropsParameters, UseSlotPropsResult } from './useS
 const { render } = createRenderer();
 
 function callUseSlotProps<
+  ElementType extends React.ElementType,
   SlotProps,
   ExternalForwardedProps,
   ExternalSlotProps,
   AdditionalProps,
   OwnerState,
 >(
   parameters: UseSlotPropsParameters<
+    ElementType,
     SlotProps,
     ExternalForwardedProps,
     ExternalSlotProps,
@@ -27,6 +29,7 @@ function callUseSlotProps<
       _: unknown,
       ref: React.Ref<
         UseSlotPropsResult<
+          ElementType,
           SlotProps,
           ExternalForwardedProps,
           ExternalSlotProps,
@@ -44,6 +47,7 @@ function callUseSlotProps<
   const ref =
     React.createRef<
       UseSlotPropsResult<
+        ElementType,
         SlotProps,
         ExternalForwardedProps,
         ExternalSlotProps,

packages/mui-base/src/utils/useSlotProps.ts

@@ -1,10 +1,11 @@
 import * as React from 'react';
 import { unstable_useForkRef as useForkRef } from '@mui/utils';
-import appendOwnerState from './appendOwnerState';
+import appendOwnerState, { AppendOwnerStateReturnType } from './appendOwnerState';
 import mergeSlotProps, { MergeSlotPropsParameters, WithRef } from './mergeSlotProps';
 import resolveComponentProps from './resolveComponentProps';
 
 export type UseSlotPropsParameters<
+  ElementType extends React.ElementType,
   SlotProps,
   ExternalForwardedProps,
   ExternalSlotProps,
@@ -17,7 +18,7 @@ export type UseSlotPropsParameters<
   /**
    * The type of the component used in the slot.
    */
-  elementType: React.ElementType;
+  elementType: ElementType;
   /**
    * The `componentsProps.*` of the unstyled component.
    */
@@ -32,16 +33,20 @@ export type UseSlotPropsParameters<
 };
 
 export type UseSlotPropsResult<
+  ElementType extends React.ElementType,
   SlotProps,
   ExternalForwardedProps,
   ExternalSlotProps,
   AdditionalProps,
   OwnerState,
-> = Omit<SlotProps & ExternalSlotProps & ExternalForwardedProps & AdditionalProps, 'ref'> & {
-  className?: string | undefined;
-  ownerState?: OwnerState | undefined;
-  ref: (instance: any | null) => void;
-};
+> = AppendOwnerStateReturnType<
+  ElementType,
+  Omit<SlotProps & ExternalSlotProps & ExternalForwardedProps & AdditionalProps, 'ref'> & {
+    className?: string | undefined;
+    ref: (instance: any | null) => void;
+  },
+  OwnerState
+>;
 
 /**
  * Builds the props to be passed into the slot of an unstyled component.
@@ -51,13 +56,15 @@ export type UseSlotPropsResult<
  * @param parameters.getSlotProps - A function that returns the props to be passed to the slot component.
  */
 export default function useSlotProps<
+  ElementType extends React.ElementType,
   SlotProps,
   ExternalForwardedProps,
   ExternalSlotProps,
   AdditionalProps,
   OwnerState,
 >(
   parameters: UseSlotPropsParameters<
+    ElementType,
     SlotProps,
     ExternalForwardedProps,
     WithRef<ExternalSlotProps>,