From 8746fe1d87485787065fa7968f82c88f69f1c232 Mon Sep 17 00:00:00 2001
From: Axel Bocciarelli <axelboc@gmail.com>
Date: Tue, 3 Sep 2024 15:30:01 +0200
Subject: [PATCH] Enforce stricter types for formatters in `MappedMatrixVis`

---
 .../core/compound/MappedCompoundVis.tsx       |  8 ++----
 .../vis-packs/core/matrix/MappedMatrixVis.tsx |  8 ++----
 .../app/src/vis-packs/core/matrix/utils.ts    | 28 +++++++++++++------
 packages/shared/src/guards.ts                 | 10 +++++++
 4 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/packages/app/src/vis-packs/core/compound/MappedCompoundVis.tsx b/packages/app/src/vis-packs/core/compound/MappedCompoundVis.tsx
index 7b9d0fa2e..6db320990 100644
--- a/packages/app/src/vis-packs/core/compound/MappedCompoundVis.tsx
+++ b/packages/app/src/vis-packs/core/compound/MappedCompoundVis.tsx
@@ -14,7 +14,7 @@ import visualizerStyles from '../../../visualizer/Visualizer.module.css';
 import { useMappedArray, useSlicedDimsAndMapping } from '../hooks';
 import type { MatrixVisConfig } from '../matrix/config';
 import MatrixToolbar from '../matrix/MatrixToolbar';
