[Base] Improve the return type of useSlotProps

Make useSlotProps (and the underlying appendOwnerState) aware of the type of element
on which the props should be applied. Based on this, ownerState will either be present
or undefined in the resulting type.

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>,