WIP persisting client state across reloads

Author: mattrunyon

plugins/ui/src/js/src/elements/UITable/UITable.tsx

@@ -1,4 +1,10 @@
-import React, { useCallback, useEffect, useMemo, useState } from 'react';
+import React, {
+  useCallback,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+} from 'react';
 import { useSelector } from 'react-redux';
 import classNames from 'classnames';
 import {
@@ -8,6 +14,8 @@ import {
   type IrisGridContextMenuData,
   IrisGridProps,
   IrisGridUtils,
+  IrisGridState,
+  DehydratedIrisGridState,
 } from '@deephaven/iris-grid';
 import {
   ColorValues,
@@ -37,6 +45,7 @@ import UITableContextMenuHandler, {
 } from './UITableContextMenuHandler';
 import UITableModel, { makeUiTableModel } from './UITableModel';
 import { UITableLayoutHints } from './JsTableProxy';
+import usePersistentState from '../hooks/usePersistentState';
 
 const log = Log.module('@deephaven/js-plugin-ui/UITable');
 
@@ -167,6 +176,8 @@ export function UITable({
   databars = EMPTY_ARRAY as unknown as DatabarConfig[],
   ...userStyleProps
 }: UITableProps): JSX.Element | null {
+  const id = useMemo(() => Math.random().toString(36).substr(2, 9), []);
+  console.log('render table', id);
   const [throwError] = useThrowError();
 
   // Margin looks wrong with ui.table, so we want to map margin to padding instead
@@ -275,6 +286,43 @@ export function UITable({
     model.setColorMap(colorMap);
   }
 
+  const prevState = useRef<IrisGridState | undefined>();
+  const [dehydratedState, setDehydratedState] = usePersistentState<
+    DehydratedIrisGridState | undefined
+  >(undefined);
+  const initialState = useRef(dehydratedState);
+
+  const onStateChange = useCallback(
+    (newState: IrisGridState) => {
+      if (
+        model == null ||
+        newState.metrics == null ||
+        newState === prevState.current ||
+        newState.sorts === prevState.current?.sorts
+      ) {
+        return;
+      }
+      // console.log('onStateChange', newState);
+      prevState.current = newState;
+      setDehydratedState(
+        utils.dehydrateIrisGridState(model, {
+          ...newState,
+          metrics: {
+            userColumnWidths: newState.metrics.userColumnWidths,
+            userRowHeights: newState.metrics.userRowHeights,
+          },
+        })
+      );
+    },
+    [model, setDehydratedState, utils]
+  );
+
+  const hydratedState = useMemo(() => {
+    if (model && initialState.current) {
+      return utils.hydrateIrisGridState(model, initialState.current);
+    }
+  }, [model, utils]);
+
   const hydratedSorts = useMemo(() => {
     if (sorts !== undefined && columns !== undefined) {
       log.debug('Hydrating sorts', sorts);
@@ -401,59 +449,66 @@ export function UITable({
     [contextMenu, alwaysFetchColumns]
   );
 
-  const irisGridProps = useMemo(
-    () =>
-      ({
-        mouseHandlers,
-        alwaysFetchColumns,
-        showSearchBar,
-        sorts: hydratedSorts,
-        quickFilters: hydratedQuickFilters,
-        isFilterBarShown: showQuickFilters,
-        reverseType: reverse
-          ? TableUtils.REVERSE_TYPE.POST_SORT
-          : TableUtils.REVERSE_TYPE.NONE,
-        density,
-        settings: { ...settings, showExtraGroupColumn: showGroupingColumn },
-        onContextMenu,
-        aggregationSettings: {
-          aggregations:
-            aggregations != null
-              ? ensureArray(aggregations).map(agg => {
-                  if (agg.cols != null && agg.ignore_cols != null) {
-                    throw new Error(
-                      'Cannot specify both cols and ignore_cols in a UI table aggregation'
-                    );
-                  }
-                  return {
-                    operation: getAggregationOperation(agg.agg),
-                    selected: ensureArray(agg.cols ?? agg.ignore_cols ?? []),
-                    // If agg.cols is set, we don't want to invert
-                    // If it is not set, then the only other options are ignore_cols or neither
-                    // In both cases, we want to invert since we are either ignoring, or selecting all as [] inverted
-                    invert: agg.cols == null,
-                  };
-                })
-              : [],
-          showOnTop: aggregationsPosition === 'top',
-        },
-      }) satisfies Partial<IrisGridProps>,
-    [
+  const irisGridProps = useMemo(() => {
+    const props = {
       mouseHandlers,
       alwaysFetchColumns,
       showSearchBar,
-      showQuickFilters,
-      hydratedSorts,
-      hydratedQuickFilters,
-      reverse,
+      sorts: hydratedSorts,
+      quickFilters: hydratedQuickFilters,
+      isFilterBarShown: showQuickFilters,
+      reverseType: reverse
+        ? TableUtils.REVERSE_TYPE.POST_SORT
+        : TableUtils.REVERSE_TYPE.NONE,
       density,
-      settings,
-      showGroupingColumn,
+      settings: { ...settings, showExtraGroupColumn: showGroupingColumn },
       onContextMenu,
-      aggregations,
-      aggregationsPosition,
-    ]
-  );
+      aggregationSettings: {
+        aggregations:
+          aggregations != null
+            ? ensureArray(aggregations).map(agg => {
+                if (agg.cols != null && agg.ignore_cols != null) {
+                  throw new Error(
+                    'Cannot specify both cols and ignore_cols in a UI table aggregation'
+                  );
+                }
+                return {
+                  operation: getAggregationOperation(agg.agg),
+                  selected: ensureArray(agg.cols ?? agg.ignore_cols ?? []),
+                  // If agg.cols is set, we don't want to invert
+                  // If it is not set, then the only other options are ignore_cols or neither
+                  // In both cases, we want to invert since we are either ignoring, or selecting all as [] inverted
+                  invert: agg.cols == null,
+                };
+              })
+            : [],
+        showOnTop: aggregationsPosition === 'top',
+      },
+    } satisfies Partial<IrisGridProps>;
+
+    // Remove any explicit undefined values
+    Object.entries(props).forEach(([key, value]) => {
+      if (value === undefined) {
+        delete props[key];
+      }
+    });
+
+    return props;
+  }, [
+    mouseHandlers,
+    alwaysFetchColumns,
+    showSearchBar,
+    showQuickFilters,
+    hydratedSorts,
+    hydratedQuickFilters,
+    reverse,
+    density,
+    settings,
+    showGroupingColumn,
+    onContextMenu,
+    aggregations,
+    aggregationsPosition,
+  ]);
 
   return model ? (
     <div
@@ -464,6 +519,9 @@ export function UITable({
       <IrisGrid
         ref={ref => setIrisGrid(ref)}
         model={model}
+        onStateChange={onStateChange}
+        // eslint-disable-next-line react/jsx-props-no-spreading
+        {...hydratedState}
         // eslint-disable-next-line react/jsx-props-no-spreading
         {...irisGridProps}
       />

plugins/ui/src/js/src/elements/hooks/usePersistentState.ts

@@ -0,0 +1,33 @@
+import {
+  useCallback,
+  useContext,
+  useState,
+  type Dispatch,
+  type SetStateAction,
+} from 'react';
+import { ReactPanelContext } from '../../layout/ReactPanelContext';
+
+export default function usePersistentState<S>(
+  initialState: S | (() => S)
+): [S, Dispatch<SetStateAction<S>>] {
+  const context = useContext(ReactPanelContext);
+  const initialPersistedState = context?.getInitialState() as S | undefined;
+  const [state, setState] = useState(initialPersistedState ?? initialState);
+
+  console.log('usePersistentState', initialPersistedState, context?.isTracking);
+  if (context && context.isTracking) {
+    context?.addPanelState?.(state);
+  }
+
+  const setter = useCallback(
+    (newState: SetStateAction<S>) => {
+      console.log('setter');
+      setState(newState);
+      context?.trigger?.();
+      // });
+    },
+    [context]
+  );
+
+  return [state, setter];
+}

plugins/ui/src/js/src/layout/ReactPanel.tsx

@@ -1,4 +1,10 @@
-import React, { useCallback, useEffect, useMemo, useRef } from 'react';
+import React, {
+  useCallback,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+} from 'react';
 import ReactDOM from 'react-dom';
 import { nanoid } from 'nanoid';
 import {
@@ -89,7 +95,8 @@ function ReactPanel({
   UNSAFE_className,
 }: Props): JSX.Element | null {
   const layoutManager = useLayoutManager();
-  const { metadata, onClose, onOpen, panelId } = useReactPanel();
+  const { metadata, onClose, onOpen, panelId, onDataChange, getInitialData } =
+    useReactPanel();
   const portalManager = usePortalPanelManager();
   const portal = portalManager.get(panelId);
   const panelTitle = title ?? metadata?.name ?? '';
@@ -201,9 +208,50 @@ function ReactPanel({
     renderedChildren = children;
   }
 
+  const [render, setRender] = useState(0);
+
+  const persistentData = useRef({
+    initial: getInitialData(),
+    prev: [] as unknown[],
+    next: [] as unknown[],
+  });
+
+  console.log('panel render', renderedChildren);
+
+  const panelContextValue = useMemo(
+    () => ({
+      panelId,
+      addPanelState(state: unknown) {
+        persistentData.current.next.push(state);
+      },
+      initialStateIterator: persistentData.current.initial[Symbol.iterator](),
+      getInitialState() {
+        // eslint-disable-next-line react/no-this-in-sfc
+        const { value, done } = this.initialStateIterator.next();
+        return value;
+      },
+      trigger() {
+        persistentData.current.next = [];
+        // eslint-disable-next-line react/no-this-in-sfc
+        this.isTracking = true;
+        setRender(prev => prev + 1);
+      },
+      isTracking: false,
+    }),
+    [panelId]
+  );
+
+  useEffect(() => {
+    console.log('panel effect', persistentData.current.next.length);
+    persistentData.current.prev = persistentData.current.next;
+    persistentData.current.next = [];
+    panelContextValue.isTracking = false;
+    onDataChange(persistentData.current.prev);
+  }, [render, panelContextValue, onDataChange]);
+
   return portal
     ? ReactDOM.createPortal(
-        <ReactPanelContext.Provider value={panelId}>
+        <ReactPanelContext.Provider value={panelContextValue}>
           <View
             height="100%"
             width="100%"
@@ -239,7 +287,9 @@ function ReactPanel({
                  * Don't render the children if there's an error with the widget. If there's an error with the widget, we can assume the children won't render properly,
                  * but we still want the panels to appear so things don't disappear/jump around.
                  */}
-                {renderedChildren ?? null}
+                {React.Children.map(renderedChildren, child =>
+                  React.cloneElement(child as React.ReactElement)
+                )}
               </ReactPanelErrorBoundary>
             </Flex>
           </View>

plugins/ui/src/js/src/layout/ReactPanelContext.ts

@@ -1,14 +1,22 @@
 import { createContext, useContext } from 'react';
 
+type PanelContextType = {
+  panelId: string;
+  addPanelState: (state: unknown) => void;
+  trigger: () => void;
+  isTracking: boolean;
+  getInitialState: () => unknown;
+};
+
 /**
  * Context that holds the ID of the panel that we are currently in.
  */
-export const ReactPanelContext = createContext<string | null>(null);
+export const ReactPanelContext = createContext<PanelContextType | null>(null);
 
 /**
  * Gets the panel ID from the nearest panel context.
  * @returns The panel ID or null if not in a panel
  */
 export function usePanelId(): string | null {
-  return useContext(ReactPanelContext);
+  return useContext(ReactPanelContext)?.panelId ?? null;
 }