React + Material UI - Best way to prevent child tree from remount when toggling parent theme

Question

Background

I wanted to follow Material UI's implementation of toggling UI's dark/light mode theme. Link.

I have encapsulated its implementation into a custom hook that returns theme-related properties when called. This is called at my App() level.

import { createContext, useState, useMemo } from 'react';
import { ThemeProvider, createTheme } from '@mui/material/styles';


export default function useTheme() {

    const ColorModeContext = createContext({ toggleColorMode: () => {}, mode: null})
    const [ mode, setMode ] = useState('light')

    const colorMode = useMemo(
        () => ({
            toggleColorMode : () => {
              setMode((prevMode) => (prevMode === 'light' ? 'dark' : 'light'))  
            },
        })
    , [])  


    // recreate theme everytime mode changes
    const appTheme = useMemo( () => (createTheme({
        palette: {
          mode: mode,
          primary: ...
        },
      })), [mode])
    
    
    
    return {
        ColorModeContext,
        colorMode,
        ThemeProvider,
        appTheme
    }
}

At the App() level, I have it returns the following App Component:

...
    const {
        ColorModeContext,
        colorMode,
        ThemeProvider,
        appTheme
    } = useTheme()

    // another custom hook to return auth related properties
    const {
        AuthContext,
        authed,
        ...
    } = useAuth()
   return (
     <ColorModeContext.Provider value={colorMode}>
            <ThemeProvider theme={appTheme}>
                <CssBaseline />
                <AuthContext.Provider value=...>
                    <Box>
                        <Router>
                            <Header 
                                ColorModeContext={ColorModeContext}
                                theme={appTheme}
                                authed={authed}
                                ...
                            /> 
                            <Routes>
                                <Route 
                                    path="/" 
                                    element={
                                        <Signin AuthContext={AuthContext} />
                                    } 
                                />
                                <Route 
                                    path="/create" 
                                    element={
                                        <RequiredAuth authed={authed}>
                                            <Create />
                                        </RequiredAuth>
                                    } 
                                />
                                <Route 
                                    path="/query" 
                                    element={
                                        <RequiredAuth authed={authed}>
                                            <Query />
                                        </RequiredAuth>
                                    } 
                                />
                            </Routes>
                            ...
                        </Router>
                    </Box>
                </AuthContext.Provider>
            </ThemeProvider>
        </ColorModeContext.Provider>
)

Lastly, in my header component (as the context consumer, to have the toggle callback function exposed in the header)

...
const { ColorModeContext, authed, ... } = props;

return (
        <ColorModeContext.Consumer>
            {
                ({toggleColorMode}) => (
                    <Box>
                        <AppBar position="static" enableColorOnDark>
                            <Toolbar>
                              ...
                                <IconButton
                                  size="small"
                                  onClick={toggleColorMode}
                                  color="inherit"
                                >
                                   {
                                       theme.palette.mode === 'dark' ? 
                                          <LightModeIcon /> : <Brightness3Icon />
                                   }
                               </IconButton>
                            ...                   

                            </Toolbar>
                        </AppBar>
                    </Box>
                )
            }
        </ColorModeContext.Consumer>
)

ISSUE

Toggling the theme works, however, If I am inside my component <Create /> or <Query /> and I toggle the theme, the entire component tree (Create or Query) is remounted, and all states within the component refreshed to its initial state.

Basically, if I am inside my <Create /> component, and I am filling out the create form, which the values of the form are maintained by useState, as soon as I toggle the theme, all of the values resets to their initial value (which was passed into the useState hook)

QUESTION

Is there a way to prevent this remount from happening? I know it's not just rerendering the component, because the state is reinitialized. If it's not a rerender issue, then would React.memo work here? If not, what would be the best way to toggle theme (context at a parent component lvl) while inside a child component without having the component remount.

---

Adding a codesandbox link. It should include a simple example of the issue that I am having. The child component remounts (as logged by the Child's useEffect) when I toggle parent's theme mode.

Answer 1

The useTheme hook should return the provided context value, not the context itself. This is recreating a context provider and remounting children.

Example:

useTheme

import { createContext, useState, useMemo, useContext } from "react";
import { ThemeProvider, createTheme } from "@mui/material/styles";

const ColorModeContext = createContext({
  toggleColorMode: () => {},
  mode: null
});

const ColorModeContextProvider = ({ children }) => {
  const [mode, setMode] = useState("light");

  const toggleColorMode = () => {
    setMode((prevMode) => (prevMode === "light" ? "dark" : "light"));
  };

  // recreate theme everytime mode changes
  const appTheme = useMemo(
    () =>
      createTheme({
        palette: {
          mode
        }
      }),
    [mode]
  );

  return (
    <ColorModeContext.Provider
      value={{
        toggleColorMode,
        mode
      }}
    >
      <ThemeProvider theme={appTheme}>{children}</ThemeProvider>
    </ColorModeContext.Provider>
  );
};

export const useTheme = () => useContext(ColorModeContext);

export default ColorModeContextProvider;

App

Import the new ColorModeContextProvider provider component to provide the theme and color mode contexts, and the useTheme hook to be used in the Header component to access the toggleColorMode callback.

import { useState, useEffect } from "react";
import { Container, CssBaseline } from "@mui/material";
import ColorModeContextProvider, { useTheme } from "./useTheme";
import useAuth from "./useAuth";

function App() {
  const { AuthContext, authed, user, login, logout } = useAuth();

  return (
    <ColorModeContextProvider>
      <CssBaseline />
      <AuthContext.Provider value={{ authed, user, login, logout }}>
        <Container maxWidth="lg">
          <Header />
        </Container>
        <Child />
      </AuthContext.Provider>
    </ColorModeContextProvider>
  );
}

function Header() {
  const { toggleColorMode } = useTheme();
  return <button onClick={toggleColorMode}>ToggleTheme</button>;
}

function Child() {
  const [counter, setCounter] = useState(0);

  useEffect(() => {
    return () => {
      console.log("Child component unmounting");
    };
  }, []);

  return (
    <>
      <button
        onClick={() => {
          setCounter((count) => count + 1);
        }}
      >
        Add
      </button>
      <div>{counter}</div>
    </>
  );
}

export default App;

The useAuth hook needs a similar refactoring applied to it as well.