@macrostrat/column-views 2.3.3 → 2.4.0

package/src/data-provider/core.ts

@@ -0,0 +1,89 @@
+import { atom, type WritableAtom } from "jotai";
+import { createIsolation } from "jotai-scope";
+import { BaseUnit } from "@macrostrat/api-types";
+import { ReactNode, useEffect, useRef } from "react";
+import h from "@macrostrat/hyper";
+
+interface StateIsolation {
+  Provider: (props: {
+    store?: any;
+    initialValues?: AtomMap;
+    children: ReactNode;
+  }) => ReactNode;
+  useStore: () => any;
+  useAtom: any;
+  useAtomValue: any;
+  useSetAtom: any;
+}
+
+export const scope: StateIsolation = createIsolation();
+
+export type AtomMap = [WritableAtom<any, any, any>, any][];
+
+type ProviderProps = {
+  children: ReactNode;
+  atoms?: AtomMap;
+  shouldUpdateAtoms?: boolean;
+};
+
+export function ScopedProvider({
+  children,
+  atoms,
+  shouldUpdateAtoms = true,
+}: ProviderProps) {
+  // Always use the same store instance in this tree
+  let val = null;
+  try {
+    val = scope.useStore();
+    // This store has already been provided
+    return children;
+  } catch {
+    return h(scope.Provider, { store: null, initialValues: atoms }, [
+      h.if(shouldUpdateAtoms)(AtomUpdater, { atoms }),
+      children,
+    ]);
+  }
+}
+
+function AtomUpdater({
+  atoms,
+}: {
+  atoms: [WritableAtom<any, any, any>, any][];
+}) {
+  /**
+   * A generic updater to sync Jotai atoms with state passed as props.
+   * Useful for scoped providers where state needs to be synced outside
+   * the current context.
+   */
+  /** TODO: this is an awkward way to keep atoms updated */
+  // The scoped store
+  const store = scope.useStore();
+
+  // Warn on atoms changing length
+  const prevLengthRef = useRef(atoms.length);
+  useEffect(() => {
+    if (prevLengthRef.current !== atoms.length) {
+      console.error("Error: number of atoms in ScopedAtomUpdater has changed.");
+      prevLengthRef.current = atoms.length;
+    }
+  }, [atoms.length]);
+
+  for (const [atom, value] of atoms) {
+    useEffect(() => {
+      store.set(atom, value);
+    }, [store, value]);
+  }
+  return null;
+}
+
+export const columnUnitsAtom = atom<BaseUnit[]>();
+
+export const columnUnitsMapAtom = atom<Map<number, BaseUnit> | null>((get) => {
+  const units = get(columnUnitsAtom);
+  if (!units) return null;
+  const unitMap = new Map<number, BaseUnit>();
+  units.forEach((unit) => {
+    unitMap.set(unit.unit_id, unit);
+  });
+  return unitMap;
+});

package/src/data-provider/store.ts

