Move viewState into a hook

src/components/Viewer.tsx

@@ -1,12 +1,11 @@
 import DeckGL from "deck.gl";
 import { type Layer, OrthographicView } from "deck.gl";
-import { type WritableAtom, useAtom } from "jotai";
 import { useAtomValue } from "jotai";
 import * as React from "react";
 
+import { useViewState } from "@/hooks";
 import type { LayerProps } from "@deck.gl/core/lib/layer";
 import type { ZarrPixelSource } from "../ZarrPixelSource";
-import type { ViewState } from "../state";
 import { layerAtoms } from "../state";
 import { fitBounds, isInterleaved } from "../utils";
 
@@ -28,11 +27,8 @@ function getLayerSize(props: Data) {
   return { height, width, maxZoom };
 }
 
-function WrappedViewStateDeck(props: {
-  layers: Array<VizarrLayer | null>;
-  viewStateAtom: WritableAtom<ViewState | undefined, ViewState>;
-}) {
-  const [viewState, setViewState] = useAtom(props.viewStateAtom);
+function WrappedViewStateDeck(props: { layers: Array<VizarrLayer | null> }) {
+  const [viewState, setViewState] = useViewState();
   const deckRef = React.useRef<DeckGL>(null);
   const firstLayerProps = props.layers[0]?.props;
 
@@ -46,30 +42,28 @@ function WrappedViewStateDeck(props: {
     setViewState(bounds);
   }
 
-  // Enables screenshots of the canvas: https://github.com/visgl/deck.gl/issues/2200
-  const glOptions: WebGLContextAttributes = {
-    preserveDrawingBuffer: true,
-  };
-
   return (
     <DeckGL
       ref={deckRef}
       layers={props.layers}
       viewState={viewState}
       onViewStateChange={(e) => setViewState(e.viewState)}
       views={[new OrthographicView({ id: "ortho", controller: true })]}
-      glOptions={glOptions}
+      glOptions={{
+        // Enables screenshots of the canvas: https://github.com/visgl/deck.gl/issues/2200
+        preserveDrawingBuffer: true,
+      }}
     />
   );
 }
 
-function Viewer({ viewStateAtom }: { viewStateAtom: WritableAtom<ViewState | undefined, ViewState> }) {
+function Viewer() {
   const layerConstructors = useAtomValue(layerAtoms);
   // @ts-expect-error - Viv types are giving up an issue
   const layers: Array<VizarrLayer | null> = layerConstructors.map((layer) => {
     return !layer.on ? null : new layer.Layer(layer.layerProps);
   });
-  return <WrappedViewStateDeck viewStateAtom={viewStateAtom} layers={layers} />;
+  return <WrappedViewStateDeck layers={layers} />;
 }
 
 export default Viewer;

src/hooks.ts

@@ -1,13 +1,15 @@
 import { type PrimitiveAtom, useAtom, useAtomValue } from "jotai";
 import * as React from "react";
 
-import { type SourceData, layerFamilyAtom } from "./state";
+import { type SourceData, type ViewState, layerFamilyAtom } from "./state";
 import { assert } from "./utils";
 
 type WithId<T> = { id: string } & T;
 type SourceDataAtom = PrimitiveAtom<WithId<SourceData>>;
+type ViewStateAtom = PrimitiveAtom<ViewState | undefined>;
 
 export const LayerContext = React.createContext<null | SourceDataAtom>(null);
+export const ViewStateContext = React.createContext<null | ViewStateAtom>(null);
 
 function useSourceAtom(): SourceDataAtom {
   const sourceAtom = React.useContext(LayerContext);
@@ -40,3 +42,9 @@ export function useLayerValue() {
   const layerAtom = useLayerAtom();
   return useAtomValue(layerAtom);
 }
+
+export function useViewState() {
+  const viewStateAtom = React.useContext(ViewStateContext);
+  assert(viewStateAtom !== null, "useViewState must be used within a ViewStateContext.Provider");
+  return useAtom(viewStateAtom);
+}

src/index.tsx

@@ -6,12 +6,13 @@ import ReactDOM from "react-dom/client";
 import Menu from "./components/Menu";
 import Viewer from "./components/Viewer";
 import "./codecs/register";
-import { type ImageLayerConfig, type ViewState, addImageAtom, atomWithEffect } from "./state";
+import { type ImageLayerConfig, type ViewState, addImageAtom } from "./state";
 import { defer, typedEmitter } from "./utils";
 
 export { version } from "../package.json";
 
 import "./index.css";
+import { ViewStateContext } from "./hooks";
 
 type Events = {
   viewStateChange: ViewState;
@@ -29,9 +30,14 @@ export interface VizarrViewer {
 export function createViewer(element: HTMLElement, options: { menuOpen?: boolean } = {}): Promise<VizarrViewer> {
   const ref = React.createRef<VizarrViewer>();
   const emitter = typedEmitter<Events>();
-  const viewStateAtom = atomWithEffect<ViewState | undefined, ViewState>(
-    atom<ViewState | undefined>(undefined),
-    ({ zoom, target }) => emitter.emit("viewStateChange", { zoom, target }),
+  const viewStateAtom = atom<ViewState | undefined>(undefined);
+  const viewStateAtomWithUrlSync = atom(
+    (get) => get(viewStateAtom),
+    (get, set, update) => {
+      const next = typeof update === "function" ? update(get(viewStateAtom)) : update;
+      next && emitter.emit("viewStateChange", { target: next.target, zoom: next.zoom });
+      set(viewStateAtom, next);
+    },
   );
   const { promise, resolve } = defer<VizarrViewer>();
 
@@ -54,10 +60,10 @@ export function createViewer(element: HTMLElement, options: { menuOpen?: boolean
       }
     }, []);
     return (
-      <>
+      <ViewStateContext.Provider value={viewStateAtomWithUrlSync}>
         <Menu open={options.menuOpen ?? true} />
-        <Viewer viewStateAtom={viewStateAtom} />
-      </>
+        <Viewer />
+      </ViewStateContext.Provider>
     );
   }
   let root = ReactDOM.createRoot(element);

src/state.ts

@@ -1,5 +1,4 @@
 import type { ImageLayer, MultiscaleImageLayer } from "@hms-dbmi/viv";
-import type { PrimitiveAtom, WritableAtom } from "jotai";
 import { atom } from "jotai";
 import { atomFamily, splitAtom, waitForAll } from "jotai/utils";
 import type { Matrix4 } from "math.gl";
@@ -14,22 +13,6 @@ export interface ViewState {
   target: [number, number];
 }
 
-export function atomWithEffect<Value, Update extends object, Result extends void | Promise<void> = void>(
-  baseAtom: WritableAtom<Value, Update | ((prev: Value) => Update), Result>,
-  callback: (data: Update) => void,
-) {
-  const derivedAtom: typeof baseAtom = atom(
-    (get) => get(baseAtom),
-    (get, set, update) => {
-      const next = typeof update === "function" ? update(get(baseAtom)) : update;
-      const result = set(baseAtom, next);
-      callback(next);
-      return result;
-    },
-  );
-  return derivedAtom;
-}
-
 interface BaseConfig {
   source: string | Readable;
   axis_labels?: string[];