docs: Fix hydration error on mobile

Author: iwoplaza

apps/typegpu-docs/src/components/ExampleLayout.tsx

@@ -1,4 +1,4 @@
-import { useAtom, useSetAtom } from 'jotai';
+import { useSetAtom } from 'jotai';
 import type { ReactNode } from 'react';
 import { useId, useRef } from 'react';
 import CrossSvg from '../assets/cross.svg';
@@ -13,16 +13,15 @@ import {
 import { SearchableExampleList } from './SearchableExampleList.tsx';
 import { Button } from './design/Button.tsx';
 import { Toggle } from './design/Toggle.tsx';
-import { useHydrated } from '../utils/useHydrated.ts';
+import { useHydratedAtom } from '../utils/useHydrated.ts';
 
 interface ExampleLayoutProps {
   children?: ReactNode | undefined;
 }
 
 export function ExampleLayout(props: ExampleLayoutProps) {
-  const [menuShown, setMenuShown] = useAtom(menuShownAtom);
-  const [codeShown, setCodeShown] = useAtom(codeEditorShownAtom);
-  const hydrated = useHydrated();
+  const [menuShown, setMenuShown] = useHydratedAtom(menuShownAtom, false);
+  const [codeShown, setCodeShown] = useHydratedAtom(codeEditorShownAtom, false);
 
   return (
     <>
@@ -35,7 +34,7 @@ export function ExampleLayout(props: ExampleLayoutProps) {
 
         <Button onClick={() => setCodeShown((prev) => !prev)}>
           {/* Applying the actual label only after the component has been hydrated */}
-          {codeShown && hydrated ? 'Preview' : 'Code'}
+          {codeShown ? 'Preview' : 'Code'}
         </Button>
       </div>
 
@@ -49,8 +48,9 @@ export function ExampleLayout(props: ExampleLayoutProps) {
 
 function SideMenu() {
   const setMenuShown = useSetAtom(menuShownAtom);
-  const [experimentalShowing, setExperimentalShowing] = useAtom(
+  const [experimentalShowing, setExperimentalShowing] = useHydratedAtom(
     experimentalExamplesShownAtom,
+    true,
   );
   const scrollRef = useRef<HTMLDivElement>(null);
   const experimentalExamplesToggleId = useId();

apps/typegpu-docs/src/utils/useHydrated.ts

@@ -1,3 +1,10 @@
+import {
+  type ExtractAtomArgs,
+  type ExtractAtomResult,
+  type ExtractAtomValue,
+  useAtom,
+  type WritableAtom,
+} from 'jotai';
 import { useEffect, useState } from 'react';
 
 /**
@@ -13,3 +20,26 @@ export function useHydrated() {
 
   return isHydrated;
 }
+
+type AnyWritableAtom = WritableAtom<unknown, never[], unknown>;
+
+type UseHydratedAtomReturn<T extends AnyWritableAtom> = [
+  ExtractAtomValue<T>,
+  (...args: ExtractAtomArgs<T>) => ExtractAtomResult<T>,
+];
+
+/**
+ * The same as `useAtom`, but returns `unhydratedValue` instead of the atom's value
+ * if the component is yet to be hydrated.
+ */
+export function useHydratedAtom<
+  T extends AnyWritableAtom,
+>(atom: T, unhydratedValue: ExtractAtomValue<T>): UseHydratedAtomReturn<T> {
+  const [atomValue, atomSetter] = useAtom(atom);
+  const hydrated = useHydrated();
+
+  return [
+    hydrated ? atomValue : unhydratedValue,
+    atomSetter,
+  ] as UseHydratedAtomReturn<T>;
+}