@@ -1,11 +1,4 @@
-import {
-  createContext,
-  ReactNode,
-  useContext,
-  useEffect,
-  useMemo,
-  useRef,
-} from "react";
+import { createContext, ReactNode, useContext, useMemo } from "react";
 import h from "@macrostrat/hyper";
 import {
   CompositeColumnScale,
@@ -13,84 +6,98 @@ import {
 } from "../prepare-units/composite-scale";
 import { ColumnAxisType } from "@macrostrat/column-components";
 import type { ExtUnit, PackageLayoutData } from "../prepare-units";
-// An isolated jotai store for Macrostrat column usage
-// TODO: there might be a better way to do this using the MacrostratDataProvider or similar
-import { createIsolation } from "jotai-scope";
-import { atom, PrimitiveAtom, type WritableAtom } from "jotai";
-
-const { Provider, useSetAtom, useAtomValue, useStore } = createIsolation();
-
-type ProviderProps = {
+import {
+  allowUnitSelectionAtom,
+  selectedUnitIDAtom,
+  UnitSelectionCallbacks,
+  UnitSelectionCallbackManager,
+} from "./unit-selection";
+import { BaseUnit } from "@macrostrat/api-types";
+import {
+  AtomMap,
+  columnUnitsAtom,
+  columnUnitsMapAtom,
+  scope,
+  ScopedProvider,
+} from "./core";
+
+export interface ColumnStateProviderProps<
+  T extends BaseUnit,
+> extends Partial<UnitSelectionCallbacks> {
   children: ReactNode;
-  initialValues?: Iterable<[WritableAtom<any, any, any>, any]>;
-};
-
-const columnUnitsAtom = atom<ExtUnit[]>();
-
-const columnUnitsMapAtom = atom<Map<number, ExtUnit> | null>((get) => {
-  const units = get(columnUnitsAtom);
-  if (!units) return null;
-  const unitMap = new Map<number, ExtUnit>();
-  units.forEach((unit) => {
-    unitMap.set(unit.unit_id, unit);
-  });
-  return unitMap;
-});
-
-function ScopedProvider({ children, ...rest }: ProviderProps) {
-  // Always use the same store instance in this tree
-  let val = null;
-  try {
-    val = useStore();
-  } catch {
-    // No store found, create a new one
-    val = null;
-  }
-  return h(Provider, { store: val, ...rest }, children);
+  units: T[];
+  selectedUnit: number | null;
+  allowUnitSelection?: boolean;
 }
 
-export function MacrostratColumnStateProvider({
+export function MacrostratColumnStateProvider<T extends BaseUnit>({
   children,
   units,
-}: {
-  children?: ReactNode;
-  units: ExtUnit[];
-}) {
+  allowUnitSelection = false,
+  onUnitSelected,
+  selectedUnit,
+}: ColumnStateProviderProps<T>) {
   /** Top-level provider for Macrostrat column data.
    * It is either provided by the Column component itself, or
    * can be hoisted higher in the tree to provide a common data context
    */
 
-  const atomMap: [[PrimitiveAtom<ExtUnit[]>, ExtUnit[]]] = [
+  /* By default, unit selection is disabled. However, if any related props are passed,
+  we enable it.
+  */
+  const _allowSelection = useMemo(() => {
+    if (allowUnitSelection) {
+      return true;
+    }
+    return selectedUnit != null || onUnitSelected != null;
+  }, []);
+
+  const atomMap: AtomMap = [
     [columnUnitsAtom, units],
+    [allowUnitSelectionAtom, _allowSelection],
+    [selectedUnitIDAtom, selectedUnit],
   ];
 
+  let selectionHandlers: ReactNode = null;
+  if (_allowSelection) {
+    selectionHandlers = h(UnitSelectionCallbackManager, { onUnitSelected });
+  }
+
   return h(
     ScopedProvider,
     {
-      initialValues: atomMap,
+      atoms: atomMap,
     },
-    [h(AtomUpdater, { atoms: atomMap }), children],
+    [selectionHandlers, children],
   );
 }
 
-export interface MacrostratColumnDataContext {
-  units: ExtUnit[];
+export interface MacrostratColumnDataContext<T extends BaseUnit> {
+  units: T[];
   sections: PackageLayoutData[];
   totalHeight?: number;
   axisType?: ColumnAxisType;
+  allowUnitSelection?: boolean;
+}
+
+export interface ColumnDataProviderProps<T extends BaseUnit>
+  extends MacrostratColumnDataContext<T>, ColumnStateProviderProps<T> {
+  children: ReactNode;
 }
 
 const MacrostratColumnDataContext =
-  createContext<MacrostratColumnDataContext>(null);
+  createContext<MacrostratColumnDataContext<any>>(null);
 
-export function MacrostratColumnDataProvider({
+export function MacrostratColumnDataProvider<T extends BaseUnit>({
   children,
   units,
   sections,
   totalHeight,
   axisType,
-}: MacrostratColumnDataContext & { children: ReactNode }) {
+  allowUnitSelection,
+  onUnitSelected,
+  selectedUnit,
+}: ColumnDataProviderProps<T>) {
   /** Internal provider for Macrostrat column data.
    * As a general rule, we want to provide data and column-axis
    * height calculations through the context, since these need to
@@ -114,11 +121,26 @@ export function MacrostratColumnDataProvider({
 
   return h(
     MacrostratColumnStateProvider,
-    { units },
+    {
+      units,
+      allowUnitSelection,
+      onUnitSelected,
+      selectedUnit,
+    },
     h(MacrostratColumnDataContext.Provider, { value }, children),
   );
 }
 
+export function UnitSelectionProvider(props: ColumnStateProviderProps<any>) {
+  /* A basic unit selection provider without column data context.
+  Mostly for use with the CorrelationChart component.
+   */
+  return h(MacrostratColumnStateProvider, {
+    ...props,
+    allowUnitSelection: true,
+  });
+}
+
 export function useMacrostratColumnData() {
   const ctx = useContext(MacrostratColumnDataContext);
   if (!ctx) {
@@ -130,12 +152,12 @@ export function useMacrostratColumnData() {
 }
 
 export function useMacrostratUnits() {
-  return useAtomValue(columnUnitsAtom);
+  return scope.useAtomValue(columnUnitsAtom);
 }
 
 export function useColumnUnitsMap(): Map<number, ExtUnit> | null {
   try {
-    return useAtomValue(columnUnitsMapAtom);
+    return scope.useAtomValue(columnUnitsMapAtom) as Map<number, ExtUnit>;
   } catch {
     return null;
   }
@@ -148,34 +170,3 @@ export function useCompositeScale(): CompositeColumnScale {
     [ctx.sections],
   );
 }
-
-function AtomUpdater({
-  atoms,
-}: {
-  atoms: [WritableAtom<any, any, any>, any][];
-}) {
-  /**
-   * A generic updater to sync Jotai atoms with state passed as props.
-   * Useful for scoped providers where state needs to be synced outside
-   * of the current context.
-   */
-  /** TODO: this is an awkward way to keep atoms updated */
-  // The scoped store
-  const store = useStore();
-
-  // Warn on atoms changing length
-  const prevLengthRef = useRef(atoms.length);
-  useEffect(() => {
-    if (prevLengthRef.current !== atoms.length) {
-      console.error("Error: number of atoms in ScopedAtomUpdater has changed.");
-      prevLengthRef.current = atoms.length;
-    }
-  }, [atoms.length]);
-
-  for (const [atom, value] of atoms) {
-    useEffect(() => {
-      store.set(atom, value);
-    }, [store, value]);
-  }
-  return null;
-}

package/src/data-provider/unit-selection.ts

@@ -1,55 +1,21 @@
 import { BaseUnit } from "@macrostrat/api-types";
-import h from "@macrostrat/hyper";
 import { useKeyHandler } from "@macrostrat/ui-components";
-import {
-  createContext,
-  useContext,
-  useEffect,
-  useState,
-  ReactNode,
-  RefObject,
-  useRef,
-  useCallback,
-} from "react";
-import { createStore, StoreApi, useStore } from "zustand";
+import { useEffect, useRef, useCallback } from "react";
 import type { RectBounds, IUnit } from "../units/types";
+import { atom } from "jotai";
+import { columnUnitsMapAtom, scope } from "./core";
 
 type UnitSelectDispatch = (
   unit: number | BaseUnit | null,
   target: HTMLElement | null,
-  event: Event | null,
 ) => void;
 
-const UnitSelectionContext = createContext<StoreApi<UnitSelectionStore>>(null);
-
 export function useUnitSelectionDispatch(): UnitSelectDispatch {
-  const store = useContext(UnitSelectionContext);
-  if (store == null) {
-    return () => {};
-  }
-  return useStore(store, (state) => state.onUnitSelected);
-}
-
-export function useUnitSelectionStore<T>(
-  selector: (state: UnitSelectionStore) => T,
-): T {
-  const store = useContext(UnitSelectionContext);
-  if (store == null) {
-    return null;
-  }
-  return useStore(store, selector);
+  return scope.useSetAtom(selectedUnitAtom);
 }
 
 export function useSelectedUnit() {
-  return useUnitSelectionStore((state) => state.selectedUnitData);
-}
-
-interface UnitSelectionStore {
-  selectedUnit: number | null;
-  selectedUnitData: BaseUnit | null;
-  overlayPosition: RectBounds | null;
-  onUnitSelected: UnitSelectDispatch;
-  setSelectedUnit: (unit: number | null) => void;
+  return scope.useAtomValue(selectedUnitAtom);
 }
 
 export interface ColumnClickData {
@@ -60,134 +26,141 @@ export interface ColumnClickData {
   // Room for boundary IDs eventually
 }
 
-export function UnitSelectionProvider<T extends BaseUnit>(props: {
-  children: ReactNode;
-  columnRef?: RefObject<HTMLElement>;
-  units: T[];
-  selectedUnit: number | null;
+export interface UnitSelectionCallbacks {
+  // It's sort of unfortunate that we need to pass in the column ref here
   onClickedColumn?: (columnClickData: ColumnClickData, event: Event) => void;
-  onUnitSelected?: (unitID: number | null, unit: T | null) => void;
-}) {
-  const [store] = useState(() =>
-    createStore<UnitSelectionStore>((set) => ({
-      selectedUnit: null,
-      selectedUnitData: null,
-      overlayPosition: null,
-      setSelectedUnit(selectedUnit: number | null | BaseUnit) {
-        if (selectedUnit == null) {
-          set({ selectedUnit: null, selectedUnitData: null });
-          return;
-          // If it's a number, set the selected unit
-        } else if (typeof selectedUnit === "number") {
-          set({ selectedUnit, selectedUnitData: null });
-        } else if ("unit_id" in selectedUnit) {
-          set({
-            selectedUnit: selectedUnit.unit_id,
-            selectedUnitData: selectedUnit,
-          });
-        }
-      },
-      onUnitSelected: (
-        input: number | T | null,
-        target: HTMLElement = null,
-        event: PointerEvent = null,
-      ) => {
-        if (input == null) {
-          props.onUnitSelected?.(null, null);
-          return set({
-            selectedUnit: null,
-            selectedUnitData: null,
-            overlayPosition: null,
-          });
-        }
-
-        let unit: T | null = null;
-        if (typeof input === "number") {
-          unit = props.units.find((u) => u.unit_id === input) || null;
-        } else {
-          unit = input;
-        }
-
-        const el = props.columnRef?.current;
-        let overlayPosition = null;
-
-        /** This is not the natural place to get positions within the column,
-         * but it will work for now.
-         */
-        if (props.onClickedColumn) {
-          // Infer height from top and bottom height of unit (because that's passed back with the call)
-          //const py = event.y;
-          //const [bottom, top] = getUnitHeightRange(unit, axisType);
-          //const height = Math.abs(bottom - top);
-
-          /** Ideally this would be defined at the column level */
-          const columnClickData: ColumnClickData = {
-            unitID: unit?.unit_id,
-            unit,
-            target,
-            height: el?.getBoundingClientRect().height || 0,
-          };
-          props.onClickedColumn(columnClickData, event);
-        }
-
-        if (unit != null && el != null && target != null) {
-          const rect = el.getBoundingClientRect();
-          const targetRect = target.getBoundingClientRect();
-
-          overlayPosition = {
-            x: targetRect.left - rect.left,
-            y: targetRect.top - rect.top,
-            width: targetRect.width,
-            height: targetRect.height,
-          };
-        }
-        props.onUnitSelected?.(unit?.unit_id, unit);
-
-        return set({
-          selectedUnit: unit?.unit_id,
-          selectedUnitData: unit,
-          overlayPosition,
-        });
-      },
-    })),
-  );
+  onUnitSelected?: <T extends BaseUnit>(
+    unitID: number | null,
+    unit: T | null,
+  ) => void;
+}
 
-  const { units, selectedUnit } = props;
+export const allowUnitSelectionAtom = atom<boolean>(true);
 
-  useEffect(() => {
-    // Synchronize store with provided props
-    if (selectedUnit != null) {
-      const unitData = units?.find((u) => u.unit_id === selectedUnit);
-      store.setState({ selectedUnit, selectedUnitData: unitData });
-    } else {
-      store.setState({ selectedUnit: null, selectedUnitData: null });
+export const selectedUnitIDAtom = atom<number | null>();
+
+const overlayPositionAtom = atom<RectBounds | null>();
+
+const columnRefAtom = atom<{ current: HTMLElement | null }>({ current: null });
+
+export function useColumnRef() {
+  return scope.useAtomValue(columnRefAtom);
+}
+
+const selectedUnitAtom = atom(
+  (get) => {
+    if (!get(allowUnitSelectionAtom)) return null;
+
+    const unitID = get(selectedUnitIDAtom);
+    if (unitID == null) return null;
+    const unitsMap = get(columnUnitsMapAtom);
+    return unitsMap?.get(unitID) || null;
+  },
+  (
+    get,
+    set,
+    selectedUnit: number | BaseUnit | null,
+    target: HTMLElement | null = null,
+  ): BaseUnit | null => {
+    if (!get(allowUnitSelectionAtom)) {
+      throw new Error("Unit selection is disabled.");
+    }
+
+    let unitID: number | null;
+    let unit: BaseUnit | null = null;
+    if (selectedUnit == null) {
+      unitID = null;
+    } else if (typeof selectedUnit === "number") {
+      unitID = selectedUnit;
+    } else if ("unit_id" in selectedUnit) {
+      unitID = selectedUnit.unit_id;
+    }
+
+    if (unitID != null) {
+      // Verify that the unit exists in the current colum, else throw
+      const unitsMap = get(columnUnitsMapAtom);
+      if (!unitsMap?.has(unitID)) {
+        throw new Error(
+          `Unit with ID ${unitID} not found in current column units.`,
+        );
+      }
+      unit = unitsMap.get(unitID) ?? null;
+    }
+
+    let overlayPosition: RectBounds | null = null;
+
+    const columnEl = get(columnRefAtom)?.current;
+
+    if (unit != null && columnEl != null && target != null) {
+      const rect = columnEl.getBoundingClientRect();
+      const targetRect = target.getBoundingClientRect();
+      overlayPosition = {
+        x: targetRect.left - rect.left,
+        y: targetRect.top - rect.top,
+        width: targetRect.width,
+        height: targetRect.height,
+      };
     }
-  }, [selectedUnit, units]);
 
-  return h(UnitSelectionContext.Provider, { value: store }, props.children);
+    set(selectedUnitIDAtom, unitID);
+    set(overlayPositionAtom, overlayPosition);
+
+    return unit;
+  },
+);
+
+export function useAtomOverlayPosition() {
+  return scope.useAtomValue(overlayPositionAtom);
+}
+
+export function UnitSelectionCallbackManager({
+  onUnitSelected,
+}: UnitSelectionCallbacks) {
+  const selectedUnit = scope.useAtomValue(selectedUnitAtom);
+  useEffect(() => {
+    onUnitSelected?.(selectedUnit?.unit_id ?? null, selectedUnit ?? null);
+  }, [selectedUnit?.unit_id]);
+  return null;
 }
 
+/** TODO: figure out if onClickedColumn is still needed
+ /** This is not the natural place to get positions within the column,
+ * but it will work for now.
+if (props.onClickedColumn) {
+  // Infer height from top and bottom height of unit (because that's passed back with the call)
+  //const py = event.y;
+  //const [bottom, top] = getUnitHeightRange(unit, axisType);
+  //const height = Math.abs(bottom - top);
+
+  const columnClickData: ColumnClickData = {
+    unitID: unit?.unit_id,
+    unit,
+    target,
+    height: el?.getBoundingClientRect().height || 0,
+  };
+  props.onClickedColumn(columnClickData, event);
+}*/
+
 export function useUnitSelectionTarget(
   unit: IUnit,
 ): [React.RefObject<HTMLElement>, boolean, (evt: Event) => void] {
   const ref = useRef<HTMLElement>();
   const selectedUnit = useSelectedUnit();
-  const onSelectUnit = useUnitSelectionStore((state) => state.onUnitSelected);
+  const selectUnit = useUnitSelectionDispatch();
   const selected = selectedUnit?.unit_id == unit.unit_id;
 
   const onClick = useCallback(
     (evt: Event) => {
-      if (onSelectUnit == null) return;
-      onSelectUnit(unit, ref.current, evt);
+      selectUnit?.(unit, ref.current);
       evt.stopPropagation();
     },
-    [unit, onSelectUnit],
+    [unit, selectUnit],
   );
 
   useEffect(() => {
-    if (!selected || onSelectUnit == null) return;
+    if (!selected || selectUnit == null) return;
     // In case we haven't set the position of the unit (if we don't have a target), set the selected unit
-    onSelectUnit(unit, ref.current, null);
+    selectUnit(unit, ref.current);
 
     // Scroll the unit into view
     ref.current?.scrollIntoView({
@@ -195,7 +168,7 @@ export function useUnitSelectionTarget(
       block: "center",
       inline: "nearest",
     });
-  }, [selected, onSelectUnit]);
+  }, [selected, selectUnit]);
 
   return [ref, selected, onClick];
 }
@@ -219,7 +192,7 @@ export function UnitKeyboardNavigation<T extends BaseUnit>({
     (event) => {
       const nextIx = keyMap[event.keyCode];
       if (nextIx == null || nextIx < 0 || nextIx >= units.length) return;
-      selectUnit(units[nextIx], null, null);
+      selectUnit(units[nextIx], null);
       event.stopPropagation();
     },
     [units, ix],

package/src/unit-details/popover.ts

@@ -1,7 +1,11 @@
 import hyper from "@macrostrat/hyper";
 import { Popover } from "@blueprintjs/core";
 import styles from "./popover.module.sass";
-import { useSelectedUnit, useUnitSelectionStore } from "../data-provider";
+import {
+  useAtomOverlayPosition,
+  useSelectedUnit,
+  useUnitSelectionDispatch,
+} from "../data-provider";
 import { UnitDetailsPanel } from "./panel";
 
 const h = hyper.styled(styles);
@@ -54,8 +58,8 @@ function InteractionBarrier({ children }) {
 
 export function UnitSelectionPopover() {
   const unit = useSelectedUnit();
-  const selectUnit = useUnitSelectionStore((state) => state.onUnitSelected);
-  const position = useUnitSelectionStore((state) => state.overlayPosition);
+  const selectUnit = useUnitSelectionDispatch();
+  const position = useAtomOverlayPosition();
   if (unit == null) {
     return null;
   }
@@ -79,7 +83,7 @@ export function UnitSelectionPopover() {
         className: "legend-panel",
         onSelectUnit: (id: number) => {
           console.log("Selected unit in popover:", id);
-          selectUnit(id, null, null);
+          selectUnit(id, null);
         },
       }),
     ),