-import { getCellWidth, getFormatter } from '../matrix/utils';
+import { getCellFormatter, getCellWidth } from '../matrix/utils';
 import { getSliceSelection } from '../utils';
 
 interface Props {
@@ -46,7 +46,7 @@ function MappedCompoundVis(props: Props) {
   );
 
   const fieldFormatters = Object.values(fields).map((field) =>
-    getFormatter(field, notation),
+    getCellFormatter(mappedArray, field, notation),
   );
 
   const { getExportURL } = useDataContext();
@@ -70,9 +70,7 @@ function MappedCompoundVis(props: Props) {
       <MatrixVis
         className={visualizerStyles.vis}
         dims={mappedArray.shape}
-        cellFormatter={(row, col) =>
-          fieldFormatters[col](mappedArray.get(row, col))
-        }
+        cellFormatter={(row, col) => fieldFormatters[col](row, col)}
         cellWidth={customCellWidth ?? cellWidth}
         columnHeaders={fieldNames}
         sticky={sticky}
diff --git a/packages/app/src/vis-packs/core/matrix/MappedMatrixVis.tsx b/packages/app/src/vis-packs/core/matrix/MappedMatrixVis.tsx
index 38e26b1b2..a810d6dfc 100644
--- a/packages/app/src/vis-packs/core/matrix/MappedMatrixVis.tsx
+++ b/packages/app/src/vis-packs/core/matrix/MappedMatrixVis.tsx
@@ -14,7 +14,7 @@ import { useMappedArray, useSlicedDimsAndMapping } from '../hooks';
 import { getSliceSelection } from '../utils';
 import type { MatrixVisConfig } from './config';
 import MatrixToolbar from './MatrixToolbar';
-import { getCellWidth, getFormatter } from './utils';
+import { getCellFormatter, getCellWidth } from './utils';
 
 interface Props {
   dataset: Dataset<ArrayShape, PrintableType>;
@@ -32,7 +32,7 @@ function MappedMatrixVis(props: Props) {
   const [slicedDims, slicedMapping] = useSlicedDimsAndMapping(dims, dimMapping);
   const mappedArray = useMappedArray(value, slicedDims, slicedMapping);
 
-  const formatter = getFormatter(type, notation);
+  const cellFormatter = getCellFormatter(mappedArray, type, notation);
   const cellWidth = getCellWidth(type);
 
   const { getExportURL } = useDataContext();
@@ -57,9 +57,7 @@ function MappedMatrixVis(props: Props) {
       <MatrixVis
         className={visualizerStyles.vis}
         dims={mappedArray.shape}
-        cellFormatter={(row: number, col: number) =>
-          formatter(mappedArray.get(row, col))
-        }
+        cellFormatter={cellFormatter}
         cellWidth={customCellWidth ?? cellWidth}
         sticky={sticky}
       />
diff --git a/packages/app/src/vis-packs/core/matrix/utils.ts b/packages/app/src/vis-packs/core/matrix/utils.ts
index 03a2c44b5..120af4b67 100644
--- a/packages/app/src/vis-packs/core/matrix/utils.ts
+++ b/packages/app/src/vis-packs/core/matrix/utils.ts
@@ -1,12 +1,13 @@
 import { Notation } from '@h5web/lib';
 import {
+  assertNdArrayValue,
   isBoolType,
   isComplexType,
   isEnumType,
   isNumericType,
 } from '@h5web/shared/guards';
 import type {
-  BooleanType,
+  ArrayValue,
   ComplexType,
   NumericType,
   PrintableCompoundType,
@@ -20,6 +21,7 @@ import {
   formatBool,
 } from '@h5web/shared/vis-utils';
 import { format } from 'd3-format';
+import type { NdArray } from 'ndarray';
 
 export function createNumericFormatter(
   notation: Notation,
@@ -45,27 +47,37 @@ export function createMatrixComplexFormatter(
   return createComplexFormatter(formatStr, true);
 }
 
-export function getFormatter(
+export function getCellFormatter(
+  dataArray: NdArray<ArrayValue<PrintableType>>,
   type: PrintableType,
   notation: Notation,
-): ValueFormatter<PrintableType> {
+): (row: number, col: number) => string {
   if (isComplexType(type)) {
-    return createMatrixComplexFormatter(notation);
+    assertNdArrayValue(type, dataArray);
+    const formatter = createMatrixComplexFormatter(notation);
+    return (row, col) => formatter(dataArray.get(row, col));
   }
 
   if (isNumericType(type)) {
-    return createNumericFormatter(notation);
+    assertNdArrayValue(type, dataArray);
+    const formatter = createNumericFormatter(notation);
+    return (row, col) => formatter(dataArray.get(row, col));
   }
 
   if (isBoolType(type)) {
-    return formatBool as ValueFormatter<BooleanType>;
+    assertNdArrayValue(type, dataArray);
+    return (row, col) => formatBool(dataArray.get(row, col));
   }
 
   if (isEnumType(type)) {
-    return createEnumFormatter(type.mapping);
+    assertNdArrayValue(type, dataArray);
+    const formatter = createEnumFormatter(type.mapping);
+    return (row, col) => formatter(dataArray.get(row, col));
   }
 
-  return (val) => (val as string).toString(); // call `toString()` for safety, in case type cast is wrong
+  // `StringType`
+  assertNdArrayValue(type, dataArray);
+  return (row, col) => dataArray.get(row, col);
 }
 
 export function getCellWidth(
diff --git a/packages/shared/src/guards.ts b/packages/shared/src/guards.ts
index d26ef51bb..b97119103 100644
--- a/packages/shared/src/guards.ts
+++ b/packages/shared/src/guards.ts
@@ -2,6 +2,7 @@ import type { Data, NdArray, TypedArray } from 'ndarray';
 
 import type {
   ArrayShape,
+  ArrayValue,
   BooleanType,
   ComplexArray,
   ComplexType,
@@ -438,6 +439,15 @@ function assertPrimitiveValue<T extends DType>(
   }
 }
 
+export function assertNdArrayValue<T extends DType>(
+  type: T,
+  value: NdArray<unknown[] | TypedArray>,
+): asserts value is NdArray<ArrayValue<T>> {
+  if (value.size > 0) {
+    assertPrimitiveValue(type, value.get(0));
+  }
+}
+
 export function assertDatasetValue<D extends Dataset<ScalarShape | ArrayShape>>(
   value: unknown,
   dataset: D,