fix: support controlled and uncontrolled states in form components

packages/odyssey-react-mui/src/@types/react-augment.d.ts

@@ -0,0 +1,19 @@
+/*!
+ * Copyright (c) 2023-present, Okta, Inc. and/or its affiliates. All rights reserved.
+ * The Okta software accompanied by this notice is provided pursuant to the Apache License, Version 2.0 (the "License.")
+ *
+ * You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *
+ * See the License for the specific language governing permissions and limitations under the License.
+ */
+
+import { FC } from "react";
+
+export interface ForwardRefWithType extends FC<WithForwardRefProps<Option>> {
+  <T extends Option>(props: WithForwardRefProps<T>): ReturnType<
+    FC<WithForwardRefProps<T>>
+  >;
+}

packages/odyssey-react-mui/src/Autocomplete.tsx

@@ -17,12 +17,16 @@ import {
   UseAutocompleteProps,
   AutocompleteValue,
 } from "@mui/material";
-import { memo, useCallback, useMemo } from "react";
+import { memo, useCallback, useMemo, useRef } from "react";
 
 import { Field } from "./Field";
 import { FieldComponentProps } from "./FieldComponentProps";
 import type { SeleniumProps } from "./SeleniumProps";
-import { useControlledState } from "./useControlledState";
+import {
+  ComponentControlledState,
+  useInputValues,
+  getControlState,
+} from "./inputUtils";
 
 export type AutocompleteProps<
   OptionType,
