Treemap looking pretty good

src/TreeMap.tsx

@@ -4,7 +4,7 @@ import * as chromatic from "d3-scale-chromatic";
 import { debounce } from "lodash";
 import { CSSProperties, useEffect, useMemo, useRef, useState} from "react";
 
-import buildNestedData from "./TreeMap/BuildData";
+import buildNestedData, {removeEmptyNodes} from "./TreeMap/BuildData";
 import {formatBytes, formatTime} from "./Util";
 
 // import SimpleD3TreeMap from "./D3TreeMap";
@@ -33,13 +33,13 @@ const svgStyle: CSSProperties = {
 export default function TreeMap({aggregate, data}: Props) {
   const modulesList = aggregate === "time" ? data.modulesByTime : data.modulesByAlloc;
 
-  const nestedData: Tree<TreeNode> = useMemo(() => buildNestedData(aggregate, modulesList), [aggregate, modulesList]);
+  const nestedData: Tree<TreeNode> = useMemo(() => {
+    const tree = buildNestedData(aggregate, modulesList);
+    return removeEmptyNodes(tree);
+  }, [aggregate, modulesList]);
 
   const ref = useRef(null);
-  const [dimensions, setDimensions] = useState<{ height: number, width: number } | null>(() => {
-    console.log("ref.current", ref.current);
-    return null;
-  });
+  const [dimensions, setDimensions] = useState<{ height: number, width: number } | null>(null);
   const elementObserver = useMemo(() => {
     return new ResizeObserver(() => {
       debounce(() => {
@@ -61,14 +61,13 @@ export default function TreeMap({aggregate, data}: Props) {
     };
   }, [ref.current, elementObserver]);
 
-  // Modes: squarify, resquarify, slice, dice, slicedice, binary, circlePack, partition, partition-pivot
-
   return (
     <div ref={ref}
          style={wrapperStyle}>
 
       {dimensions &&
        <D3TreeMap<Tree<TreeNode>>
+         key={aggregate}
          id="myTreeMap"
          width={dimensions.width + paddingPx}
          svgStyle={svgStyle}

src/TreeMap/BuildData.ts

@@ -2,9 +2,11 @@
 import {find, forEach} from "lodash";
 
 
+type ModuleItem = {module: string, alloc: number} | {module: string, time: number};
+
 export default function buildNestedData(
   aggregate: Aggregate,
-  modulesList: Array<{module: string, alloc: number}> | Array<{module: string, time: number}>
+  modulesList: Array<ModuleItem>
 ) {
   const ret: Tree<TreeNode> = {
     name: "Top-level",
@@ -14,26 +16,7 @@ export default function buildNestedData(
   };
 
   for (const m of modulesList) {
-    let current = ret;
-    let soFar = "";
-    const parts = m.module.split(".");
-    for (let i = 0; i < parts.length; i += 1) {
-      const part = parts[i];
-      soFar += (i === 0 ? part : "." + part);
-
-      let item = find(current.children, (x) => x.part === part);
-      if (!item) {
-        item = {
-          name: soFar,
-          part,
-          value: 0,
-          children: [],
-        };
-        current.children.push(item);
-      }
-      current = item;
-      current.value += (aggregate === "time" ? m["time"]: m["alloc"]);
-    }
+    addModuleToTree(aggregate, m, ret);
   }
   forEach(ret.children, (x) => {
     ret.value += x.value;
@@ -43,3 +26,44 @@ export default function buildNestedData(
 
   return ret;
 }
+
+
+function addModuleToTree(aggregate: Aggregate, m: ModuleItem, ret: Tree<TreeNode>) {
+  let current = ret;
+  let soFar = "";
+  const parts = m.module.split(".");
+
+  for (let i = 0; i < parts.length; i += 1) {
+    const part = parts[i];
+    soFar += (i === 0 ? part : "." + part);
+
+    let item = find(current.children, (x) => x.part === part);
+    if (!item) {
+      item = {
+        name: soFar,
+        part,
+        value: 0,
+        children: [],
+      };
+      current.children.push(item);
+    }
+    current = item;
+
+    // Add the value to the bottommost node
+    if (i === parts.length - 1) {
+      current.value += (aggregate === "time" ? m["time"]: m["alloc"]);
+    }
+  }
+}
+
+export function removeEmptyNodes(tree: Tree<TreeNode>): Tree<TreeNode> {
+  for (let i = 0; i < tree.children.length; i += 1) {
+    tree.children[i] = removeEmptyNodes(tree.children[i]);
+  }
+
+  if (tree.value === 0 && tree.children.length === 1) {
+    return tree.children[0];
+  } else {
+    return tree;
+  }
+}