Add symbol mapping & diff UI improvements

webview/views/DiffView.module.css

@@ -0,0 +1,23 @@
+.content {
+  flex: 1 1 0;
+  display: flex;
+  flex-flow: row;
+}
+
+.column {
+  flex: 1 1 0;
+  margin: 0;
+  padding: 0;
+  overflow: auto;
+}
+
+.no-object {
+  color: var(--color-blue);
+
+  font-family: var(--code-font-family);
+  font-weight: var(--code-font-weight);
+  font-size: var(--code-font-size);
+  font-variant-ligatures: var(--code-font-variant-ligatures);
+  text-wrap: nowrap;
+  white-space: pre;
+}

webview/views/FunctionView.module.css

@@ -5,8 +5,6 @@
 .instruction-row {
   display: flex;
   flex-direction: row;
-  gap: 0.5em;
-  height: var(--list-row-height);
 }
 
 .instruction-cell {

webview/views/FunctionView.tsx

@@ -4,26 +4,20 @@ import styles from './FunctionView.module.css';
 import clsx from 'clsx';
 import memoizeOne from 'memoize-one';
 import { type diff, display } from 'objdiff-wasm';
-import { memo, useMemo } from 'react';
-import AutoSizer from 'react-virtualized-auto-sizer';
+import { memo, useCallback, useMemo } from 'react';
 import { FixedSizeList, areEqual } from 'react-window';
 import type { ListChildComponentProps } from 'react-window';
 import { useShallow } from 'zustand/react/shallow';
 import { createContextMenu, renderContextItems } from '../common/ContextMenu';
-import TooltipShared from '../common/TooltipShared';
-import {
-  buildDiffConfig,
-  runBuild,
-  useAppStore,
-  useExtensionStore,
-} from '../state';
+import { createTooltip } from '../common/TooltipShared';
+import { buildDiffConfig, useAppStore, useExtensionStore } from '../state';
 import {
   type HighlightState,
   highlightColumn,
   highlightMatches,
   updateHighlight,
 } from '../util/highlight';
-import { percentClass, useFontSize } from '../util/util';
+import { useFontSize } from '../util/util';
 
 const ROTATION_CLASSES = [
   styles.rotation0,
@@ -37,10 +31,20 @@ const ROTATION_CLASSES = [
   styles.rotation8,
 ];
 
-const { ContextMenuProvider, useContextMenu } = createContextMenu<{
+export type InstructionTooltipContent = {
   column: number;
   row: number;
-}>();
+};
+
+export const {
+  Tooltip: InstructionTooltip,
+  useTooltip: useInstructionTooltip,
+} = createTooltip<InstructionTooltipContent>();
+
+export const {
+  ContextMenuProvider: InstructionContextMenuProvider,
+  useContextMenu: useInstructionContextMenu,
+} = createContextMenu<InstructionTooltipContent>();
 
 const AsmCell = ({
   obj,
@@ -53,15 +57,28 @@ const AsmCell = ({
 }: {
   obj: diff.ObjectDiff | undefined;
   config: diff.DiffConfig;
-  symbol: display.SectionDisplaySymbol | null;
+  symbol: display.SymbolRef | null;
   row: number;
   column: number;
   highlight: HighlightState;
   setHighlight: (highlight: HighlightState) => void;
 }) => {
-  const onContextMenu = useContextMenu();
+  const onContextMenu = useInstructionContextMenu();
+  const tooltipContent: InstructionTooltipContent = useMemo(
+    () => ({
+      column,
+      row,
+    }),
+    [column, row],
+  );
+  const tooltipProps = useInstructionTooltip(tooltipContent);
+  const onContextMenuMemo = useCallback(
+    (e: React.MouseEvent<HTMLElement>) => onContextMenu(e, tooltipContent),
+    [onContextMenu, tooltipContent],
+  );
+
   if (!obj || !symbol) {
-    return <div className={styles.instructionCell} />;
+    return null;
   }
 
   const highlight = highlightColumn(highlightState, column);
@@ -198,33 +215,24 @@ const AsmCell = ({
     return <div className={clsx(classes)} />;
   }
 
-  const tooltipContent: InstructionTooltipContent = { column, row };
   return (
     <div
       className={clsx(classes)}
-      data-tooltip-id="instruction-tooltip"
-      data-tooltip-content={JSON.stringify(tooltipContent)}
-      onContextMenu={(e) => onContextMenu(e, { column, row })}
+      onContextMenu={onContextMenuMemo}
+      {...tooltipProps}
     >
       {out}
     </div>
   );
 };
 
-type InstructionTooltipContent = {
-  column: number;
-  row: number;
-};
-
 type ItemData = {
   itemCount: number;
   symbolName: string;
   result: diff.DiffResult;
   config: diff.DiffConfig;
   matchPercent?: number;
-  left: display.SectionDisplaySymbol | null;
   leftSymbol: display.SymbolDisplay | null;
-  right: display.SectionDisplaySymbol | null;
   rightSymbol: display.SymbolDisplay | null;
   highlight: HighlightState;
   setHighlight: (highlight: HighlightState) => void;
@@ -234,7 +242,7 @@ const AsmRow = memo(
   ({
     index,
     style,
-    data: { result, config, left, right, highlight, setHighlight },
+    data: { result, config, leftSymbol, rightSymbol, highlight, setHighlight },
   }: ListChildComponentProps<ItemData>) => {
     return (
       <div
@@ -254,7 +262,7 @@ const AsmRow = memo(
         <AsmCell
           obj={result.left}
           config={config}
-          symbol={left}
+          symbol={leftSymbol?.info.id ?? null}
           row={index}
           column={0}
           highlight={highlight}
@@ -263,7 +271,7 @@ const AsmRow = memo(
         <AsmCell
           obj={result.right}
           config={config}
-          symbol={right}
+          symbol={rightSymbol?.info.id ?? null}
           row={index}
           column={1}
           highlight={highlight}
@@ -278,20 +286,16 @@ const AsmRow = memo(
 const createItemData = memoizeOne(
   (
     result: diff.DiffResult,
-    left: display.SectionDisplaySymbol | null,
-    right: display.SectionDisplaySymbol | null,
+    leftSymbol: display.SymbolDisplay | null,
+    rightSymbol: display.SymbolDisplay | null,
     highlight: HighlightState,
     setHighlight: (highlight: HighlightState) => void,
   ): ItemData => {
-    const leftSymbol = left ? display.displaySymbol(result.left!, left) : null;
-    const rightSymbol = right
-      ? display.displaySymbol(result.right!, right)
-      : null;
     const itemCount = Math.max(
       leftSymbol?.rowCount || 0,
       rightSymbol?.rowCount || 0,
     );
-    const symbolName = leftSymbol?.name || rightSymbol?.name || '';
+    const symbolName = leftSymbol?.info.name || rightSymbol?.info.name || '';
     const config = buildDiffConfig(null);
     const matchPercent = rightSymbol?.matchPercent;
     return {
@@ -300,9 +304,7 @@ const createItemData = memoizeOne(
       result,
       config,
       matchPercent,
-      left,
       leftSymbol,
-      right,
       rightSymbol,
       highlight,
       setHighlight,
@@ -322,7 +324,7 @@ const SymbolLabel = ({
       </span>
     );
   }
-  const displayName = symbol.demangledName || symbol.name;
+  const displayName = symbol.info.demangledName || symbol.info.name;
   return (
     <span
       className={clsx(headerStyles.label, headerStyles.emphasized)}
@@ -333,36 +335,35 @@ const SymbolLabel = ({
   );
 };
 
-const FunctionView = ({
+export const InstructionList = ({
+  height,
+  width,
   diff,
-  left,
-  right,
+  leftSymbol,
+  rightSymbol,
 }: {
+  height: number;
+  width: number;
   diff: diff.DiffResult;
-  left: display.SectionDisplaySymbol | null;
-  right: display.SectionDisplaySymbol | null;
+  leftSymbol: display.SymbolDisplay | null;
+  rightSymbol: display.SymbolDisplay | null;
 }) => {
-  const { buildRunning, currentUnit, lastBuilt, hasProjectConfig } =
-    useExtensionStore(
-      useShallow((state) => ({
-        buildRunning: state.buildRunning,
-        currentUnit: state.currentUnit,
-        lastBuilt: state.lastBuilt,
-        hasProjectConfig: state.projectConfig != null,
-      })),
-    );
+  const currentUnit = useExtensionStore((state) => state.currentUnit);
+  const { highlight, setSymbolScrollOffset, setHighlight } = useAppStore(
+    useShallow((state) => ({
+      highlight: state.highlight,
+      setSymbolScrollOffset: state.setSymbolScrollOffset,
+      setHighlight: state.setHighlight,
+    })),
+  );
+  const itemData = createItemData(
+    diff,
+    leftSymbol,
+    rightSymbol,
+    highlight,
+    setHighlight,
+  );
   const currentUnitName = currentUnit?.name || '';
-  const { highlight, setSelectedSymbol, setSymbolScrollOffset, setHighlight } =
-    useAppStore(
-      useShallow((state) => ({
-        highlight: state.highlight,
-        setSelectedSymbol: state.setSelectedSymbol,
-        setSymbolScrollOffset: state.setSymbolScrollOffset,
-        setHighlight: state.setHighlight,
-      })),
-    );
-
-  const itemData = createItemData(diff, left, right, highlight, setHighlight);
   const initialScrollOffset = useMemo(
     () =>
       useAppStore.getState().getUnitState(currentUnitName).symbolScrollOffsets[
@@ -370,140 +371,25 @@ const FunctionView = ({
       ] || 0,
     [currentUnitName, itemData.symbolName],
   );
-
   const itemSize = useFontSize() * 1.33;
   return (
-    <>
-      <div className={headerStyles.header}>
-        <div className={headerStyles.column}>
-          <div className={headerStyles.row}>
-            <button title="Back" onClick={() => setSelectedSymbol(null, null)}>
-              <span className="codicon codicon-chevron-left" />
-            </button>
-          </div>
-          <div className={headerStyles.row}>
-            <SymbolLabel symbol={itemData.leftSymbol} />
-          </div>
-        </div>
-        <div className={headerStyles.column}>
-          <div className={headerStyles.row}>
-            {hasProjectConfig && (
-              <button
-                title="Build"
-                onClick={() => runBuild()}
-                disabled={buildRunning}
-              >
-                <span className="codicon codicon-refresh" />
-              </button>
-            )}
-            {lastBuilt && (
-              <span className={headerStyles.label}>
-                Last built: {new Date(lastBuilt).toLocaleTimeString('en-US')}
-              </span>
-            )}
-          </div>
-          <div className={headerStyles.row}>
-            {itemData.matchPercent !== undefined && (
-              <>
-                <span
-                  className={clsx(
-                    headerStyles.label,
-                    percentClass(itemData.matchPercent),
-                  )}
-                >
-                  {Math.floor(itemData.matchPercent).toFixed(0)}%
-                </span>
-                {' | '}
-              </>
-            )}
-            <SymbolLabel symbol={itemData.rightSymbol} />
-          </div>
-        </div>
-      </div>
-      <div className={styles.instructionList}>
-        <ContextMenuProvider
-          render={({ data }, close) => {
-            let obj: diff.ObjectDiff | undefined;
-            let symbol: display.SectionDisplaySymbol | undefined;
-            switch (data.column) {
-              case 0:
-                obj = diff.left;
-                symbol = itemData.left ?? undefined;
-                break;
-              case 1:
-                obj = diff.right;
-                symbol = itemData.right ?? undefined;
-                break;
-              default:
-                break;
-            }
-            if (!obj || !symbol) {
-              return null;
-            }
-            const items = display.instructionContext(
-              obj,
-              symbol,
-              data.row,
-              itemData.config,
-            );
-            return renderContextItems(items, close);
-          }}
-        >
-          <AutoSizer>
-            {({ height, width }) => (
-              <FixedSizeList
-                height={height}
-                itemCount={itemData.itemCount}
-                itemSize={itemSize}
-                width={width}
-                itemData={itemData}
-                overscanCount={20}
-                onScroll={(e) => {
-                  setSymbolScrollOffset(
-                    currentUnitName,
-                    itemData.symbolName,
-                    e.scrollOffset,
-                  );
-                }}
-                initialScrollOffset={initialScrollOffset}
-              >
-                {AsmRow}
-              </FixedSizeList>
-            )}
-          </AutoSizer>
-        </ContextMenuProvider>
-      </div>
-      <TooltipShared
-        id="instruction-tooltip"
-        callback={(content) => {
-          const data: InstructionTooltipContent = JSON.parse(content);
-          let obj: diff.ObjectDiff | undefined;
-          let symbol: display.SectionDisplaySymbol | undefined;
-          switch (data.column) {
-            case 0:
-              obj = diff.left;
-              symbol = itemData.left ?? undefined;
-              break;
-            case 1:
-              obj = diff.right;
-              symbol = itemData.right ?? undefined;
-              break;
-            default:
-              break;
-          }
-          if (!obj || !symbol) {
-            return null;
-          }
-          return display.instructionHover(
-            obj,
-            symbol,
-            data.row,
-            itemData.config,
-          );
-        }}
-      />
-    </>
+    <FixedSizeList
+      height={height}
+      itemCount={itemData.itemCount}
+      itemSize={itemSize}
+      width={width}
+      itemData={itemData}
+      overscanCount={20}
+      onScroll={(e) => {
+        setSymbolScrollOffset(
+          currentUnitName,
+          itemData.symbolName,
+          e.scrollOffset,
+        );
+      }}
+      initialScrollOffset={initialScrollOffset}
+    >
+      {AsmRow}
+    </FixedSizeList>
   );
 };
-
-export default FunctionView;

webview/views/SymbolsView.module.css

@@ -11,41 +11,6 @@
   padding: 0;
 }
 
-.symbol-list-row {
-  cursor: pointer;
-  user-select: none;
-  height: var(--list-row-height);
-
-  font-family: var(--code-font-family);
-  font-weight: var(--code-font-weight);
-  font-size: var(--code-font-size);
-  font-variant-ligatures: var(--code-font-variant-ligatures);
-  text-wrap: nowrap;
-  white-space: pre;
-
-  &:hover {
-    background-color: var(--list-row-hover-background);
-  }
-}
-
-.section {
-  padding-left: 0.5em;
-
-  &::before {
-    content: "▼ ";
-  }
-  &.collapsed {
-    opacity: 0.75;
-    &::before {
-      content: "▶ ";
-    }
-  }
-}
-
-.symbol {
-  padding-left: 2em;
-}
-
 .flag-local {
   color: inherit;
 }
@@ -62,6 +27,10 @@
   color: var(--color-blue);
 }
 
+.flag-hidden {
+  color: var(--color-muted);
+}
+
 .symbol-name {
   color: var(--color-bright);
 }
@@ -76,3 +45,19 @@
   text-wrap: nowrap;
   white-space: pre;
 }
+
+.selected {
+  background-color: var(--list-row-selection-background) !important;
+}
+
+.highlighted {
+  background-color: var(--list-row-highlight-background) !important;
+}
+
+.selected.highlighted {
+  background-color: color-mix(
+    in srgb,
+    var(--list-row-selection-background) 50%,
+    var(--list-row-highlight-background) 50%
+  ) !important;
+}

webview/views/TreeView.module.css

@@ -0,0 +1,42 @@
+.row {
+  display: flex;
+  cursor: pointer;
+  user-select: none;
+  height: var(--list-row-height);
+  padding-left: 0.5em;
+
+  font-family: var(--code-font-family);
+  font-weight: var(--code-font-weight);
+  font-size: var(--code-font-size);
+  font-variant-ligatures: var(--code-font-variant-ligatures);
+  text-wrap: nowrap;
+  white-space: pre;
+
+  &:hover {
+    background-color: var(--list-row-hover-background);
+  }
+}
+
+.indent {
+  display: inline-block;
+  height: var(--list-row-height);
+  width: 0;
+  border-right: 1px solid var(--panel-separator);
+  margin-left: 0.5em;
+  margin-right: 0.5em;
+}
+
+.indent-highlighted {
+  border-right-color: var(--color-muted);
+}
+
+.toggle {
+  display: inline-block;
+  height: var(--list-row-height);
+  width: 1em;
+  margin-right: 0.25em;
+}
+
+.collapsed {
+  color: color-mix(in srgb, var(--foreground) 75%, transparent);
+}

webview/views/TreeView.tsx

@@ -0,0 +1,215 @@
+import clsx from 'clsx';
+import type { ListChildComponentProps } from 'react-window';
+import styles from './TreeView.module.css';
+
+export type NodeInner<T> = {
+  id: string;
+  indent: number;
+  path: string[];
+  data: T;
+};
+
+export type BranchNode<T> = NodeInner<T> & {
+  type: 'branch';
+  collapsed: boolean;
+};
+
+export type LeafNode<T> = NodeInner<T> & {
+  type: 'leaf';
+};
+
+export type Node<B, L> = BranchNode<B> | LeafNode<L>;
+
+export type TreeData<B, L> = {
+  leafCount: number;
+  nodes: Node<B, L>[];
+  highlightedPath: string | null;
+  setHighlightedPath: (id: string | null) => void;
+};
+
+export type TreeRowProps<B, L> = ListChildComponentProps<TreeData<B, L>> & {
+  rowProps?: React.HTMLProps<HTMLDivElement> | null;
+  render: (item: Node<B, L>) => React.ReactNode;
+  getClasses?: (item: Node<B, L>) => string[];
+  onLeafClick?: (item: LeafNode<L>) => void;
+  onHover?: (item: Node<B, L>) => void;
+  setBranchCollapsed?: (id: string, collapsed: boolean) => void;
+};
+
+export function TreeRow<B, L>({
+  index,
+  style,
+  data: { nodes, highlightedPath, setHighlightedPath },
+  rowProps,
+  render,
+  getClasses,
+  onLeafClick,
+  onHover,
+  setBranchCollapsed,
+}: TreeRowProps<B, L>) {
+  const node = nodes[index];
+  const classes = [styles.row];
+  if (node.type === 'branch' && node.collapsed) {
+    classes.push(styles.collapsed);
+  }
+  classes.push(...(getClasses?.(node) ?? []));
+  const indentItems = [];
+  for (let i = 0; i < node.indent; i++) {
+    indentItems.push(
+      <span
+        key={i}
+        className={clsx(
+          styles.indent,
+          node.path[i] === highlightedPath && styles.indentHighlighted,
+        )}
+      />,
+    );
+  }
+  if (node.type === 'branch') {
+    indentItems.push(
+      <span
+        key="toggle"
+        className={clsx(
+          styles.toggle,
+          'codicon',
+          node.collapsed ? 'codicon-chevron-right' : 'codicon-chevron-down',
+        )}
+      />,
+    );
+  }
+  return (
+    <div
+      className={clsx(classes)}
+      style={style}
+      {...rowProps}
+      onClick={() => {
+        if (node.type === 'leaf') {
+          onLeafClick?.(node);
+        } else {
+          const collapsed = !node.collapsed;
+          setBranchCollapsed?.(node.id, collapsed);
+          setHighlightedPath(
+            collapsed ? node.path[node.path.length - 1] : node.id,
+          );
+        }
+      }}
+      onMouseEnter={() => {
+        onHover?.(node);
+        setHighlightedPath(
+          node.type === 'leaf' || node.collapsed
+            ? node.path[node.path.length - 1]
+            : node.id,
+        );
+      }}
+    >
+      {indentItems}
+      {render(node)}
+    </div>
+  );
+}
+
+type NodeWithChildren<B, L> = BranchNode<B> & {
+  children: (NodeWithChildren<B, L> | LeafNode<L>)[];
+};
+
+export type SimpleTreeNodeData = { label: string };
+export type SimpleTreeData<L> = TreeData<
+  SimpleTreeNodeData,
+  L & SimpleTreeNodeData
+>;
+
+// Build a simple tree structure by splitting the path of each item
+// into components (splitting on '/') and creating a tree node for
+// each component.
+export function buildSimpleTree<L>(
+  items: L[],
+  getPath: (item: L) => string,
+  collapsed: Record<string, boolean>,
+  highlightedPath: string | null,
+  setHighlightedPath: (id: string | null) => void,
+): SimpleTreeData<L> {
+  type SimpleTreeNode = NodeWithChildren<
+    SimpleTreeNodeData,
+    L & SimpleTreeNodeData
+  >;
+  const map = new Map<string, SimpleTreeNode>();
+  const rootNodes = [];
+  for (const item of items) {
+    const path = getPath(item);
+    const split = path.split('/');
+    let parent: SimpleTreeNode | null = null;
+    for (let i = 0; i < split.length - 1; i++) {
+      const name = split[i];
+      const dirPath = split.slice(0, i + 1);
+      const key = dirPath.join('/');
+      let node = map.get(key);
+      if (!node) {
+        node = {
+          type: 'branch',
+          id: key,
+          indent: i,
+          collapsed: !!collapsed[key],
+          children: [],
+          path: buildPath(dirPath),
+          data: { label: name },
+        };
+        if (parent) {
+          parent.children.push(node);
+        } else {
+          rootNodes.push(node);
+        }
+        map.set(key, node);
+      }
+      parent = node;
+    }
+    const node: LeafNode<L & SimpleTreeNodeData> = {
+      type: 'leaf',
+      id: path,
+      indent: split.length - 1,
+      path: buildPath(split),
+      data: {
+        ...item,
+        label: split[split.length - 1],
+      },
+    };
+    if (parent) {
+      parent.children.push(node);
+    } else {
+      rootNodes.push(node);
+    }
+  }
+  const nodes: SimpleTreeData<L>['nodes'] = [];
+  for (const node of rootNodes) {
+    pushNodes(node, nodes);
+  }
+  return {
+    leafCount: items.length,
+    nodes,
+    highlightedPath,
+    setHighlightedPath,
+  };
+}
+
+function pushNodes<B, L>(
+  item: NodeWithChildren<B, L> | LeafNode<L>,
+  out: Node<B, L>[],
+) {
+  if (item.type === 'branch') {
+    out.push(item);
+    if (!item.collapsed) {
+      for (const child of item.children) {
+        pushNodes(child, out);
+      }
+    }
+  } else if (item.type === 'leaf') {
+    out.push(item);
+  }
+}
+
+function buildPath(split: string[]) {
+  const path = [];
+  for (let i = 0; i < split.length - 1; i++) {
+    path.push(split.slice(0, i + 1).join('/'));
+  }
+  return path;
+}

webview/views/UnitsView.module.css

@@ -5,49 +5,8 @@
   overflow: hidden;
 }
 
-.row {
-  display: flex;
-  cursor: pointer;
-  user-select: none;
-  height: var(--list-row-height);
-  padding-left: 0.5em;
-
-  font-family: var(--code-font-family);
-  font-weight: var(--code-font-weight);
-  font-size: var(--code-font-size);
-  font-variant-ligatures: var(--code-font-variant-ligatures);
-  text-wrap: nowrap;
-  white-space: pre;
-
-  &:hover {
-    background-color: var(--list-row-hover-background);
-  }
-}
-
-.indent {
-  display: inline-block;
-  height: var(--list-row-height);
-  width: 0;
-  border-right: 1px solid var(--panel-separator);
-  margin-left: 0.5em;
-  margin-right: 0.5em;
-}
-
-.indent-highlighted {
-  border-right-color: var(--color-muted);
-}
-
-.toggle {
-  display: inline-block;
-  height: var(--list-row-height);
-  width: 1em;
-  margin-right: 0.25em;
-}
-
-.unit {
-  .label {
-    margin-left: 0.4em;
-  }
+.unit-label {
+  margin-left: 0.4em;
 }
 
 .complete {
@@ -57,7 +16,3 @@
 .incomplete {
   color: var(--color-red);
 }
-
-.collapsed {
-  color: color-mix(in srgb, var(--foreground) 75%, transparent);
-}

webview/views/UnitsView.tsx

@@ -1,6 +1,3 @@
-import styles from './UnitsView.module.css';
-
-import clsx from 'clsx';
 import memoizeOne from 'memoize-one';
 import { memo, useMemo, useState } from 'react';
 import AutoSizer from 'react-virtualized-auto-sizer';
@@ -19,135 +16,40 @@ import {
   useExtensionStore,
 } from '../state';
 import { useFontSize } from '../util/util';
+import { type SimpleTreeData, TreeRow, buildSimpleTree } from './TreeView';
+import styles from './UnitsView.module.css';
 
-type DirectoryItem = {
-  type: 'directory';
-  indent: number;
-  label: string;
-  id: string;
-  collapsed: boolean;
-  path: string[];
-};
-
-type UnitItem = {
-  type: 'unit';
-  indent: number;
-  label: string;
-  id: string;
-  unit: Unit;
-  path: string[];
-};
-
-type Item = DirectoryItem | UnitItem;
-
-type ItemData = {
-  unitsCount: number;
-  items: Item[];
-  highlightedPath: string | null;
-  setHighlightedPath: (id: string | null) => void;
-};
-
-const UnitRow = memo(
-  ({
-    index,
-    style,
-    data: { items, highlightedPath, setHighlightedPath },
-  }: ListChildComponentProps<ItemData>) => {
-    const setCollapsedUnit = useAppStore((state) => state.setCollapsedUnit);
-    const item = items[index];
-    const classes = [styles.row];
-    if (item.type === 'unit') {
-      classes.push(styles.unit);
-      if (item.unit.metadata?.complete !== undefined) {
-        if (item.unit.metadata.complete) {
-          classes.push(styles.complete);
-        } else {
-          classes.push(styles.incomplete);
-        }
-      }
-    } else {
-      // classes.push(styles.directory);
-      if (item.collapsed) {
-        classes.push(styles.collapsed);
-      }
-    }
-    const indentItems = [];
-    for (let i = 0; i < item.indent; i++) {
-      indentItems.push(
-        <span
-          key={i}
-          className={clsx(
-            styles.indent,
-            item.path[i] === highlightedPath && styles.indentHighlighted,
-          )}
-        />,
-      );
-    }
-    if (item.type === 'directory') {
-      indentItems.push(
-        <span
-          key="toggle"
-          className={clsx(
-            styles.toggle,
-            'codicon',
-            item.collapsed ? 'codicon-chevron-right' : 'codicon-chevron-down',
-          )}
-        />,
-      );
-    }
-    return (
-      <div
-        className={clsx(classes)}
-        style={style}
-        onClick={() => {
-          if (item.type === 'unit') {
-            setCurrentUnit(item.unit);
-          } else {
-            const collapsed = !item.collapsed;
-            setCollapsedUnit(item.id, collapsed);
-            setHighlightedPath(
-              collapsed ? item.path[item.path.length - 1] : item.id,
-            );
+const UnitRow = memo((props: ListChildComponentProps<SimpleTreeData<Unit>>) => {
+  const setCollapsedUnit = useAppStore((state) => state.setCollapsedUnit);
+  return (
+    <TreeRow
+      {...props}
+      getClasses={(item) => {
+        if (
+          item.type === 'leaf' &&
+          item.data.metadata?.complete !== undefined
+        ) {
+          if (item.data.metadata.complete) {
+            return [styles.complete];
           }
-        }}
-        onMouseEnter={() => {
-          setHighlightedPath(
-            item.type === 'unit' || item.collapsed
-              ? item.path[item.path.length - 1]
-              : item.id,
-          );
-        }}
-      >
-        {indentItems}
-        <span className={styles.label}>{item.label}</span>
-      </div>
-    );
-  },
-  areEqual,
-);
-
-type TreeItem = DirectoryItem & { children: (TreeItem | UnitItem)[] };
-
-function pushTreeItems(item: TreeItem | UnitItem, out: Item[]) {
-  if (item.type === 'directory') {
-    out.push(item);
-    if (!item.collapsed) {
-      for (const child of item.children) {
-        pushTreeItems(child, out);
-      }
-    }
-  } else if (item.type === 'unit') {
-    out.push(item);
-  }
-}
-
-const buildPath = (split: string[]) => {
-  const path = [];
-  for (let i = 0; i < split.length - 1; i++) {
-    path.push(split.slice(0, i + 1).join('/'));
-  }
-  return path;
-};
+          return [styles.incomplete];
+        }
+        return [];
+      }}
+      onLeafClick={(item) => {
+        setCurrentUnit(item.data);
+      }}
+      setBranchCollapsed={setCollapsedUnit}
+      render={(item) => {
+        return (
+          <span className={item.type === 'leaf' ? styles.unitLabel : undefined}>
+            {item.data.label}
+          </span>
+        );
+      }}
+    />
+  );
+}, areEqual);
 
 const createItemData = memoizeOne(
   (
@@ -155,62 +57,14 @@ const createItemData = memoizeOne(
     collapsedUnits: Record<string, boolean>,
     highlightedPath: string | null,
     setHighlightedPath: (id: string | null) => void,
-  ): ItemData => {
-    const units = config?.units ?? [];
-    const map = new Map<string, TreeItem>();
-    const rootItems = [];
-    for (const unit of units) {
-      const path = unit.name || '';
-      const split = path.split('/');
-      let parent: TreeItem | null = null;
-      for (let i = 0; i < split.length - 1; i++) {
-        const name = split[i];
-        const dirPath = split.slice(0, i + 1);
-        const key = dirPath.join('/');
-        let item = map.get(key);
-        if (!item) {
-          item = {
-            type: 'directory',
-            indent: i,
-            label: name,
-            id: key,
-            collapsed: !!collapsedUnits[key],
-            children: [],
-            path: buildPath(dirPath),
-          };
-          if (parent) {
-            parent.children.push(item);
-          } else {
-            rootItems.push(item);
-          }
-          map.set(key, item);
-        }
-        parent = item;
-      }
-      const unitItem: UnitItem = {
-        type: 'unit',
-        indent: split.length - 1,
-        label: split[split.length - 1],
-        id: path,
-        unit,
-        path: buildPath(split),
-      };
-      if (parent) {
-        parent.children.push(unitItem);
-      } else {
-        rootItems.push(unitItem);
-      }
-    }
-    const items: Item[] = [];
-    for (const item of rootItems) {
-      pushTreeItems(item, items);
-    }
-    return {
-      unitsCount: units.length,
-      items,
+  ): SimpleTreeData<Unit> => {
+    return buildSimpleTree(
+      config?.units ?? [],
+      (unit) => unit.name || '',
+      collapsedUnits,
       highlightedPath,
       setHighlightedPath,
-    };
+    );
   },
 );
 
@@ -250,7 +104,7 @@ const UnitsView = () => {
           </div>
           <div className={headerStyles.row}>
             <span className={headerStyles.label}>
-              {itemData.unitsCount} unit{itemData.unitsCount === 1 ? '' : 's'}
+              {itemData.leafCount} unit{itemData.leafCount === 1 ? '' : 's'}
             </span>
           </div>
         </div>
@@ -260,7 +114,7 @@ const UnitsView = () => {
           {({ height, width }) => (
             <FixedSizeList
               height={height}
-              itemCount={itemData.items.length}
+              itemCount={itemData.nodes.length}
               itemSize={itemSize}
               width={width}
               itemData={itemData}