added proper hierarchy from parsed data

packages/mermaid/src/diagrams/treemap/parser.ts

@@ -4,6 +4,7 @@ import { log } from '../../logger.js';
 import { populateCommonDb } from '../common/populateCommonDb.js';
 import { db } from './db.js';
 import type { TreemapNode } from './types.js';
+import { buildHierarchy } from './utils.js';
 
 /**
  * Populates the database with data from the Treemap AST
@@ -12,11 +13,8 @@ import type { TreemapNode } from './types.js';
 const populate = (ast: any) => {
   populateCommonDb(ast, db);
 
-  // Process rows
-  let lastLevel = 0;
-  let lastNode: TreemapNode | undefined;
-
-  // Process each row in the treemap, building the node hierarchy
+  const items = [];
+  // Extract data from each row in the treemap
   for (const row of ast.TreemapRows || []) {
     const item = row.item;
     if (!item) {
@@ -25,57 +23,26 @@ const populate = (ast: any) => {
 
     const level = row.indent ? parseInt(row.indent) : 0;
     const name = getItemName(item);
-
-    // Create the node
-    const node: TreemapNode = {
-      name,
-      children: [],
-    };
-
-    // If it's a leaf node, add the value
-    if (item.$type === 'Leaf') {
-      node.value = item.value;
-    }
-
-    // Add to the right place in hierarchy
-    if (level === 0) {
-      // Root node
-      db.addNode(node, level);
-    } else if (level > lastLevel) {
-      // Child of the last node
-      if (lastNode) {
-        lastNode.children = lastNode.children || [];
-        lastNode.children.push(node);
-        node.parent = lastNode;
-      }
-      db.addNode(node, level);
-    } else if (level === lastLevel) {
-      // Sibling of the last node
-      if (lastNode?.parent) {
-        lastNode.parent.children = lastNode.parent.children || [];
-        lastNode.parent.children.push(node);
-        node.parent = lastNode.parent;
-      }
-      db.addNode(node, level);
-    } else if (level < lastLevel) {
-      // Go up in the hierarchy
-      let parent = lastNode ? lastNode.parent : undefined;
-      for (let i = lastLevel; i > level; i--) {
-        if (parent) {
-          parent = parent.parent;
-        }
-      }
-      if (parent) {
-        parent.children = parent.children || [];
-        parent.children.push(node);
-        node.parent = parent;
-      }
-      db.addNode(node, level);
-    }
-
-    lastLevel = level;
-    lastNode = node;
+    const itemData = { level, name, type: item.$type, value: item.value };
+    items.push(itemData);
   }
+
+  // Convert flat structure to hierarchical
+  const hierarchyNodes = buildHierarchy(items);
+
+  // Add all nodes to the database
+  const addNodesRecursively = (nodes: TreemapNode[], level: number) => {
+    for (const node of nodes) {
+      db.addNode(node, level);
+      if (node.children && node.children.length > 0) {
+        addNodesRecursively(node.children, level + 1);
+      }
+    }
+  };
+
+  addNodesRecursively(hierarchyNodes, 0);
+
+  log.debug('Processed items:', items);
 };
 
 /**

packages/mermaid/src/diagrams/treemap/renderer.ts

@@ -3,10 +3,9 @@ import type { DiagramRenderer, DrawDefinition } from '../../diagram-api/types.js
 import { selectSvgElement } from '../../rendering-util/selectSvgElement.js';
 import { configureSvgSize } from '../../setupGraphViewbox.js';
 import type { TreemapDB, TreemapNode } from './types.js';
+import { scaleOrdinal, treemap, hierarchy, format } from 'd3';
 
-const DEFAULT_PADDING = 10;
-const DEFAULT_NODE_WIDTH = 100;
-const DEFAULT_NODE_HEIGHT = 40;
+const DEFAULT_PADDING = 1;
 
 /**
  * Draws the treemap diagram
@@ -23,136 +22,144 @@ const draw: DrawDefinition = (_text, id, _version, diagram: Diagram) => {
   }
 
   const svg = selectSvgElement(id);
-
-  // Calculate the size of the treemap
-  const { width, height } = calculateTreemapSize(root, config);
+  // Use config dimensions or defaults
+  const width = config.nodeWidth ? config.nodeWidth * 10 : 960;
+  const height = config.nodeHeight ? config.nodeHeight * 10 : 500;
   const titleHeight = title ? 30 : 0;
-  const svgWidth = width + padding * 2;
-  const svgHeight = height + padding * 2 + titleHeight;
+  const svgWidth = width;
+  const svgHeight = height + titleHeight;
 
   // Set the SVG size
   svg.attr('viewBox', `0 0 ${svgWidth} ${svgHeight}`);
   configureSvgSize(svg, svgHeight, svgWidth, config.useMaxWidth);
 
+  // Format for displaying values
+  const valueFormat = format(',d');
+
+  // Create color scale
+  const colorScale = scaleOrdinal<string>().range([
+    '#8dd3c7',
+    '#ffffb3',
+    '#bebada',
+    '#fb8072',
+    '#80b1d3',
+    '#fdb462',
+    '#b3de69',
+    '#fccde5',
+    '#d9d9d9',
+    '#bc80bd',
+  ]);
+
   // Create a container group to hold all elements
-  const g = svg.append('g').attr('transform', `translate(${padding}, ${padding + titleHeight})`);
+  const g = svg.append('g').attr('transform', `translate(0, ${titleHeight})`);
 
   // Draw the title if it exists
   if (title) {
     svg
       .append('text')
       .attr('x', svgWidth / 2)
-      .attr('y', padding + titleHeight / 2)
+      .attr('y', titleHeight / 2)
       .attr('class', 'treemapTitle')
       .attr('text-anchor', 'middle')
       .attr('dominant-baseline', 'middle')
       .text(title);
   }
 
-  // Draw the treemap recursively
-  drawNode(g, root, 0, 0, width, height, config);
-};
+  // Convert data to hierarchical structure
+  const hierarchyRoot = hierarchy<TreemapNode>(root)
+    .sum((d) => d.value || 0)
+    .sort((a, b) => (b.value || 0) - (a.value || 0));
 
-/**
- * Calculates the size of the treemap
- */
-const calculateTreemapSize = (
-  root: TreemapNode,
-  config: any
-): { width: number; height: number } => {
-  // If we have a value, use it as the size
-  if (root.value) {
-    return {
-      width: config.nodeWidth || DEFAULT_NODE_WIDTH,
-      height: config.nodeHeight || DEFAULT_NODE_HEIGHT,
-    };
-  }
+  // Create treemap layout
+  const treemapLayout = treemap<TreemapNode>().size([width, height]).padding(padding).round(true);
 
-  // Otherwise, layout the children
-  if (!root.children || root.children.length === 0) {
-    return {
-      width: config.nodeWidth || DEFAULT_NODE_WIDTH,
-      height: config.nodeHeight || DEFAULT_NODE_HEIGHT,
-    };
-  }
+  // Apply the treemap layout to the hierarchy
+  const treemapData = treemapLayout(hierarchyRoot);
 
-  // Calculate based on children
-  let totalWidth = 0;
-  let maxHeight = 0;
+  // Draw ALL nodes, not just leaves
+  const allNodes = treemapData.descendants();
 
-  // Arrange in a simple tiled layout
-  for (const child of root.children) {
-    const { width, height } = calculateTreemapSize(child, config);
-    totalWidth += width + (config.padding || DEFAULT_PADDING);
-    maxHeight = Math.max(maxHeight, height);
-  }
+  // Draw section nodes (non-leaf nodes)
+  const sections = g
+    .selectAll('.treemapSection')
+    .data(allNodes.filter((d) => d.children && d.children.length > 0))
+    .enter()
+    .append('g')
+    .attr('class', 'treemapSection')
+    .attr('transform', (d) => `translate(${d.x0},${d.y0})`);
 
-  // Remove the last padding
-  totalWidth -= config.padding || DEFAULT_PADDING;
-
-  return {
-    width: Math.max(totalWidth, config.nodeWidth || DEFAULT_NODE_WIDTH),
-    height: Math.max(
-      maxHeight + (config.padding || DEFAULT_PADDING) * 2,
-      config.nodeHeight || DEFAULT_NODE_HEIGHT
-    ),
-  };
-};
-
-/**
- * Recursively draws a node and its children in the treemap
- */
-const drawNode = (
-  parent: any,
-  node: TreemapNode,
-  x: number,
-  y: number,
-  width: number,
-  height: number,
-  config: any
-) => {
-  // Add rectangle
-  parent
+  // Add rectangles for the sections
+  sections
     .append('rect')
-    .attr('x', x)
-    .attr('y', y)
-    .attr('width', width)
-    .attr('height', height)
-    .attr('class', `treemapNode ${node.value ? 'treemapLeaf' : 'treemapSection'}`);
+    .attr('width', (d) => d.x1 - d.x0)
+    .attr('height', (d) => d.y1 - d.y0)
+    .attr('class', 'treemapSectionRect')
+    .attr('fill', (d) => colorScale(d.data.name))
+    .attr('fill-opacity', 0.2)
+    .attr('stroke', (d) => colorScale(d.data.name))
+    .attr('stroke-width', 1);
 
-  // Add the label
-  parent
+  // Add section labels
+  sections
     .append('text')
-    .attr('x', x + width / 2)
-    .attr('y', y + 20) // Position the label at the top
-    .attr('class', 'treemapLabel')
-    .attr('text-anchor', 'middle')
-    .text(node.name);
+    .attr('class', 'treemapSectionLabel')
+    .attr('x', 4)
+    .attr('y', 14)
+    .text((d) => d.data.name)
+    .attr('font-weight', 'bold');
 
-  // Add the value if it exists and should be shown
-  if (node.value !== undefined && config.showValues !== false) {
-    parent
+  // Add section values if enabled
+  if (config.showValues !== false) {
+    sections
       .append('text')
-      .attr('x', x + width / 2)
-      .attr('y', y + height - 10) // Position the value at the bottom
-      .attr('class', 'treemapValue')
-      .attr('text-anchor', 'middle')
-      .text(node.value);
+      .attr('class', 'treemapSectionValue')
+      .attr('x', 4)
+      .attr('y', 28)
+      .text((d) => (d.value ? valueFormat(d.value) : ''))
+      .attr('font-style', 'italic');
   }
 
-  // If this is a section with children, layout and draw the children
-  if (!node.value && node.children && node.children.length > 0) {
-    // Simple tiled layout for children
-    const padding = config.padding || DEFAULT_PADDING;
-    let currentX = x + padding;
-    const innerY = y + 30; // Allow space for the label
-    const innerHeight = height - 40; // Allow space for label
+  // Draw the leaf nodes (nodes with no children)
+  const cell = g
+    .selectAll('.treemapLeaf')
+    .data(treemapData.leaves())
+    .enter()
+    .append('g')
+    .attr('class', 'treemapNode')
+    .attr('transform', (d) => `translate(${d.x0},${d.y0})`);
 
-    for (const child of node.children) {
-      const childWidth = width / node.children.length - padding;
-      drawNode(parent, child, currentX, innerY, childWidth, innerHeight, config);
-      currentX += childWidth + padding;
-    }
+  // Add rectangle for each leaf node
+  cell
+    .append('rect')
+    .attr('width', (d) => d.x1 - d.x0)
+    .attr('height', (d) => d.y1 - d.y0)
+    .attr('class', 'treemapLeaf')
+    .attr('fill', (d) => {
+      // Go up to parent for color
+      let current = d;
+      while (current.depth > 1 && current.parent) {
+        current = current.parent;
+      }
+      return colorScale(current.data.name);
+    })
+    .attr('fill-opacity', 0.8);
+
+  // Add node labels
+  cell
+    .append('text')
+    .attr('class', 'treemapLabel')
+    .attr('x', 4)
+    .attr('y', 14)
+    .text((d) => d.data.name);
+
+  // Add node values if enabled
+  if (config.showValues !== false) {
+    cell
+      .append('text')
+      .attr('class', 'treemapValue')
+      .attr('x', 4)
+      .attr('y', 26)
+      .text((d) => (d.value ? valueFormat(d.value) : ''));
   }
 };
 

packages/mermaid/src/diagrams/treemap/styles.ts

@@ -21,6 +21,11 @@ export const getStyles: DiagramStylesProvider = ({
   const options = cleanAndMerge(defaultPacketStyleOptions, packet);
 
   return `
+  .treemapNode {
+    fill: pink;
+    stroke: black;
+    stroke-width: 1;
+  }
 	.packetByte {
 		font-size: ${options.byteFontSize};
 	}

packages/mermaid/src/diagrams/treemap/utils.test.ts

@@ -0,0 +1,100 @@
+import { describe, it, expect } from 'vitest';
+import { buildHierarchy } from './utils.js';
+import type { TreemapNode } from './types.js';
+
+describe('treemap utilities', () => {
+  describe('buildHierarchy', () => {
+    it('should convert a flat array into a hierarchical structure', () => {
+      // Input flat structure
+      const flatItems = [
+        { level: 0, name: 'Root', type: 'Section' },
+        { level: 4, name: 'Branch 1', type: 'Section' },
+        { level: 8, name: 'Leaf 1.1', type: 'Leaf', value: 10 },
+        { level: 8, name: 'Leaf 1.2', type: 'Leaf', value: 15 },
+        { level: 4, name: 'Branch 2', type: 'Section' },
+        { level: 8, name: 'Leaf 2.1', type: 'Leaf', value: 20 },
+        { level: 8, name: 'Leaf 2.2', type: 'Leaf', value: 25 },
+        { level: 8, name: 'Leaf 2.3', type: 'Leaf', value: 30 },
+      ];
+
+      // Expected hierarchical structure
+      const expectedHierarchy: TreemapNode[] = [
+        {
+          name: 'Root',
+          children: [
+            {
+              name: 'Branch 1',
+              children: [
+                { name: 'Leaf 1.1', value: 10 },
+                { name: 'Leaf 1.2', value: 15 },
+              ],
+            },
+            {
+              name: 'Branch 2',
+              children: [
+                { name: 'Leaf 2.1', value: 20 },
+                { name: 'Leaf 2.2', value: 25 },
+                { name: 'Leaf 2.3', value: 30 },
+              ],
+            },
+          ],
+        },
+      ];
+
+      const result = buildHierarchy(flatItems);
+      expect(result).toEqual(expectedHierarchy);
+    });
+
+    it('should handle empty input', () => {
+      expect(buildHierarchy([])).toEqual([]);
+    });
+
+    it('should handle only root nodes', () => {
+      const flatItems = [
+        { level: 0, name: 'Root 1', type: 'Section' },
+        { level: 0, name: 'Root 2', type: 'Section' },
+      ];
+
+      const expected = [
+        { name: 'Root 1', children: [] },
+        { name: 'Root 2', children: [] },
+      ];
+
+      expect(buildHierarchy(flatItems)).toEqual(expected);
+    });
+
+    it('should handle complex nesting levels', () => {
+      const flatItems = [
+        { level: 0, name: 'Root', type: 'Section' },
+        { level: 2, name: 'Level 1', type: 'Section' },
+        { level: 4, name: 'Level 2', type: 'Section' },
+        { level: 6, name: 'Leaf 1', type: 'Leaf', value: 10 },
+        { level: 4, name: 'Level 2 again', type: 'Section' },
+        { level: 6, name: 'Leaf 2', type: 'Leaf', value: 20 },
+      ];
+
+      const expected = [
+        {
+          name: 'Root',
+          children: [
+            {
+              name: 'Level 1',
+              children: [
+                {
+                  name: 'Level 2',
+                  children: [{ name: 'Leaf 1', value: 10 }],
+                },
+                {
+                  name: 'Level 2 again',
+                  children: [{ name: 'Leaf 2', value: 20 }],
+                },
+              ],
+            },
+          ],
+        },
+      ];
+
+      expect(buildHierarchy(flatItems)).toEqual(expected);
+    });
+  });
+});

packages/mermaid/src/diagrams/treemap/utils.ts

@@ -0,0 +1,53 @@
+import type { TreemapNode } from './types.js';
+
+/**
+ * Converts a flat array of treemap items into a hierarchical structure
+ * @param items - Array of flat treemap items with level, name, type, and optional value
+ * @returns A hierarchical tree structure
+ */
+export function buildHierarchy(
+  items: { level: number; name: string; type: string; value?: number }[]
+): TreemapNode[] {
+  if (!items.length) {
+    return [];
+  }
+
+  const root: TreemapNode[] = [];
+  const stack: { node: TreemapNode; level: number }[] = [];
+
+  items.forEach((item) => {
+    const node: TreemapNode = {
+      name: item.name,
+      children: item.type === 'Leaf' ? undefined : [],
+    };
+
+    if (item.type === 'Leaf' && item.value !== undefined) {
+      node.value = item.value;
+    }
+
+    // Find the right parent for this node
+    while (stack.length > 0 && stack[stack.length - 1].level >= item.level) {
+      stack.pop();
+    }
+
+    if (stack.length === 0) {
+      // This is a root node
+      root.push(node);
+    } else {
+      // Add as child to the parent
+      const parent = stack[stack.length - 1].node;
+      if (parent.children) {
+        parent.children.push(node);
+      } else {
+        parent.children = [node];
+      }
+    }
+
+    // Only add to stack if it can have children
+    if (item.type !== 'Leaf') {
+      stack.push({ node, level: item.level });
+    }
+  });
+
+  return root;
+}