diff --git a/packages/mermaid/src/diagrams/treemap/parser.ts b/packages/mermaid/src/diagrams/treemap/parser.ts index 943f6622a..88fb647f0 100644 --- a/packages/mermaid/src/diagrams/treemap/parser.ts +++ b/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); }; /** diff --git a/packages/mermaid/src/diagrams/treemap/renderer.ts b/packages/mermaid/src/diagrams/treemap/renderer.ts index 24a512d36..6f7daecc7 100644 --- a/packages/mermaid/src/diagrams/treemap/renderer.ts +++ b/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().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(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().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) : '')); } }; diff --git a/packages/mermaid/src/diagrams/treemap/styles.ts b/packages/mermaid/src/diagrams/treemap/styles.ts index 03c6328a8..5c80e7810 100644 --- a/packages/mermaid/src/diagrams/treemap/styles.ts +++ b/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}; } diff --git a/packages/mermaid/src/diagrams/treemap/utils.test.ts b/packages/mermaid/src/diagrams/treemap/utils.test.ts new file mode 100644 index 000000000..bfbd74c59 --- /dev/null +++ b/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); + }); + }); +}); diff --git a/packages/mermaid/src/diagrams/treemap/utils.ts b/packages/mermaid/src/diagrams/treemap/utils.ts new file mode 100644 index 000000000..74c4f793a --- /dev/null +++ b/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; +}