Progress Report Grid - Expansion tweaks

src/components/Comments/CommentsContext.tsx

@@ -1,4 +1,4 @@
-import { useLocalStorageState, useSet } from 'ahooks';
+import { useLocalStorageState } from 'ahooks';
 import { noop } from 'lodash';
 import {
   createContext,
@@ -9,16 +9,12 @@ import {
   useState,
 } from 'react';
 import { ChildrenProp } from '~/common';
-
-type ExpandedThreads = ReturnType<typeof useSet<string>>[1] & {
-  has: (threadId: string) => boolean;
-  toggle: (threadId: string, next?: boolean) => void;
-};
+import { SetHook, useSet } from '~/hooks';
 
 const initialCommentsBarContext = {
   toggleCommentsBar: noop as (state?: boolean) => void,
   isCommentsBarOpen: false,
-  expandedThreads: {} as unknown as ExpandedThreads,
+  expandedThreads: {} as unknown as SetHook<string>,
   resourceId: undefined as string | undefined,
   setResourceId: noop as (resourceId: string | undefined) => void,
 };
@@ -31,18 +27,7 @@ export const CommentsProvider = ({ children }: ChildrenProp) => {
 
   const [resourceId, setResourceId] = useState<string | undefined>(undefined);
 
-  const [currentExpandedThreads, setExpandedThreads] = useSet<string>();
-  const expandedThreads = useMemo(
-    (): ExpandedThreads => ({
-      ...setExpandedThreads,
-      has: currentExpandedThreads.has.bind(currentExpandedThreads),
-      toggle: (threadId: string, next?: boolean) => {
-        next = next ?? !currentExpandedThreads.has(threadId);
-        setExpandedThreads[next ? 'add' : 'remove'](threadId);
-      },
-    }),
-    [currentExpandedThreads, setExpandedThreads]
-  );
+  const expandedThreads = useSet<string>();
 
   const toggleCommentsBar = useCallback(
     (state?: boolean) => {

src/hooks/index.ts

@@ -1,4 +1,5 @@
 export * from './useFirstMountState';
+export * from './useSet';
 export * from './useUserAgent';
 export * from './useQueryParams';
 export * from './useIsomorphicEffect';

src/hooks/useSet.ts

@@ -0,0 +1,96 @@
+/* eslint-disable react-hooks/exhaustive-deps */
+import { useLatest } from 'ahooks';
+import { useMemo, useState } from 'react';
+import { Merge } from 'type-fest';
+
+export type SetHook<T> = Merge<Set<T>, SetModifiers<T>>;
+export interface SetModifiers<T> {
+  /**
+   * Add item to the set.
+   * No change if the item already exists, which the return value conveys.
+   */
+  readonly add: (item: T) => boolean &
+    // Not really, just so this can be typed as Set<T>
+    Set<T>;
+  /**
+   * Alias of {@link delete}.
+   */
+  readonly remove: (item: T) => boolean;
+  /**
+   * Remove an item from the set.
+   * No change if the item already exists, which the return value conveys.
+   */
+  readonly delete: (item: T) => boolean;
+  /**
+   * Toggle adding/removing an item from the set.
+   * Optionally, explicitly state if the value should be added/removed with the `next` arg.
+   * No change if the item already exists, which the return value conveys.
+   */
+  readonly toggle: (item: T, next?: boolean) => boolean;
+  /**
+   * Replace the entries with the ones given here.
+   * This always produces an identity change, even if the entry values are the same.
+   */
+  readonly set: (items: Iterable<T>) => void;
+  /**
+   * Remove all the entries.
+   * No change if the set is already empty.
+   */
+  readonly clear: () => void;
+  /**
+   * Reset the entries to the ones given in the hook creation.
+   * This always produces an identity change, even if the entry values are the same.
+   */
+  readonly reset: () => void;
+}
+
+/**
+ * Provides a Set in React state.
+ *
+ * Each change produces a new Set.
+ * However, the modifier functions (their identities) don't change
+ * and will always reference the current entries when needed.
+ *
+ * We do differ with `add()` return value, though.
+ * It returns whether an entry was added instead of returning `this`.
+ */
+export function useSet<T>(initialValue?: Iterable<T>): SetHook<T> {
+  const getInitValue = useLatest(() => new Set<T>(initialValue));
+  const [current, set] = useState(getInitValue.current);
+  const ref = useLatest(current);
+
+  const modifiers = useMemo((): SetModifiers<T> => {
+    const toggle = (item: T, next?: boolean) => {
+      // Keeping this found check independent of the setter scope below.
+      // React handles when it should be called and it could be not synchronous.
+      const found = ref.current.has(item);
+
+      set((prev) => {
+        // Check again here, because maybe prev has changed, and we don't want
+        // to change identity redundantly.
+        const currentlyIn = prev.has(item);
+        next = next ?? !currentlyIn;
+        if ((next && currentlyIn) || (!next && !currentlyIn)) {
+          return prev;
+        }
+
+        const temp = new Set(prev);
+        temp[next ? 'add' : 'delete'](item);
+        return temp;
+      });
+
+      return found;
+    };
+    return {
+      toggle,
+      add: (item) => toggle(item, true) as any,
+      remove: (item) => toggle(item, false),
+      delete: (item) => toggle(item, false),
+      set: (items) => set(new Set(items)),
+      reset: () => set(getInitValue.current),
+      clear: () => set((prev) => (prev.size === 0 ? prev : new Set())),
+    };
+  }, []);
+
+  return useMemo(() => Object.assign(current, modifiers), [current]);
+}

src/scenes/Dashboard/ProgressReportsWidget/ExpansionCell.tsx

@@ -1,23 +1,15 @@
 import { Box } from '@mui/material';
-import {
-  GridState,
-  GridRenderCellParams as RenderCellParams,
-  useGridSelector,
-} from '@mui/x-data-grid';
+import { GridRenderCellParams as RenderCellParams } from '@mui/x-data-grid';
 import { ChildrenProp, extendSx, StyleProps } from '~/common';
+import { useExpanded } from './expansionState';
 
 export const ExpansionCell = ({
   id,
-  api,
   sx,
   className,
   children,
 }: Pick<RenderCellParams, 'id' | 'api'> & StyleProps & ChildrenProp) => {
-  const selectedRows = useGridSelector(
-    { current: api },
-    (state: GridState) => state.rowSelection
-  );
-  const isExpanded = selectedRows.includes(id);
+  const isExpanded = useExpanded().has(id);
 
   return (
     <Box

src/scenes/Dashboard/ProgressReportsWidget/ProgressReportsExpandedGrid.tsx

@@ -3,13 +3,12 @@ import {
   DataGridProProps as DataGridProps,
   GridColDef,
   GridRenderCellParams,
-  GridRowId,
   GridToolbarColumnsButton,
   GridToolbarFilterButton,
   useGridApiRef,
 } from '@mui/x-data-grid-pro';
 import { entries } from '@seedcompany/common';
-import { useState } from 'react';
+import { useMemo } from 'react';
 import { extendSx } from '~/common';
 import {
   getInitialVisibility,
@@ -19,6 +18,12 @@ import {
   Toolbar,
   useFilterToggle,
 } from '~/components/Grid';
+import {
+  CollapseAllButton,
+  ExpandAllButton,
+  ExpansionContext,
+  useExpandedSetup,
+} from './expansionState';
 import {
   ExpansionMarker,
   ProgressReportsColumnMap,
@@ -95,6 +100,8 @@ const ProgressReportsToolbar = () => (
         Pinned
       </QuickFilterButton>
     </QuickFilters>
+    <CollapseAllButton />
+    <ExpandAllButton />
   </Toolbar>
 );
 
@@ -107,30 +114,41 @@ export const ProgressReportsExpandedGrid = (
 ) => {
   const apiRef = useGridApiRef();
 
-  const [selected, setSelected] = useState<GridRowId[]>([]);
+  const { expanded, onMouseDown, onRowClick } = useExpandedSetup();
+
+  const slotProps = useMemo(
+    (): DataGridProps['slotProps'] => ({
+      row: {
+        onMouseDown,
+      },
+    }),
+    [onMouseDown]
+  );
 
   return (
-    <ProgressReportsGrid
-      {...props}
-      density="standard"
-      slots={slots}
-      apiRef={apiRef}
-      columns={columns}
-      initialState={initialState}
-      onRowClick={({ id }) => setSelected(selected.length > 0 ? [] : [id])}
-      rowSelectionModel={selected}
-      getRowHeight={(params) =>
-        apiRef.current.isRowSelected(params.id) ? 'auto' : COLLAPSED_ROW_HEIGHT
-      }
-      sx={[
-        {
-          // Don't want 'auto' to shrink below this when the cell is empty
-          '.MuiDataGrid-cell': {
-            minHeight: COLLAPSED_ROW_HEIGHT,
+    <ExpansionContext.Provider value={expanded}>
+      <ProgressReportsGrid
+        {...props}
+        density="standard"
+        slots={slots}
+        apiRef={apiRef}
+        columns={columns}
+        initialState={initialState}
+        slotProps={slotProps}
+        onRowClick={onRowClick}
+        getRowHeight={(params) =>
+          expanded.has(params.id) ? 'auto' : COLLAPSED_ROW_HEIGHT
+        }
+        sx={[
+          {
+            // Don't want 'auto' to shrink below this when the cell is empty
+            '.MuiDataGrid-cell': {
+              minHeight: COLLAPSED_ROW_HEIGHT,
+            },
           },
-        },
-        ...extendSx(props.sx),
-      ]}
-    />
+          ...extendSx(props.sx),
+        ]}
+      />
+    </ExpansionContext.Provider>
   );
 };

src/scenes/Dashboard/ProgressReportsWidget/ProgressReportsWidget.tsx

@@ -40,6 +40,7 @@ export const ProgressReportsWidget = ({
           disablePortal
           options={quarter.available}
           getOptionLabel={(q) => `Q${q.fiscalQuarter} FY${q.fiscalYear}`}
+          isOptionEqualToValue={(a, b) => +a === +b}
           value={quarter.current}
           onChange={(_, q) => quarter.set(q)}
           disableClearable