From 1006db5522d051ecc7260fe6725d752893394fe1 Mon Sep 17 00:00:00 2001
From: Tim Jennison <tjennison@verily.com>
Date: Mon, 3 Feb 2025 19:32:48 +0000
Subject: [PATCH] Refactor TreeGrid to have typed items and data

* Both items and data can be typed separately, though it mostly
  only makes sense to customize one or the other. EntityGroup
  criteria have generic data and typed items to store extra
  associated info such as the entity group but the study pages
  just have typed data because it's simpler.
* Pass the item to rowCustomization callbacks since it contains
  the data as well. Type it appropriately.
* Split the root children out so there's no need for an awkward
  root object.
---
 ui/src/addByCode.tsx                          |  25 ++---
 ui/src/addCohort.tsx                          |  13 +--
 ui/src/addCriteria.tsx                        |  35 +++---
 ui/src/cohortReview/cohortReviewList.tsx      |  11 +-
 ui/src/cohortReview/participantsList.tsx      |  16 +--
 .../cohortReview/plugins/occurrenceTable.tsx  |  29 ++---
 ui/src/components/treegrid.test.tsx           |  68 ++++++------
 ui/src/components/treegrid.tsx                |  50 +++++----
 ui/src/criteria/classification.tsx            |  47 +++-----
 ui/src/criteria/survey.tsx                    | 100 ++++++++----------
 ui/src/export.tsx                             |  19 ++--
 ui/src/featureSet/addFeatureSet.tsx           |  19 +---
 ui/src/featureSet/featureSet.tsx              |  18 ++--
 ui/src/sampleApp/studiesList.tsx              |  13 ++-
 ui/src/sampleApp/studyOverview.tsx            |  13 ++-
 15 files changed, 211 insertions(+), 265 deletions(-)

diff --git a/ui/src/addByCode.tsx b/ui/src/addByCode.tsx
index 54ff37858..f58008951 100644
--- a/ui/src/addByCode.tsx
+++ b/ui/src/addByCode.tsx
@@ -7,6 +7,7 @@ import {
   TreeGrid,
   TreeGridColumn,
   TreeGridId,
+  TreeGridItem,
   useArrayAsTreeGridData,
 } from "components/treegrid";
 import ActionBar from "actionBar";
@@ -24,7 +25,7 @@ import { cohortURL, useIsSecondBlock } from "router";
 import { insertCohortCriteria, useCohortContext } from "cohortContext";
 import { isValid } from "util/valid";
 
