From cfd6e4f527adc3661c4807ba381b5b43b64d35ff Mon Sep 17 00:00:00 2001
From: Sterling Camden <rusty.camden@gmail.com>
Date: Fri, 12 Apr 2024 21:42:06 +0000
Subject: [PATCH] allow the setHash option to be overridden on a per atom set
 basis

---
 __tests__/atomWithHash_spec.tsx | 68 +++++++++++++++++++++++++++++++++
 package.json                    |  2 +
 src/atomWithHash.ts             | 65 ++++++++++++++++++++++---------
 3 files changed, 116 insertions(+), 19 deletions(-)

diff --git a/__tests__/atomWithHash_spec.tsx b/__tests__/atomWithHash_spec.tsx
index 1491fa9..34eb5ae 100644
--- a/__tests__/atomWithHash_spec.tsx
+++ b/__tests__/atomWithHash_spec.tsx
@@ -133,6 +133,74 @@ describe('atomWithHash', () => {
       expect(window.location.search).toEqual('?q=foo');
       expect(window.location.hash).toEqual('#count=2');
     });
+    window.history.back();
+    await waitFor(() => {
+      expect(window.location.pathname).toEqual('/');
+      expect(window.location.search).toEqual('');
+      expect(window.location.hash).toEqual('');
+    });
+  });
+
+  it('keeping current path only for one set', async () => {
+    const countAtom = atomWithHash('count', 0);
+
+    const Counter = () => {
+      const [count, setCount] = useAtom(countAtom);
+      useEffect(() => {
+        setCount(1, { setHash: 'replaceState' });
+      }, []);
+      return (
+        <>
+          <div>count: {count}</div>
+          <button type="button" onClick={() => setCount((c) => c + 1)}>
+            button
+          </button>
+        </>
+      );
+    };
+
+    window.history.pushState(null, '', '/?q=foo');
+    const { findByText, getByText } = render(
+      <StrictMode>
+        <Counter />
+      </StrictMode>,
+    );
+
+    await findByText('count: 1');
+    await waitFor(() => {
+      expect(window.location.pathname).toEqual('/');
+      expect(window.location.search).toEqual('?q=foo');
+      expect(window.location.hash).toEqual('#count=1');
+    });
+    fireEvent.click(getByText('button'));
+    await findByText('count: 2');
+    expect(window.location.pathname).toEqual('/');
+    expect(window.location.search).toEqual('?q=foo');
+    expect(window.location.hash).toEqual('#count=2');
+
+    window.history.pushState(null, '', '/another');
+    await waitFor(() => {
+      expect(window.location.pathname).toEqual('/another');
+    });
+
+    window.history.back();
+    await waitFor(() => {
+      expect(window.location.pathname).toEqual('/');
+      expect(window.location.search).toEqual('?q=foo');
+      expect(window.location.hash).toEqual('#count=2');
+    });
+    window.history.back();
+    await waitFor(() => {
+      expect(window.location.pathname).toEqual('/');
+      expect(window.location.search).toEqual('?q=foo');
+      expect(window.location.hash).toEqual('#count=1');
+    });
+    window.history.back();
+    await waitFor(() => {
+      expect(window.location.pathname).toEqual('/');
+      expect(window.location.search).toEqual('');
+      expect(window.location.hash).toEqual('');
+    });
   });
 
   it('should optimize value to prevent unnecessary re-renders', async () => {
diff --git a/package.json b/package.json
index 06e1838..3d2edeb 100644
--- a/package.json
+++ b/package.json
@@ -29,8 +29,10 @@
     "compile": "microbundle build -f modern,umd --globals react=React",
     "postcompile": "cp dist/index.modern.mjs dist/index.modern.js && cp dist/index.modern.mjs.map dist/index.modern.js.map",
     "test": "run-s eslint tsc-test jest",
+    "test:debug": "run-s eslint tsc-test jest:debug",
     "eslint": "eslint --ext .js,.ts,.tsx .",
     "jest": "jest",
+    "jest:debug" : "node --inspect ./node_modules/jest/bin/jest.js --runInBand",
     "tsc-test": "tsc --project . --noEmit",
     "examples:01_minimal": "DIR=01_minimal EXT=js webpack serve",
     "examples:02_typescript": "DIR=02_typescript EXT=tsx webpack serve",
diff --git a/src/atomWithHash.ts b/src/atomWithHash.ts
index 76b5d58..1d81b6c 100644
--- a/src/atomWithHash.ts
+++ b/src/atomWithHash.ts
@@ -15,6 +15,37 @@ const safeJSONParse = (initialValue: unknown) => (str: string) => {
   }
 };
 
+export type SetHashOption =
+  | 'default'
+  | 'replaceState'
+  | ((searchParams: string) => void);
+
+export type AtomWithHashSetOptions = {
+  setHash?: SetHashOption;
+};
+
+export const setHashWithPush = (searchParams: string) => {
+  window.location.hash = searchParams;
+};
+
+export const setHashWithReplace = (searchParams: string): void => {
+  window.history.replaceState(
+    window.history.state,
+    '',
+    `${window.location.pathname}${window.location.search}#${searchParams}`,
+  );
+};
+
+function getSetHashFn(setHashOption?: SetHashOption) {
+  if (setHashOption === 'replaceState') {
+    return setHashWithReplace;
+  }
+  if (typeof setHashOption === 'function') {
+    return setHashOption;
+  }
+  return setHashWithPush;
+}
+
 export function atomWithHash<Value>(
   key: string,
   initialValue: Value,
@@ -22,9 +53,14 @@ export function atomWithHash<Value>(
     serialize?: (val: Value) => string;
     deserialize?: (str: string) => Value;
     subscribe?: (callback: () => void) => () => void;
-    setHash?: 'default' | 'replaceState' | ((searchParams: string) => void);
+    setHash?: SetHashOption;
   },
-): WritableAtom<Value, [SetStateActionWithReset<Value>], void> {
+): WritableAtom<
+  Value,
+  | [SetStateActionWithReset<Value>]
+  | [SetStateActionWithReset<Value>, AtomWithHashSetOptions],
+  void
+> {
   const serialize = options?.serialize || JSON.stringify;
 
   const deserialize = options?.deserialize || safeJSONParse(initialValue);
@@ -36,22 +72,7 @@ export function atomWithHash<Value>(
         window.removeEventListener('hashchange', callback);
       };
     });
