Store model columns in IrisGridMetricCalculator

vbabich · vbabich · commit 0c1f05fba2b7 · 2025-09-17T14:28:28.000-06:00
diff --git a/packages/iris-grid/src/IrisGrid.tsx b/packages/iris-grid/src/IrisGrid.tsx
@@ -125,7 +125,6 @@ import {
 } from './mousehandlers';
 import ToastBottomBar from './ToastBottomBar';
 import IrisGridMetricCalculator, {
-  getColumnNameMap,
   IrisGridMetricState,
 } from './IrisGridMetricCalculator';
 import IrisGridModelUpdater from './IrisGridModelUpdater';
diff --git a/packages/iris-grid/src/IrisGridMetricCalculator.test.ts b/packages/iris-grid/src/IrisGridMetricCalculator.test.ts
@@ -2,7 +2,6 @@ import { GridMetricCalculator } from '@deephaven/grid';
 import { dh } from '@deephaven/jsapi-types';
 import { TestUtils } from '@deephaven/utils';
 import {
-  getColumnNameMap,
   IrisGridMetricCalculator,
   type IrisGridMetricState,
 } from './IrisGridMetricCalculator';
@@ -49,9 +48,7 @@ describe('IrisGridMetricCalculator', () => {
         return index !== -1 ? index : undefined;
       },
     });
-    calculator = new IrisGridMetricCalculator({
-      columnNameMap: getColumnNameMap(model),
-    });
+    calculator = new IrisGridMetricCalculator();
     state = makeGridMetricState(model);
   });
 
