Merge pull request #6738 from mermaid-js/treemap-diagram-to-use-the-new-class-based-approach

Updated code to use class based approach for treemap

.changeset/pretty-falcons-say.md

@@ -0,0 +1,5 @@
+---
+'mermaid': patch
+---
+
+chore: Updated TreeMapDB to use class based approach

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

@@ -1,10 +1,10 @@
-import { getConfig as commonGetConfig } from '../../config.js';
-import DEFAULT_CONFIG from '../../defaultConfig.js';
+import type { DiagramDB } from '../../diagram-api/types.js';
 import type { DiagramStyleClassDef } from '../../diagram-api/types.js';
-import { isLabelStyle } from '../../rendering-util/rendering-elements/shapes/handDrawnShapeStyles.js';
-
+import type { TreemapDiagramConfig, TreemapNode } from './types.js';
+import DEFAULT_CONFIG from '../../defaultConfig.js';
+import { getConfig as commonGetConfig } from '../../config.js';
 import { cleanAndMerge } from '../../utils.js';
-import { ImperativeState } from '../../utils/imperativeState.js';
+import { isLabelStyle } from '../../rendering-util/rendering-elements/shapes/handDrawnShapeStyles.js';
 import {
   clear as commonClear,
   getAccDescription,
@@ -14,99 +14,82 @@ import {
   setAccTitle,
   setDiagramTitle,
 } from '../common/commonDb.js';
-import type { TreemapDB, TreemapData, TreemapDiagramConfig, TreemapNode } from './types.js';
+export class TreeMapDB implements DiagramDB {
+  private nodes: TreemapNode[] = [];
+  private levels: Map<TreemapNode, number> = new Map<TreemapNode, number>();
+  private outerNodes: TreemapNode[] = [];
+  private classes: Map<string, DiagramStyleClassDef> = new Map<string, DiagramStyleClassDef>();
+  private root?: TreemapNode;
 
-const defaultTreemapData: TreemapData = {
-  nodes: [],
-  levels: new Map(),
-  outerNodes: [],
-  classes: new Map(),
-};
-
-const state = new ImperativeState<TreemapData>(() => structuredClone(defaultTreemapData));
-
-const getConfig = (): Required<TreemapDiagramConfig> => {
-  // Use type assertion with unknown as intermediate step
-  const defaultConfig = DEFAULT_CONFIG as unknown as { treemap: Required<TreemapDiagramConfig> };
-  const userConfig = commonGetConfig() as unknown as { treemap?: Partial<TreemapDiagramConfig> };
-
-  return cleanAndMerge({
-    ...defaultConfig.treemap,
-    ...(userConfig.treemap ?? {}),
-  }) as Required<TreemapDiagramConfig>;
-};
-
-const getNodes = (): TreemapNode[] => state.records.nodes;
-
-const addNode = (node: TreemapNode, level: number) => {
-  const data = state.records;
-  data.nodes.push(node);
-  data.levels.set(node, level);
-
-  if (level === 0) {
-    data.outerNodes.push(node);
+  public getNodes() {
+    return this.nodes;
   }
 
-  // Set the root node if this is a level 0 node and we don't have a root yet
-  if (level === 0 && !data.root) {
-    data.root = node;
+  public getConfig() {
+    const defaultConfig = DEFAULT_CONFIG as unknown as { treemap: Required<TreemapDiagramConfig> };
+    const userConfig = commonGetConfig() as unknown as { treemap?: Partial<TreemapDiagramConfig> };
+    return cleanAndMerge({
+      ...defaultConfig.treemap,
+      ...(userConfig.treemap ?? {}),
+    }) as Required<TreemapDiagramConfig>;
   }
-};
 
-const getRoot = (): TreemapNode | undefined => ({ name: '', children: state.records.outerNodes });
+  public addNode(node: TreemapNode, level: number) {
+    this.nodes.push(node);
+    this.levels.set(node, level);
+    if (level === 0) {
+      this.outerNodes.push(node);
+      this.root ??= node;
+    }
+  }
 
-const addClass = (id: string, _style: string) => {
-  const classes = state.records.classes;
-  const styleClass = classes.get(id) ?? { id, styles: [], textStyles: [] };
-  classes.set(id, styleClass);
+  public getRoot() {
+    return { name: '', children: this.outerNodes };
+  }
 
-  const styles = _style.replace(/\\,/g, '§§§').replace(/,/g, ';').replace(/§§§/g, ',').split(';');
-
-  if (styles) {
-    styles.forEach((s) => {
-      if (isLabelStyle(s)) {
-        if (styleClass?.textStyles) {
-          styleClass.textStyles.push(s);
-        } else {
-          styleClass.textStyles = [s];
+  public addClass(id: string, _style: string) {
+    const styleClass = this.classes.get(id) ?? { id, styles: [], textStyles: [] };
+    const styles = _style.replace(/\\,/g, '§§§').replace(/,/g, ';').replace(/§§§/g, ',').split(';');
+    if (styles) {
+      styles.forEach((s) => {
+        if (isLabelStyle(s)) {
+          if (styleClass?.textStyles) {
+            styleClass.textStyles.push(s);
+          } else {
+            styleClass.textStyles = [s];
+          }
         }
-      }
-      if (styleClass?.styles) {
-        styleClass.styles.push(s);
-      } else {
-        styleClass.styles = [s];
-      }
-    });
+        if (styleClass?.styles) {
+          styleClass.styles.push(s);
+        } else {
+          styleClass.styles = [s];
+        }
+      });
+    }
+    this.classes.set(id, styleClass);
   }
 
-  classes.set(id, styleClass);
-};
-const getClasses = (): Map<string, DiagramStyleClassDef> => {
-  return state.records.classes;
-};
+  public getClasses() {
+    return this.classes;
+  }
 
-const getStylesForClass = (classSelector: string): string[] => {
-  return state.records.classes.get(classSelector)?.styles ?? [];
-};
+  public getStylesForClass(classSelector: string): string[] {
+    return this.classes.get(classSelector)?.styles ?? [];
+  }
 
-const clear = () => {
-  commonClear();
-  state.reset();
-};
+  public clear() {
+    commonClear();
+    this.nodes = [];
+    this.levels = new Map();
+    this.outerNodes = [];
+    this.classes = new Map();
+    this.root = undefined;
+  }
 
-export const db: TreemapDB = {
-  getNodes,
-  addNode,
-  getRoot,
-  getConfig,
-  clear,
-  setAccTitle,
-  getAccTitle,
-  setDiagramTitle,
-  getDiagramTitle,
-  getAccDescription,
-  setAccDescription,
-  addClass,
-  getClasses,
-  getStylesForClass,
-};
+  public setAccTitle = setAccTitle;
+  public getAccTitle = getAccTitle;
+  public setDiagramTitle = setDiagramTitle;
+  public getDiagramTitle = getDiagramTitle;
+  public getAccDescription = getAccDescription;
+  public setAccDescription = setAccDescription;
+}

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

@@ -1,12 +1,14 @@
 import type { DiagramDefinition } from '../../diagram-api/types.js';
-import { db } from './db.js';
+import { TreeMapDB } from './db.js';
 import { parser } from './parser.js';
 import { renderer } from './renderer.js';
 import styles from './styles.js';
 
 export const diagram: DiagramDefinition = {
   parser,
-  db,
+  get db() {
+    return new TreeMapDB();
+  },
   renderer,
   styles,
 };

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

@@ -2,15 +2,15 @@ import { parse } from '@mermaid-js/parser';
 import type { ParserDefinition } from '../../diagram-api/types.js';
 import { log } from '../../logger.js';
 import { populateCommonDb } from '../common/populateCommonDb.js';
-import { db } from './db.js';
-import type { TreemapNode, TreemapAst } from './types.js';
+import type { TreemapNode, TreemapAst, TreemapDB } from './types.js';
 import { buildHierarchy } from './utils.js';
+import { TreeMapDB } from './db.js';
 
 /**
  * Populates the database with data from the Treemap AST
  * @param ast - The Treemap AST
  */
-const populate = (ast: TreemapAst) => {
+const populate = (ast: TreemapAst, db: TreemapDB) => {
   // We need to bypass the type checking for populateCommonDb
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   populateCommonDb(ast as any, db);
@@ -84,6 +84,8 @@ const getItemName = (item: { name?: string | number }): string => {
 };
 
 export const parser: ParserDefinition = {
+  // @ts-expect-error - TreeMapDB is not assignable to DiagramDB
+  parser: { yy: undefined },
   parse: async (text: string): Promise<void> => {
     try {
       // Use a generic parse that accepts any diagram type
@@ -91,7 +93,13 @@ export const parser: ParserDefinition = {
       const parseFunc = parse as (diagramType: string, text: string) => Promise<TreemapAst>;
       const ast = await parseFunc('treemap', text);
       log.debug('Treemap AST:', ast);
-      populate(ast);
+      const db = parser.parser?.yy;
+      if (!(db instanceof TreeMapDB)) {
+        throw new Error(
+          'parser.parser?.yy was not a TreemapDB. This is due to a bug within Mermaid, please report this issue at https://github.com/mermaid-js/mermaid/issues.'
+        );
+      }
+      populate(ast, db);
     } catch (error) {
       log.error('Error parsing treemap:', error);
       throw error;