-  const setHashOption = options?.setHash;
-  let setHash = (searchParams: string) => {
-    window.location.hash = searchParams;
-  };
-  if (setHashOption === 'replaceState') {
-    setHash = (searchParams) => {
-      window.history.replaceState(
-        window.history.state,
-        '',
-        `${window.location.pathname}${window.location.search}#${searchParams}`,
-      );
-    };
-  }
-  if (typeof setHashOption === 'function') {
-    setHash = setHashOption;
-  }
+
   const isLocationAvailable =
     typeof window !== 'undefined' && !!window.location;
 
@@ -77,7 +98,12 @@ export function atomWithHash<Value>(
   });
   return atom(
     (get) => get(valueAtom),
-    (get, set, update: SetStateActionWithReset<Value>) => {
+    (
+      get,
+      set,
+      update: SetStateActionWithReset<Value>,
+      setOptions?: AtomWithHashSetOptions,
+    ) => {
       const nextValue =
         typeof update === 'function'
           ? (update as (prev: Value) => Value | typeof RESET)(get(valueAtom))
@@ -91,6 +117,7 @@ export function atomWithHash<Value>(
         set(strAtom, str);
         searchParams.set(key, str);
       }
+      const setHash = getSetHashFn(setOptions?.setHash ?? options?.setHash);
       setHash(searchParams.toString());
     },
   );