@@ -31,7 +35,6 @@ export type AutocompleteProps<
 > = {
   /**
    * The default value. Use when the component is not controlled.
-   * @default props.multiple ? [] : null
    */
   defaultValue?: UseAutocompleteProps<
     OptionType,
@@ -189,6 +192,45 @@ const Autocomplete = <
   getIsOptionEqualToValue,
   testId,
 }: AutocompleteProps<OptionType, HasMultipleChoices, IsCustomValueAllowed>) => {
+  const controlledStateRef = useRef(
+    getControlState({ controlledValue: value, uncontrolledValue: defaultValue })
+  );
+  const defaultValueProp = useMemo<
+    | AutocompleteValue<
+        OptionType,
+        HasMultipleChoices,
+        undefined,
+        IsCustomValueAllowed
+      >
+    | undefined
+  >(() => {
+    if (hasMultipleChoices) {
+      if (value === undefined) {
+        return defaultValue;
+      }
+      return [] as AutocompleteValue<
+        OptionType,
+        HasMultipleChoices,
+        undefined,
+        IsCustomValueAllowed
+      >;
+    }
+    return value === undefined ? defaultValue : undefined;
+  }, [defaultValue, hasMultipleChoices, value]);
+
+  const valueProps = useInputValues({
+    defaultValue: defaultValueProp,
+    value: value,
+    controlState: controlledStateRef.current,
+  });
+
+  const inputValueProp = useMemo(() => {
+    if (controlledStateRef.current === ComponentControlledState.CONTROLLED) {
+      return { inputValue };
+    }
+    return undefined;
+  }, [inputValue]);
+
   const renderInput = useCallback(
     ({ InputLabelProps, InputProps, ...params }) => (
       <Field
@@ -223,39 +265,6 @@ const Autocomplete = <
     ),
     [errorMessage, hint, isOptional, label, nameOverride]
   );
-
-  const defaultValuesProp = useMemo<
-    | AutocompleteValue<
-        OptionType,
-        HasMultipleChoices,
-        undefined,
-        IsCustomValueAllowed
-      >
-    | undefined
-  >(() => {
-    if (hasMultipleChoices) {
-      return defaultValue === undefined
-        ? ([] as AutocompleteValue<
-            OptionType,
-            HasMultipleChoices,
-            undefined,
-            IsCustomValueAllowed
-          >)
-        : defaultValue;
-    }
-    return defaultValue ?? undefined;
-  }, [defaultValue, hasMultipleChoices]);
-
-  const [localValue, setLocalValue] = useControlledState({
-    controlledValue: value,
-    uncontrolledValue: defaultValuesProp,
-  });
-
-  const [localInputValue, setLocalInputValue] = useControlledState({
-    controlledValue: inputValue,
-    uncontrolledValue: undefined,
-  });
-
   const onChange = useCallback<
     NonNullable<
       UseAutocompleteProps<
@@ -267,10 +276,9 @@ const Autocomplete = <
     >
   >(
     (event, value, reason, details) => {
-      setLocalValue(value);
       onChangeProp?.(event, value, reason, details);
     },
-    [onChangeProp, setLocalValue]
+    [onChangeProp]
   );
 
   const onInputChange = useCallback<
@@ -284,18 +292,18 @@ const Autocomplete = <
     >
   >(
     (event, value, reason) => {
-      setLocalInputValue(value);
       onInputChangeProp?.(event, value, reason);
     },
-    [onInputChangeProp, setLocalInputValue]
+    [onInputChangeProp]
   );
 
   return (
     <MuiAutocomplete
+      {...valueProps}
+      {...inputValueProp}
       // AutoComplete is wrapped in a div within MUI which does not get the disabled attr. So this aria-disabled gets set in the div
       aria-disabled={isDisabled}
       data-se={testId}
-      defaultValue={defaultValuesProp}
       disableCloseOnSelect={hasMultipleChoices}
       disabled={isDisabled}
       freeSolo={isCustomValueAllowed}
@@ -310,8 +318,6 @@ const Autocomplete = <
       options={options}
       readOnly={isReadOnly}
       renderInput={renderInput}
-      value={localValue}
-      inputValue={localInputValue}
       isOptionEqualToValue={getIsOptionEqualToValue}
     />
   );

packages/odyssey-react-mui/src/Checkbox.tsx

@@ -11,7 +11,7 @@
  */
 
 import { useTranslation } from "react-i18next";
-import { memo, useCallback, useMemo } from "react";
+import { memo, useCallback, useMemo, useRef } from "react";
 import {
   Checkbox as MuiCheckbox,
   CheckboxProps as MuiCheckboxProps,
@@ -22,7 +22,7 @@ import {
 import { FieldComponentProps } from "./FieldComponentProps";
 import { Typography } from "./Typography";
 import type { SeleniumProps } from "./SeleniumProps";
-import { useControlledState } from "./useControlledState";
+import { ComponentControlledState, getControlState } from "./inputUtils";
 import { CheckedFieldProps } from "./FormCheckedProps";
 
 export const checkboxValidityValues = ["valid", "invalid", "inherit"] as const;
@@ -90,10 +90,18 @@ const Checkbox = ({
   value,
 }: CheckboxProps) => {
   const { t } = useTranslation();
-  const [isLocalChecked, setIsLocalChecked] = useControlledState({
-    controlledValue: isChecked,
-    uncontrolledValue: isDefaultChecked,
-  });
+  const controlledStateRef = useRef(
+    getControlState({
+      controlledValue: isChecked,
+      uncontrolledValue: isDefaultChecked,
+    })
+  );
+  const inputValues = useMemo(() => {
+    if (controlledStateRef.current === ComponentControlledState.CONTROLLED) {
+      return { checked: isChecked };
+    }
+    return { defaultChecked: isDefaultChecked };
+  }, [isDefaultChecked, isChecked]);
 
   const label = useMemo(() => {
     return (
@@ -114,10 +122,9 @@ const Checkbox = ({
 
   const onChange = useCallback<NonNullable<MuiCheckboxProps["onChange"]>>(
     (event, checked) => {
-      setIsLocalChecked(checked);
       onChangeProp?.(event, checked);
     },
-    [onChangeProp, setIsLocalChecked]
+    [onChangeProp]
   );
 
   return (
@@ -134,7 +141,7 @@ const Checkbox = ({
       }
       control={
         <MuiCheckbox
-          checked={isLocalChecked}
+          {...inputValues}
           indeterminate={isIndeterminate}
           onChange={onChange}
           required={isRequired}