-type LookupEntryItem = {
+type LookupEntryData = {
   config?: JSX.Element;
   code: DataKey;
   name?: string;
@@ -48,7 +49,7 @@ export function AddByCode() {
     [underlay.criteriaSelectors]
   );
 
-  const lookupEntriesState = useSWRMutation<LookupEntryItem[]>(
+  const lookupEntriesState = useSWRMutation<LookupEntryData[]>(
     {
       component: "AddByCode",
       query,
@@ -110,7 +111,7 @@ export function AddByCode() {
   const data = useArrayAsTreeGridData(lookupEntriesState?.data ?? [], "code");
 
   const onInsert = useCallback(() => {
-    const configMap = new Map<string, LookupEntryItem[]>();
+    const configMap = new Map<string, LookupEntryData[]>();
     lookupEntriesState.data?.forEach((e) => {
       if (!e.entry || !selected.has(e.code)) {
         return;
@@ -196,19 +197,13 @@ export function AddByCode() {
         </GridBox>
         <Loading immediate showProgressOnMutate status={lookupEntriesState}>
           {lookupEntriesState.data?.length ? (
-            <TreeGrid
+            <TreeGrid<TreeGridItem<LookupEntryData>>
               columns={columns}
               data={data}
-              rowCustomization={(id: TreeGridId) => {
-                if (!lookupEntriesState.data) {
-                  return undefined;
-                }
-
-                const item = data.get(id)?.data as LookupEntryItem;
-                if (!item) {
-                  return undefined;
-                }
-
+              rowCustomization={(
+                id: TreeGridId,
+                { data }: TreeGridItem<LookupEntryData>
+              ) => {
                 const sel = selected.has(id);
                 return [
                   {
@@ -218,7 +213,7 @@ export function AddByCode() {
                         size="small"
                         fontSize="inherit"
                         checked={sel}
-                        disabled={!item.entry}
+                        disabled={!data.entry}
                         onChange={() => {
                           updateSelected((selected) => {
                             if (sel) {
diff --git a/ui/src/addCohort.tsx b/ui/src/addCohort.tsx
index 40dc87749..cd7c176ef 100644
--- a/ui/src/addCohort.tsx
+++ b/ui/src/addCohort.tsx
@@ -51,7 +51,7 @@ export function AddCohort() {
   const navigate = useNavigate();
   const params = useBaseParams();
 
-  const cohortsState = useSWR(
+  const cohortsState = useSWR<CohortData[]>(
     {
       type: "cohorts",
       studyId,
@@ -112,16 +112,7 @@ export function AddCohort() {
             <TreeGrid
               data={data}
               columns={columns}
-              rowCustomization={(id: TreeGridId) => {
-                if (!cohortsState.data) {
-                  return undefined;
-                }
-
-                const cohortData = data.get(id)?.data as CohortData;
-                if (!cohortData) {
-                  return undefined;
-                }
-
+              rowCustomization={(id: TreeGridId, { data: cohortData }) => {
                 return [
                   {
                     column: columns.length - 2,
diff --git a/ui/src/addCriteria.tsx b/ui/src/addCriteria.tsx
index 0bb903992..6662206bc 100644
--- a/ui/src/addCriteria.tsx
+++ b/ui/src/addCriteria.tsx
@@ -23,7 +23,7 @@ import {
   TreeGrid,
   TreeGridId,
   TreeGridItem,
-  TreeGridRowData,
+  TreeGridData,
 } from "components/treegrid";
 import { MergedItem } from "data/mergeLists";
 import { Criteria, isTemporalSection } from "data/source";
@@ -352,9 +352,7 @@ function AddCriteria(props: AddCriteriaProps) {
 
   const search = useCallback(async () => {
     const children: DataKey[] = [];
-    const data = new Map<TreeGridId, CriteriaItem>([
-      ["root", { data: {}, children }],
-    ]);
+    const rows = new Map<TreeGridId, CriteriaItem>();
 
     if (query) {
       const res = await searchCriteria(
@@ -376,13 +374,16 @@ function AddCriteria(props: AddCriteriaProps) {
           },
           entry: entry,
         };
-        data.set(key, item);
+        rows.set(key, item);
       });
     }
 
-    return data;
+    return {
+      rows,
+      children,
+    };
   }, [underlaySource, query, selectedOptions, optionsMap]);
-  const searchState = useSWRImmutable<Map<TreeGridId, CriteriaItem>>(
+  const searchState = useSWRImmutable<TreeGridData<CriteriaItem>>(
     {
       component: "AddCriteria",
       underlayName: underlay.name,
@@ -522,7 +523,7 @@ function AddCriteria(props: AddCriteriaProps) {
         {!!query ? (
           <Paper>
             <Loading status={searchState}>
-              {!searchState.data?.get("root")?.children?.length ? (
+              {!searchState.data?.children?.length ? (
                 <Empty
                   minHeight="300px"
                   image={emptyImage}
@@ -531,17 +532,9 @@ function AddCriteria(props: AddCriteriaProps) {
               ) : (
                 <TreeGrid
                   columns={columns}
-                  data={searchState.data ?? new Map()}
-                  rowCustomization={(
-                    id: TreeGridId,
-                    rowData: TreeGridRowData
-                  ) => {
-                    if (!searchState.data) {
-                      return undefined;
-                    }
-
-                    const item = searchState.data.get(id);
-                    const option = optionsMap.get(item?.entry?.source ?? "");
+                  data={searchState.data}
+                  rowCustomization={(id, item) => {
+                    const option = optionsMap.get(item.entry?.source ?? "");
                     if (!option || !item?.entry?.source) {
                       throw new Error(
                         `Item source "${item?.entry?.source}" doesn't match any criteria config ID.`
@@ -554,7 +547,9 @@ function AddCriteria(props: AddCriteriaProps) {
                         content: (
                           <GridLayout colAlign="center">
                             <Button
-                              data-testid={rowData[searchConfig.columns[0].key]}
+                              data-testid={
+                                item.data[searchConfig.columns[0].key]
+                              }
                               onClick={() => onClick(option, item.entry?.data)}
                               variant="outlined"
                             >
diff --git a/ui/src/cohortReview/cohortReviewList.tsx b/ui/src/cohortReview/cohortReviewList.tsx
index a33615c68..7def2c2e4 100644
--- a/ui/src/cohortReview/cohortReviewList.tsx
+++ b/ui/src/cohortReview/cohortReviewList.tsx
@@ -687,16 +687,7 @@ function Annotations() {
             <TreeGrid
               data={data}
               columns={columns}
-              rowCustomization={(id: TreeGridId) => {
-                if (!annotationsState.data) {
-                  return undefined;
-                }
-
-                const annotation = data.get(id)?.data;
-                if (!annotation) {
-                  return undefined;
-                }
-
+              rowCustomization={(id: TreeGridId, { data: annotation }) => {
                 return [
                   {
                     column: columns.length - 1,
diff --git a/ui/src/cohortReview/participantsList.tsx b/ui/src/cohortReview/participantsList.tsx
index edaf5b77b..435b1f5a2 100644
--- a/ui/src/cohortReview/participantsList.tsx
+++ b/ui/src/cohortReview/participantsList.tsx
@@ -11,7 +11,7 @@ import {
   useReviewSearchState,
 } from "cohortReview/reviewHooks";
 import Loading from "components/loading";
-import { TreeGrid, TreeGridData, TreeGridId } from "components/treegrid";
+import { TreeGrid, TreeGridId } from "components/treegrid";
 import GridLayout from "layout/gridLayout";
 import React, { ReactNode, useMemo, useState } from "react";
 
@@ -125,7 +125,8 @@ function ParticipantsList(props: ParticipantsListProps) {
   );
 
   const data = useMemo(() => {
-    const data: TreeGridData = new Map([["root", { data: {}, children: [] }]]);
+    const children: TreeGridId[] = [];
+    const rows = new Map();
 
     instancesState.data
       ?.slice(
@@ -134,13 +135,13 @@ function ParticipantsList(props: ParticipantsListProps) {
       )
       .forEach((instance) => {
         const key = instance.data.key;
-        data.set(key, { data: { ...instance.data } });
-        data.get("root")?.children?.push(key);
+        rows.set(key, { data: { ...instance.data } });
+        children.push(key);
 
         annotationsState.data?.forEach((a) => {
           const values = instance.annotations.get(a.id);
           if (values) {
-            const valueData = data.get(key)?.data;
+            const valueData = rows.get(key)?.data;
             if (valueData) {
               valueData[`t_${a.id}`] = values[values.length - 1].value;
             }
@@ -148,7 +149,10 @@ function ParticipantsList(props: ParticipantsListProps) {
         });
       });
 
-    return data;
+    return {
+      rows,
+      children,
+    };
   }, [instancesState, annotationsState, props.page, props.rowsPerPage]);
 
   return (
diff --git a/ui/src/cohortReview/plugins/occurrenceTable.tsx b/ui/src/cohortReview/plugins/occurrenceTable.tsx
index 7ebcee65c..026bcffde 100644
--- a/ui/src/cohortReview/plugins/occurrenceTable.tsx
+++ b/ui/src/cohortReview/plugins/occurrenceTable.tsx
@@ -6,7 +6,6 @@ import {
 import {
   TreeGrid,
   TreeGridColumn,
-  TreeGridData,
   TreeGridSortDirection,
   TreeGridSortOrder,
 } from "components/treegrid";
@@ -62,14 +61,17 @@ function OccurrenceTable({ id, config }: { id: string; config: Config }) {
 
   const data = useMemo(() => {
     const children: DataKey[] = [];
-    const data: TreeGridData = new Map([["root", { data: {}, children }]]);
+    const rows = new Map();
 
     context.rows[config.entity]?.forEach((o) => {
-      data.set(o.key, { data: o });
+      rows.set(o.key, { data: o });
       children.push(o.key);
     });
 
-    return data;
+    return {
+      rows,
+      children,
+    };
   }, [context]);
 
   const filterRegExps = useMemo(() => {
@@ -82,25 +84,24 @@ function OccurrenceTable({ id, config }: { id: string; config: Config }) {
 
   const sortedData = useMemo(() => {
     return produce(data, (data) => {
-      const root = data.get("root");
-      if (!root) {
-        return;
-      }
-
-      root.children = (root.children ?? []).filter((child) =>
+      data.children = (data.children ?? []).filter((child) =>
         Object.entries(filterRegExps ?? {}).reduce(
           (cur: boolean, [col, re]) =>
             cur &&
             re.test(
-              stringifyDataValue(data.get(child)?.data?.[col] as DataValue)
+              stringifyDataValue(data.rows.get(child)?.data?.[col] as DataValue)
             ),
           true
         )
       );
-      root.children.sort((a, b) => {
+      data.children.sort((a, b) => {
         for (const o of searchState.sortOrders ?? []) {
-          const valA = data.get(a)?.data?.[o.column] as DataValue | undefined;
-          const valB = data.get(b)?.data?.[o.column] as DataValue | undefined;
+          const valA = data.rows.get(a)?.data?.[o.column] as
+            | DataValue
+            | undefined;
+          const valB = data.rows.get(b)?.data?.[o.column] as
+            | DataValue
+            | undefined;
           const c = compareDataValues(valA, valB);
           if (c !== 0) {
             return o.direction === TreeGridSortDirection.Asc ? c : -c;
diff --git a/ui/src/components/treegrid.test.tsx b/ui/src/components/treegrid.test.tsx
index 5a5cbe2be..251df2a50 100644
--- a/ui/src/components/treegrid.test.tsx
+++ b/ui/src/components/treegrid.test.tsx
@@ -32,45 +32,41 @@ test("Table renders correctly", async () => {
     },
   ];
 
-  const data = new Map<TreeGridId, TreeGridItem>([
-    [
-      "root",
-      {
-        data: {},
-        children: [1, 3],
-      },
-    ],
-    [
-      1,
-      {
-        data: {
-          col1: "1-col1",
-          col2: "1-col2",
+  const data: TreeGridData = {
+    children: [1, 3],
+    rows: new Map([
+      [
+        1,
+        {
+          data: {
+            col1: "1-col1",
+            col2: "1-col2",
+          },
+          children: [2],
         },
-        children: [2],
-      },
-    ],
-    [
-      2,
-      {
-        data: {
-          col1: "2-col1",
-          col2: "2-col2",
-          col3: "2-col3",
+      ],
+      [
+        2,
+        {
+          data: {
+            col1: "2-col1",
+            col2: "2-col2",
+            col3: "2-col3",
+          },
         },
-      },
-    ],
-    [
-      3,
-      {
-        data: {
-          col1: "3-col1",
-          col2: "3-col2",
-          col3: <AccountTreeIcon />,
+      ],
+      [
+        3,
+        {
+          data: {
+            col1: "3-col1",
+            col2: "3-col2",
+            col3: <AccountTreeIcon />,
+          },
         },
-      },
-    ],
-  ]);
+      ],
+    ]),
+  };
 
   const rowCustomization = (id: TreeGridId) => {
     // Add variety to the rows.
diff --git a/ui/src/components/treegrid.tsx b/ui/src/components/treegrid.tsx
index 0116f7eeb..152a08d91 100644
--- a/ui/src/components/treegrid.tsx
+++ b/ui/src/components/treegrid.tsx
@@ -52,10 +52,10 @@ export type TreeGridItem<RowType extends TreeGridRowData = TreeGridRowData> = {
   data: RowType;
 };
 
-export type TreeGridData<ItemType extends TreeGridItem = TreeGridItem> = Map<
-  TreeGridId,
-  ItemType
->;
+export type TreeGridData<ItemType extends TreeGridItem = TreeGridItem> = {
+  rows: Map<TreeGridId, ItemType>;
+  children: TreeGridId[];
+};
 
 export type TreeGridColumn = {
   key: string;
@@ -95,7 +95,7 @@ export type TreeGridProps<ItemType extends TreeGridItem = TreeGridItem> = {
   highlightId?: TreeGridId;
   rowCustomization?: (
     id: TreeGridId,
-    data: TreeGridRowData
+    item: ItemType
   ) => ColumnCustomization[] | undefined;
   loadChildren?: (id: TreeGridId) => Promise<void>;
   minWidth?: boolean;
@@ -373,8 +373,6 @@ export function TreeGrid<ItemType extends TreeGridItem = TreeGridItem>(
             state,
             (id: TreeGridId) =>
               updateState((draft) => toggleExpanded(draft, id)),
-            "root",
-            "root",
             0,
             false, // collapse
             true, // first
@@ -393,19 +391,18 @@ export function useArrayAsTreeGridData<
 >(array: T[], key: K): TreeGridData<TreeGridItem<T>> {
   return useMemo(() => {
     const children: TreeGridId[] = [];
-    const data = new Map<TreeGridId, TreeGridItem<T>>([
-      // Force empty data here since it's never accessed but making it optional
-      // is a pain for "data" access everywhere.
-      ["root", { data: {} as T, children }],
-    ]);
+    const rows = new Map<TreeGridId, TreeGridItem<T>>();
 
     array?.forEach((a) => {
       const k = a[key] as TreeGridId;
-      data.set(k, { data: a });
+      rows.set(k, { data: a });
       children.push(k);
     });
 
-    return data;
+    return {
+      rows,
+      children,
+    };
   }, [array]);
 }
 
@@ -421,23 +418,24 @@ export function fromProtoColumns(
   }));
 }
 
-function renderChildren(
+function renderChildren<ItemType extends TreeGridItem = TreeGridItem>(
   theme: Theme,
-  props: TreeGridProps,
+  props: TreeGridProps<ItemType>,
   state: TreeGridState,
   toggleExpanded: (id: TreeGridId) => void,
-  id: TreeGridId,
-  key: string,
   indent: number,
   collapse: boolean,
   first: boolean,
   highlightRef: MutableRefObject<HTMLTableRowElement | null>,
-  highlightId?: TreeGridId
+  highlightId?: TreeGridId,
+  id?: TreeGridId,
+  key?: string
 ): JSX.Element[] {
   const results: JSX.Element[] = [];
 
-  props.data.get(id)?.children?.forEach((childId) => {
-    const child = props.data.get(childId);
+  const children = id ? props.data.rows.get(id)?.children : props.data.children;
+  children?.forEach((childId) => {
+    const child = props.data.rows.get(childId);
     if (!child) {
       return;
     }
@@ -447,7 +445,7 @@ function renderChildren(
     if (!child.data) {
       throw new Error(`'data' is undefined for ${JSON.stringify(child)}`);
     }
-    const rowCustomization = props.rowCustomization?.(childId, child.data);
+    const rowCustomization = props.rowCustomization?.(childId, child);
 
     const renderColumn = (
       column: number,
@@ -543,7 +541,7 @@ function renderChildren(
       );
     };
 
-    const childKey = `${key}~${childId}`;
+    const childKey = `${key ?? "root"}~${childId}`;
     results.push(
       <tr
         key={childKey}
@@ -631,13 +629,13 @@ function renderChildren(
         props,
         state,
         toggleExpanded,
-        childId,
-        childKey,
         indent + 1,
         collapse || childState?.status !== Status.Expanded,
         false,
         highlightRef,
-        highlightId
+        highlightId,
+        childId,
+        childKey
       )
     );
   });
diff --git a/ui/src/criteria/classification.tsx b/ui/src/criteria/classification.tsx
index f9ada89bf..cbad42ab2 100644
--- a/ui/src/criteria/classification.tsx
+++ b/ui/src/criteria/classification.tsx
@@ -315,9 +315,9 @@ function ClassificationEdit(props: ClassificationEditProps) {
       nodes: MergedItem<EntityNode>[],
       hierarchy?: DataKey[],
       parent?: DataKey,
-      prevData?: TreeGridData
-    ): Map<TreeGridId, EntityNodeItem> => {
-      const data = prevData ?? new Map();
+      prevData?: TreeGridData<EntityNodeItem>
+    ): TreeGridData<EntityNodeItem> => {
+      const data = prevData ?? { rows: new Map(), children: [] };
 
       const children: DataKey[] = [];
       nodes.forEach(({ source: entityGroup, data: node }) => {
@@ -339,7 +339,7 @@ function ClassificationEdit(props: ClassificationEditProps) {
         // Copy over existing children in case they're being loaded in
         // parallel. Grouping children and non-grouping, non-hierarchy nodes
         // always have no children.
-        let childChildren = data.get(key)?.children;
+        let childChildren = data.rows.get(key)?.children;
         if (!childChildren) {
           if (
             (!hierarchy && group && parent) ||
@@ -357,17 +357,14 @@ function ClassificationEdit(props: ClassificationEditProps) {
           entityGroup,
           groupingNode: !!group && !parent,
         };
-        data.set(key, cItem);
+        data.rows.set(key, cItem);
       });
 
       if (parent) {
         // Store children even if the data isn't loaded yet.
-        data.set(parent, { ...data.get(parent), children });
+        data.rows.set(parent, { ...data.rows.get(parent), children });
       } else {
-        data.set("root", {
-          children: [...(data?.get("root")?.children || []), ...children],
-          data: {},
-        });
+        data.children = [...(data.children ?? []), ...children];
       }
 
       return data;
@@ -521,7 +518,7 @@ function ClassificationEdit(props: ClassificationEditProps) {
 
   // Partially update the state when expanding rows in the hierarchy view.
   const updateData = useCallback(
-    (data: Map<TreeGridId, EntityNodeItem>) => {
+    (data) => {
       instancesState.mutate(data, { revalidate: false });
     },
     [instancesState]
@@ -625,7 +622,7 @@ function ClassificationEdit(props: ClassificationEditProps) {
             </GridBox>
           </GridBox>
           <Loading status={instancesState}>
-            {!instancesState.data?.get("root")?.children?.length ? (
+            {!instancesState.data?.children?.length ? (
               <Empty
                 minHeight="300px"
                 image={emptyImage}
@@ -641,23 +638,12 @@ function ClassificationEdit(props: ClassificationEditProps) {
                 highlightId={searchState?.highlightId}
                 expandable
                 reserveExpansionSpacing={!!searchState?.hierarchy}
-                rowCustomization={(
-                  id: TreeGridId,
-                  rowData: TreeGridRowData
-                ) => {
-                  if (!instancesState.data) {
-                    return undefined;
-                  }
-
-                  // TODO(tjennison): Make TreeGridData's type generic so we can
-                  // avoid this type assertion. Also consider passing the
-                  // TreeGridItem to the callback instead of the TreeGridRowData.
-                  const item = instancesState.data.get(id) as EntityNodeItem;
-                  if (!item || item.groupingNode) {
+                rowCustomization={(id, item) => {
+                  if (item.groupingNode) {
                     return undefined;
                   }
 
-                  const name = rowData[nameAttribute(props.config)];
+                  const name = item.data[nameAttribute(props.config)];
                   const newItem = {
                     key: item.node.data.key,
                     name: !!name ? String(name) : "",
@@ -674,9 +660,10 @@ function ClassificationEdit(props: ClassificationEditProps) {
                       onClick={() => {
                         setHierarchyQuery(undefined);
                         updateSearchState((data: SearchState) => {
-                          if (rowData.view_hierarchy) {
+                          if (item.data.view_hierarchy) {
                             data.hierarchyEntityGroup = item.entityGroup;
-                            data.hierarchy = rowData.view_hierarchy as string[];
+                            data.hierarchy = item.data
+                              .view_hierarchy as string[];
                             data.highlightId = id;
                           }
                         });
@@ -776,14 +763,14 @@ function ClassificationEdit(props: ClassificationEditProps) {
 
                   return [...listContent, ...hierarchyContent];
                 }}
-                loadChildren={(id: TreeGridId) => {
+                loadChildren={(id) => {
                   const data = instancesState.data;
                   if (!data) {
                     return Promise.resolve();
                   }
 
                   const findEntityGroupConfig = () => {
-                    const item = data.get(id) as EntityNodeItem;
+                    const item = data.rows.get(id);
                     if (item) {
                       const config = configEntityGroups(props.config).find(
                         (eg) => eg.id === item.entityGroup
diff --git a/ui/src/criteria/survey.tsx b/ui/src/criteria/survey.tsx
index b0933401d..6bb67846b 100644
--- a/ui/src/criteria/survey.tsx
+++ b/ui/src/criteria/survey.tsx
@@ -13,7 +13,6 @@ import {
   fromProtoColumns,
   TreeGrid,
   TreeGridColumn,
-  TreeGridData,
   TreeGridId,
   TreeGridItem,
   TreeGridRowData,
@@ -79,8 +78,6 @@ type EntityNodeItem = TreeGridItem & {
   parentKey?: DataKey;
 };
 
-type EntityTreeGridData = TreeGridData<EntityNodeItem>;
-
 // Exported for testing purposes.
 export interface Data {
   selected: Selection[];
@@ -276,60 +273,68 @@ function SurveyEdit(props: SurveyEditProps) {
 
   const processEntities = useCallback(
     (allEntityGroups: [string, EntityNode[]][]) => {
-      const data: EntityTreeGridData = new Map();
+      const children: TreeGridId[] = [];
+      const rows = new Map<TreeGridId, EntityNodeItem>();
 
       allEntityGroups.forEach(([entityGroup, nodes]) => {
         nodes.forEach((node) => {
           const rowData: TreeGridRowData = { ...node.data };
           const key = dataKeyToKey(node.data.key, entityGroup);
 
-          let parentKey = "root";
+          let parentKey: string | undefined;
           if (node.ancestors?.length) {
             parentKey = dataKeyToKey(node.ancestors[0], entityGroup);
           }
 
           const cItem: EntityNodeItem = {
             data: rowData,
-            children: data.get(key)?.children ?? [],
+            children: rows.get(key)?.children ?? [],
             node: node,
             entityGroup,
             type:
-              data.get(key)?.type ??
+              rows.get(key)?.type ??
               (node.childCount === 0
                 ? EntityNodeItemType.Answer
                 : EntityNodeItemType.Topic),
             parentKey: parentKey,
           };
-          data.set(key, cItem);
-
-          const parent = data.get(parentKey);
-          if (parent) {
-            parent.children?.push(key);
-            if (cItem.type === EntityNodeItemType.Answer) {
-              parent.type = EntityNodeItemType.Question;
+          rows.set(key, cItem);
+
+          if (parentKey) {
+            const parent = rows.get(parentKey);
+            if (parent) {
+              parent.children?.push(key);
+              if (cItem.type === EntityNodeItemType.Answer) {
+                parent.type = EntityNodeItemType.Question;
+              }
+            } else {
+              const d = {
+                key: parentKey,
+              };
+              rows.set(parentKey, {
+                data: d,
+                node: {
+                  data: d,
+                  entity: "loading",
+                },
+                entityGroup,
+                type:
+                  cItem.type === EntityNodeItemType.Answer
+                    ? EntityNodeItemType.Question
+                    : EntityNodeItemType.Topic,
+                children: [key],
+              });
             }
           } else {
-            const d = {
-              key: parentKey,
-            };
-            data.set(parentKey, {
-              data: d,
-              node: {
-                data: d,
-                entity: "loading",
-              },
-              entityGroup,
-              type:
-                cItem.type === EntityNodeItemType.Answer
-                  ? EntityNodeItemType.Question
-                  : EntityNodeItemType.Topic,
-              children: [key],
-            });
+            children.push(key);
           }
         });
       });
 
-      return data;
+      return {
+        rows,
+        children,
+      };
     },
     []
   );
@@ -397,10 +402,10 @@ function SurveyEdit(props: SurveyEditProps) {
   const [filteredData, defaultExpanded] = useMemo(() => {
     const data = instancesState.data;
     if (!data) {
-      return [new Map(), []];
+      return [{ rows: new Map(), children: [] }, []];
     }
     if (!searchState?.query) {
-      return [data, data.get("root")?.children ?? []];
+      return [data, data.children];
     }
 
     // TODO(tjennison): Handle RegExp errors.
@@ -409,7 +414,7 @@ function SurveyEdit(props: SurveyEditProps) {
     const ancestors = new Set<TreeGridId>();
 
     const matchNode = (key: TreeGridId) => {
-      const node = data.get(key);
+      const node = data.rows.get(key);
       if (!node) {
         throw new Error(`Tree node not found for ${key}.`);
       }
@@ -432,7 +437,7 @@ function SurveyEdit(props: SurveyEditProps) {
 
     const filtered = produce(data, (data) => {
       for (const key in data) {
-        const node = data.get(key);
+        const node = data.rows.get(key);
         if (!node) {
           throw new Error(`Tree node not found for ${key}.`);
         }
@@ -482,7 +487,7 @@ function SurveyEdit(props: SurveyEditProps) {
             />
           </GridBox>
           <Loading status={instancesState}>
-            {!filteredData?.get("root")?.children?.length ? (
+            {!filteredData?.children?.length ? (
               <Empty
                 minHeight="300px"
                 image={emptyImage}
@@ -495,22 +500,7 @@ function SurveyEdit(props: SurveyEditProps) {
                 expandable
                 defaultExpanded={defaultExpanded}
                 reserveExpansionSpacing
-                rowCustomization={(
-                  id: TreeGridId,
-                  rowData: TreeGridRowData
-                ) => {
-                  if (!instancesState.data) {
-                    return undefined;
-                  }
-
-                  // TODO(tjennison): Make TreeGridData's type generic so we can
-                  // avoid this type assertion. Also consider passing the
-                  // TreeGridItem to the callback instead of the TreeGridRowData.
-                  const item = instancesState.data.get(id);
-                  if (!item) {
-                    return undefined;
-                  }
-
+                rowCustomization={(id: TreeGridId, item) => {
                   const entityGroupSet = selectedSets.get(item.entityGroup);
                   const found = !!entityGroupSet?.has(item.node.data.key);
                   const foundAncestor = !!item.node.ancestors?.reduce(
@@ -539,10 +529,12 @@ function SurveyEdit(props: SurveyEditProps) {
                                 const question =
                                   item.parentKey &&
                                   item.type === EntityNodeItemType.Answer
-                                    ? instancesState.data?.get(item.parentKey)
+                                    ? instancesState.data?.rows?.get(
+                                        item.parentKey
+                                      )
                                     : undefined;
                                 const name =
-                                  rowData[nameAttribute(props.config)];
+                                  item.data[nameAttribute(props.config)];
                                 const questionName =
                                   question?.node?.data?.[
                                     nameAttribute(props.config)
diff --git a/ui/src/export.tsx b/ui/src/export.tsx
index 213c204d3..1b1ae5551 100644
--- a/ui/src/export.tsx
+++ b/ui/src/export.tsx
@@ -27,7 +27,7 @@ import Checkbox from "components/checkbox";
 import Empty from "components/empty";
 import Loading from "components/loading";
 import { Tabs as TanagraTabs } from "components/tabs";
-import { TreeGrid, TreeGridData } from "components/treegrid";
+import { TreeGrid, TreeGridId, TreeGridData } from "components/treegrid";
 import { Cohort, ExportModel, ExportResultLink, FeatureSet } from "data/source";
 import { useStudySource } from "data/studySourceContext";
 import { useUnderlaySource } from "data/underlaySourceContext";
@@ -496,19 +496,21 @@ function PreviewTable(props: PreviewTableProps) {
             props.featureSets
           );
 
-          const data: TreeGridData = new Map([
-            ["root", { data: {}, children: [] }],
-          ]);
+          const children: TreeGridId[] = [];
+          const rows = new Map();
 
           res.data.forEach((entry, i) => {
-            data.set(i, { data: entry });
-            data.get("root")?.children?.push(i);
+            rows.set(i, { data: entry });
+            children.push(i);
           });
 
           return {
             name: filters.name,
             sql: filters.sql ?? res.sql,
-            data: data,
+            data: {
+              rows,
+              children,
+            },
           };
         })
       );
@@ -572,8 +574,7 @@ function PreviewTable(props: PreviewTableProps) {
                   </Typography>
                 </GridBox>
               ) : tabDataState.data?.[tab]?.data ? (
-                tabDataState.data?.[tab]?.data?.get("root")?.children
-                  ?.length ? (
+                tabDataState.data?.[tab]?.data?.children?.length ? (
                   <TreeGrid
                     data={tabDataState.data?.[tab]?.data}
                     columns={props.occurrenceFilters[tab]?.attributes.map(
diff --git a/ui/src/featureSet/addFeatureSet.tsx b/ui/src/featureSet/addFeatureSet.tsx
index 518831629..ca094d5f1 100644
--- a/ui/src/featureSet/addFeatureSet.tsx
+++ b/ui/src/featureSet/addFeatureSet.tsx
@@ -7,11 +7,7 @@ import { generateId, getCriteriaTitle, sectionName } from "cohort";
 import { insertCohortCriteria, useCohortContext } from "cohortContext";
 import Empty from "components/empty";
 import Loading from "components/loading";
-import {
-  TreeGrid,
-  TreeGridId,
-  useArrayAsTreeGridData,
-} from "components/treegrid";
+import { TreeGrid, useArrayAsTreeGridData } from "components/treegrid";
 import { FeatureSet } from "data/source";
 import { useStudySource } from "data/studySourceContext";
 import { useUnderlaySource } from "data/underlaySourceContext";
@@ -49,7 +45,7 @@ export function AddFeatureSet() {
   const params = useBaseParams();
   const secondBlock = useIsSecondBlock();
 
-  const featureSetsState = useSWR(
+  const featureSetsState = useSWR<FeatureSetData[]>(
     {
       type: "featureSets",
       studyId,
@@ -117,16 +113,7 @@ export function AddFeatureSet() {
             <TreeGrid
               data={data}
               columns={columns}
-              rowCustomization={(id: TreeGridId) => {
-                if (!featureSetsState.data) {
-                  return undefined;
-                }
-
-                const featureSetData = data.get(id)?.data as FeatureSetData;
-                if (!featureSetData) {
-                  return undefined;
-                }
-
+              rowCustomization={(id, { data: featureSetData }) => {
                 return [
                   {
                     column: columns.length - 2,
diff --git a/ui/src/featureSet/featureSet.tsx b/ui/src/featureSet/featureSet.tsx
index dd282dee0..c59ebb039 100644
--- a/ui/src/featureSet/featureSet.tsx
+++ b/ui/src/featureSet/featureSet.tsx
@@ -23,7 +23,7 @@ import { SaveStatus } from "components/saveStatus";
 import { useSimpleDialog } from "components/simpleDialog";
 import { Tabs } from "components/tabs";
 import { useTextInputDialog } from "components/textInputDialog";
-import { TreeGrid, TreeGridData } from "components/treegrid";
+import { TreeGrid, TreeGridId, TreeGridData } from "components/treegrid";
 import { Criteria } from "data/source";
 import { useStudySource } from "data/studySourceContext";
 import { useUnderlaySource } from "data/underlaySourceContext";
@@ -340,19 +340,21 @@ function Preview() {
             true
           );
 
-          const data: TreeGridData = new Map([
-            ["root", { data: {}, children: [] }],
-          ]);
+          const children: TreeGridId[] = [];
+          const rows = new Map();
 
           res.data.forEach((entry, i) => {
-            data.set(i, { data: entry });
-            data.get("root")?.children?.push(i);
+            rows.set(i, { data: entry });
+            children?.push(i);
           });
 
           return {
             id: params.id,
             name: params.name,
-            data: data,
+            data: {
+              rows,
+              children,
+            },
             attributes: params.attributes,
           };
         })
@@ -371,7 +373,7 @@ function Preview() {
                 id: data.name,
                 title: data.name,
                 render: () =>
-                  data.data.get("root")?.children?.length ? (
+                  data.data.children?.length ? (
                     previewOccurrences[i] ? (
                       <PreviewTable
                         occurrence={previewOccurrences[i]}
diff --git a/ui/src/sampleApp/studiesList.tsx b/ui/src/sampleApp/studiesList.tsx
index 1f30e369a..36538ec85 100644
--- a/ui/src/sampleApp/studiesList.tsx
+++ b/ui/src/sampleApp/studiesList.tsx
@@ -6,7 +6,7 @@ import Empty from "components/empty";
 import Loading from "components/loading";
 import { useSimpleDialog } from "components/simpleDialog";
 import { useTextInputDialog } from "components/textInputDialog";
-import { TreeGrid, TreeGridData } from "components/treegrid";
+import { TreeGrid } from "components/treegrid";
 import { useStudySource } from "data/studySourceContext";
 import { DataKey } from "data/types";
 import GridLayout from "layout/gridLayout";
@@ -78,7 +78,7 @@ export function StudiesList() {
 
   const data = useMemo(() => {
     const children: DataKey[] = [];
-    const data: TreeGridData = new Map([["root", { data: {}, children }]]);
+    const rows = new Map();
 
     studiesState.data?.forEach((study) => {
       const key = study.id;
@@ -120,10 +120,13 @@ export function StudiesList() {
           ),
         },
       };
-      data.set(key, item);
+      rows.set(key, item);
     });
 
-    return data;
+    return {
+      rows,
+      children,
+    };
   }, [studySource, studiesState.data]);
 
   return (
@@ -134,7 +137,7 @@ export function StudiesList() {
       <Header />
       <Loading status={studiesState}>
         <GridLayout rows spacing={4}>
-          {!!data?.get("root")?.children?.length ? (
+          {!!data?.children?.length ? (
             <TreeGrid columns={columns} data={data} />
           ) : (
             <Empty
diff --git a/ui/src/sampleApp/studyOverview.tsx b/ui/src/sampleApp/studyOverview.tsx
index 62273fa5f..5aad37d93 100644
--- a/ui/src/sampleApp/studyOverview.tsx
+++ b/ui/src/sampleApp/studyOverview.tsx
@@ -5,7 +5,7 @@ import IconButton from "@mui/material/IconButton";
 import Empty from "components/empty";
 import Loading from "components/loading";
 import { useSimpleDialog } from "components/simpleDialog";
-import { TreeGrid, TreeGridData, TreeGridId } from "components/treegrid";
+import { TreeGrid, TreeGridId } from "components/treegrid";
 import { useStudySource } from "data/studySourceContext";
 import { DataKey } from "data/types";
 import { useStudyId, useUnderlay } from "hooks";
@@ -110,7 +110,7 @@ export function StudyOverview() {
 
   const data = useMemo(() => {
     const children: DataKey[] = [];
-    const data: TreeGridData = new Map([["root", { data: {}, children }]]);
+    const rows = new Map();
 
     artifactsState.data?.forEach((artifact) => {
       const key = `${artifact.type}~${artifact.id}`;
@@ -142,10 +142,13 @@ export function StudyOverview() {
           ),
         },
       };
-      data.set(key, item);
+      rows.set(key, item);
     });
 
-    return data;
+    return {
+      rows,
+      children,
+    };
   }, [studySource, artifactsState.data]);
 
   const navigate = useNavigate();
@@ -200,7 +203,7 @@ export function StudyOverview() {
           </GridLayout>
         </GridBox>
         <Loading status={artifactsState}>
-          {!!data?.get("root")?.children?.length ? (
+          {!!data?.children?.length ? (
             <TreeGrid
               columns={columns}
               data={data}