@@ -60,29 +57,22 @@ describe('IrisGridMetricCalculator', () => {
 
     expect(model.getColumnIndexByName('Column1')).toBe(0);
 
-    calculator.setColumnWidth(
-      state,
-      model.getColumnIndexByName('Column1'),
-      100
-    );
-    calculator.setColumnWidth(
-      state,
-      model.getColumnIndexByName('Column2'),
-      200
-    );
-    calculator.setColumnWidth(
-      state,
-      model.getColumnIndexByName('Column3'),
-      300
-    );
+    // setColumnWidth requires getMetrics call
+    calculator.getMetrics(state);
+    calculator.setColumnWidth(model.getColumnIndexByName('Column1'), 100);
+    calculator.setColumnWidth(model.getColumnIndexByName('Column2'), 200);
+    calculator.setColumnWidth(model.getColumnIndexByName('Column3'), 300);
 
+    // Calling getMetrics to update user column widths
+    calculator.getMetrics(state);
     expect(calculator.getUserColumnWidths().size).toBe(3);
 
     // Delete Column2
     columns = columns.filter(col => col.name !== 'Column2');
 
     expect(state.model.columns[1].name).toBe('Column3');
 
+    calculator.getMetrics(state);
     expect([...calculator.getUserColumnWidths().entries()]).toEqual([
       [0, 100],
       [1, 300],
@@ -98,6 +88,7 @@ describe('IrisGridMetricCalculator', () => {
       ...columns.slice(1),
     ];
 
+    calculator.getMetrics(state);
     expect([...calculator.getUserColumnWidths().entries()]).toEqual([
       [0, 100],
       [1, 200],
@@ -106,20 +97,13 @@ describe('IrisGridMetricCalculator', () => {
   });
 
   it('setColumnWidth updates user column width', () => {
-    calculator.setColumnWidth(
-      state,
-      model.getColumnIndexByName('Column1'),
-      100
-    );
+    calculator.getMetrics(state);
+    calculator.setColumnWidth(model.getColumnIndexByName('Column1'), 100);
 
     expect(calculator.getUserColumnWidths().size).toBe(1);
     expect(calculator.getUserColumnWidths().get(0)).toBe(100);
 
-    calculator.setColumnWidth(
-      state,
-      model.getColumnIndexByName('Column1'),
-      150
-    );
+    calculator.setColumnWidth(model.getColumnIndexByName('Column1'), 150);
     expect(calculator.getUserColumnWidths().get(0)).toBe(150);
   });
 });
diff --git a/packages/iris-grid/src/IrisGridMetricCalculator.ts b/packages/iris-grid/src/IrisGridMetricCalculator.ts
@@ -1,31 +1,87 @@
-import { GridMetricCalculator, ModelSizeMap } from '@deephaven/grid';
+import {
+  GridMetricCalculator,
+  GridMetrics,
+  ModelIndex,
+  ModelSizeMap,
+  trimMap,
+} from '@deephaven/grid';
 import type { GridMetricState } from '@deephaven/grid';
 import type { dh } from '@deephaven/jsapi-types';
+import { assertNotNull } from '@deephaven/utils';
 import type IrisGridModel from './IrisGridModel';
 import { IrisGridThemeType } from './IrisGridTheme';
+import {
+  ColumnName,
+  ReadonlyAdvancedFilterMap,
+  ReadonlyQuickFilterMap,
+} from './CommonTypes';
 
 export interface IrisGridMetricState extends GridMetricState {
   model: IrisGridModel;
   theme: IrisGridThemeType;
   isFilterBarShown: boolean;
-  advancedFilters: Map<
-    string,
-    { options: unknown; filter: dh.FilterCondition | null }
-  >;
-  quickFilters: Map<
-    string,
-    { text: string; filter: dh.FilterCondition | null }
-  >;
-  sorts: dh.Sort[];
+  advancedFilters: ReadonlyAdvancedFilterMap;
+  quickFilters: ReadonlyQuickFilterMap;
+  sorts: readonly dh.Sort[];
   reverse: boolean;
 }
 
-/**
- * Class to calculate all the metrics for a grid.
- * Call getMetrics() with the state to get metrics
- */
 export class IrisGridMetricCalculator extends GridMetricCalculator {
+  // Column widths by name to keep track of columns going in and out of viewport
+  userColumnWidthsByName: Map<ColumnName, number> = new Map();
+
+  private cachedModelColumns: readonly dh.Column[] | undefined;
+
+  // constructor(...args: IrisGridMetricCalculatorParameters) {
+  //   super(...args);
+  //   const [config] = args;
+  //   const { model, userColumnWidths } = config;
+
+  //   this.model = model;
+
+  //   // Populate userColumnWidthsByName from index-based userColumnWidths arg
+  //   // and the map of existing column names.
+  //   // We need to persist widths for all columns, even the ones that are not currently present because the model can expand and collapse
+  //   // TODO: this will be problematic for user width persistence since the model can initialize with partial set of columns
+  //   // and populates the rest on model update.
+  //   // TODO: change the arg type to Map<ColumnName, number>?
+  //   this.userColumnWidthsByName = new Map<ColumnName, number>();
+  //   userColumnWidths?.forEach((width, modelIndex) => {
+  //     const name = model.columns[modelIndex]?.name;
+  //     if (name != null) {
+  //       this.userColumnWidthsByName.set(name, width);
+  //     }
+  //   });
+  // }
+
+  /**
+   * Updates the user column widths based on the current model state
+   * @param model The current IrisGridModel
+   */
+  private updateUserColumnWidths(model: IrisGridModel): void {
+    this.userColumnWidths = new Map<ModelIndex, number>();
+    this.userColumnWidthsByName.forEach((width, name) => {
+      const modelIndex = model.getColumnIndexByName(name);
+      if (modelIndex != null) {
+        super.setColumnWidth(modelIndex, width);
+      }
+    });
+  }
+
+  /**
+   * Updates the user and calculated column widths if the model columns have changed
+   * @param model The current IrisGridModel
+   */
+  private updateColumnWidthsIfNecessary(model: IrisGridModel): void {
+    if (model.columns !== this.cachedModelColumns) {
+      this.resetCalculatedColumnWidths();
+      this.updateUserColumnWidths(model);
+    }
+  }
+
   getGridY(state: IrisGridMetricState): number {
+    // The state here seems to be a GridMetricState with stateOverrides passed from IrisGrid in the props,
+    // not guaranteed to be IrisGridMetricState
     let gridY = super.getGridY(state);
     const {
       isFilterBarShown,
@@ -50,7 +106,69 @@ export class IrisGridMetricCalculator extends GridMetricCalculator {
     return gridY;
   }
 
+  /**
+   * Gets the metrics for the current state. This method has to be called to update cachedModelColumns on model columns change.
+   * @param state The current IrisGridMetricState
+   * @returns The metrics for the current state
+   */
+  getMetrics(state: IrisGridMetricState): GridMetrics {
+    const { model } = state;
+    // Update column widths if columns in the cached model don't match the current model passed in the state
+    this.updateColumnWidthsIfNecessary(model);
+
+    this.cachedModelColumns = model.columns;
+    return super.getMetrics(state);
+  }
+
+  /**
+   * Sets the width for a specific column by index
+   * @param column The index of the column to set
+   * @param size The new width for the column
+   */
+  setColumnWidth(column: number, size: number): void {
+    super.setColumnWidth(column, size);
+    assertNotNull(
+      this.cachedModelColumns,
+      'setColumnWidth should be called after getMetrics'
+    );
+    const name = this.cachedModelColumns[column]?.name;
+    if (name != null) {
+      this.userColumnWidthsByName.set(name, size);
+      trimMap(this.userColumnWidthsByName);
+    }
+  }
+
+  /**
+   * Resets the width for a specific column by index
+   * @param column The index of the column to reset
+   */
+  resetColumnWidth(column: number): void {
+    super.resetColumnWidth(column);
+    assertNotNull(
+      this.cachedModelColumns,
+      'resetColumnWidth should be called after getMetrics'
+    );
+    const name = this.cachedModelColumns[column]?.name;
+    if (name != null) {
+      this.userColumnWidthsByName.delete(name);
+    }
+  }
+
+  /**
+   * Resets all user column widths
+   */
+  resetAllColumnWidths(): void {
+    this.userColumnWidths = new Map<ModelIndex, number>();
+    this.userColumnWidthsByName = new Map<ColumnName, number>();
+  }
+
+  /**
+   * Gets the user column widths for the current state
+   * @param state The current IrisGridMetricState
+   * @returns A map of user column widths
+   */
   getUserColumnWidths(): ModelSizeMap {
+    // This might return stale data if getMetrics hasn't been called
     return this.userColumnWidths;
   }
 }
