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 type { DiagramDB } from '../../diagram-api/types.js';
-import DEFAULT_CONFIG from '../../defaultConfig.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,54 +14,42 @@ 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 = {
+  public getNodes() {
-  nodes: [],
+    return this.nodes;
-  levels: new Map(),
+  }
-  outerNodes: [],
-  classes: new Map(),
-};
 
-const state = new ImperativeState<TreemapData>(() => structuredClone(defaultTreemapData));
+  public getConfig() {
-
-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);
 
+  public addNode(node: TreemapNode, level: number) {
+    this.nodes.push(node);
+    this.levels.set(node, level);
     if (level === 0) {
-    data.outerNodes.push(node);
+      this.outerNodes.push(node);
+      this.root ??= node;
+    }
   }
 
-  // Set the root node if this is a level 0 node and we don't have a root yet
+  public getRoot() {
-  if (level === 0 && !data.root) {
+    return { name: '', children: this.outerNodes };
-    data.root = node;
   }
-};
-
-const getRoot = (): TreemapNode | undefined => ({ name: '', children: state.records.outerNodes });
-
-const addClass = (id: string, _style: string) => {
-  const classes = state.records.classes;
-  const styleClass = classes.get(id) ?? { id, styles: [], textStyles: [] };
-  classes.set(id, styleClass);
 
+  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)) {
@@ -78,35 +66,30 @@ const addClass = (id: string, _style: string) => {
         }
       });
     }
+    this.classes.set(id, styleClass);
+  }
 
-  classes.set(id, styleClass);
+  public getClasses() {
-};
+    return this.classes;
-const getClasses = (): Map<string, DiagramStyleClassDef> => {
+  }
-  return state.records.classes;
-};
 
-const getStylesForClass = (classSelector: string): string[] => {
+  public getStylesForClass(classSelector: string): string[] {
-  return state.records.classes.get(classSelector)?.styles ?? [];
+    return this.classes.get(classSelector)?.styles ?? [];
-};
+  }
 
-const clear = () => {
+  public clear() {
     commonClear();
-  state.reset();
+    this.nodes = [];
-};
+    this.levels = new Map();
+    this.outerNodes = [];
+    this.classes = new Map();
+    this.root = undefined;
+  }
 
-export const db: TreemapDB = {
+  public setAccTitle = setAccTitle;
-  getNodes,
+  public getAccTitle = getAccTitle;
-  addNode,
+  public setDiagramTitle = setDiagramTitle;
-  getRoot,
+  public getDiagramTitle = getDiagramTitle;
-  getConfig,
+  public getAccDescription = getAccDescription;
-  clear,
+  public setAccDescription = setAccDescription;
-  setAccTitle,
+}
-  getAccTitle,
-  setDiagramTitle,
-  getDiagramTitle,
-  getAccDescription,
-  setAccDescription,
-  addClass,
-  getClasses,
-  getStylesForClass,
-};

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, TreemapDB } from './types.js';
-import type { TreemapNode, TreemapAst } 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;