Allow previous state reuse in usePersistedState (#645)

swan-io · Dec 6, 2024 · dc06cd6 · dc06cd6
1 parent ec8a3b9
commit dc06cd6
Showing 1 changed file with 41 additions and 33 deletions.
diff --git a/packages/lake/src/hooks/usePersistedState.ts b/packages/lake/src/hooks/usePersistedState.ts
@@ -1,56 +1,64 @@
 import { Option, Result } from "@swan-io/boxed";
-import { useCallback, useEffect, useState } from "react";
+import { useCallback, useEffect, useMemo, useState } from "react";
 
 const getItem = (key: string) => Result.fromExecution(() => localStorage.getItem(key)).getOr(null);
 
+const setItem = (key: string, value: string | null) => {
+  try {
+    if (value != null) {
+      localStorage.setItem(key, value);
+    } else {
+      localStorage.removeItem(key);
+    }
+  } catch {
+    // ignore
+  }
+};
+
 const parseRawValue = <T>(rawValue: string | null, defaultValue: T) =>
   Result.fromExecution(() => (rawValue != null ? (JSON.parse(rawValue) as T) : rawValue))
     .toOption()
     .flatMap(Option.fromNullable)
     .getOr(defaultValue);
 
+const stringifyValue = <T>(value: T) =>
+  Result.fromExecution<string | null>(() => JSON.stringify(value))
+    .toOption()
+    .getOr(null);
+
 export const usePersistedState = <T>(key: string, defaultValue: T) => {
-  const [state, setState] = useState(() => {
-    const rawValue = getItem(key);
-    const value = parseRawValue(rawValue, defaultValue);
-    return { defaultValue, value, rawValue };
-  });
+  const [stableDefaultValue] = useState(() => defaultValue);
+  const [rawValue, setRawValue] = useState(() => getItem(key));
 
-  const updateRawValue = useCallback((rawValue: string | null) => {
-    setState(prevState => {
-      if (rawValue === prevState.rawValue) {
-        return prevState; // skip update if rawValue didn't changed
-      } else {
-        const { defaultValue } = prevState;
-        const value = parseRawValue(rawValue, defaultValue);
-        return { defaultValue, value, rawValue };
-      }
-    });
-  }, []);
+  const value = useMemo(
+    () => parseRawValue(rawValue, stableDefaultValue),
+    [rawValue, stableDefaultValue],
+  );
 
   const setPersistedState = useCallback(
-    (value: T | null) => {
-      try {
-        if (value != null) {
-          const rawValue = JSON.stringify(value);
-          updateRawValue(rawValue);
-          localStorage.setItem(key, rawValue);
-        } else {
-          updateRawValue(null);
-          localStorage.removeItem(key);
-        }
-      } catch {
-        // ignore
+    (value: T | null | ((prevState: T) => T)): void => {
+      if (typeof value !== "function") {
+        const rawValue = value != null ? stringifyValue(value) : null;
+        setItem(key, rawValue);
+        setRawValue(rawValue);
+      } else {
+        setRawValue(prevState => {
+          const prevValue = parseRawValue(prevState, stableDefaultValue);
+          const nextValue = (value as (prevState: T) => T)(prevValue);
+          const rawValue = stringifyValue(nextValue);
+          setItem(key, rawValue);
+          return rawValue;
+        });
       }
     },
-    [key, updateRawValue],
+    [key, stableDefaultValue],
   );
 
   useEffect(() => {
     const listener = (event: StorageEvent) => {
       if (event.storageArea === localStorage && (event.key === key || event.key === null)) {
         const rawValue = getItem(key);
-        updateRawValue(rawValue);
+        setRawValue(rawValue);
       }
     };
 
@@ -59,7 +67,7 @@ export const usePersistedState = <T>(key: string, defaultValue: T) => {
     return () => {
       window.removeEventListener("storage", listener);
     };
-  }, [key, updateRawValue]);
+  }, [key]);
 
-  return [state.value, setPersistedState] as const;
+  return [value, setPersistedState] as const;
 };