Updated code 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 { cleanAndMerge } from '../../utils.js';
 import { ImperativeState } from '../../utils/imperativeState.js';
+import type { TreemapData, TreemapDiagramConfig, TreemapNode } from './types.js';
+import DEFAULT_CONFIG from '../../defaultConfig.js';
+import { getConfig as commonGetConfig } from '../../config.js';
+import { cleanAndMerge } from '../../utils.js';
+import { isLabelStyle } from '../../rendering-util/rendering-elements/shapes/handDrawnShapeStyles.js';
 import {
   clear as commonClear,
   getAccDescription,
@@ -14,18 +14,26 @@ import {
   setAccTitle,
   setDiagramTitle,
 } from '../common/commonDb.js';
-import type { TreemapDB, TreemapData, TreemapDiagramConfig, TreemapNode } from './types.js';
+export class TreeMapDB implements DiagramDB {
+  private data: TreemapData;
+  private state: ImperativeState<TreemapData>;
 
-const defaultTreemapData: TreemapData = {
+  constructor() {
+    this.data = {
       nodes: [],
       levels: new Map(),
       outerNodes: [],
       classes: new Map(),
     };
 
-const state = new ImperativeState<TreemapData>(() => structuredClone(defaultTreemapData));
+    this.state = new ImperativeState<TreemapData>(() => structuredClone(this.data));
+  }
 
-const getConfig = (): Required<TreemapDiagramConfig> => {
+  public getNodes() {
+    return this.state.records.nodes;
+  }
+
+  public getConfig() {
     // 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> };
@@ -34,12 +42,10 @@ const getConfig = (): Required<TreemapDiagramConfig> => {
       ...defaultConfig.treemap,
       ...(userConfig.treemap ?? {}),
     }) as Required<TreemapDiagramConfig>;
-};
+  }
 
-const getNodes = (): TreemapNode[] => state.records.nodes;
+  public addNode(node: TreemapNode, level: number) {
-
+    const data = this.state.records;
-const addNode = (node: TreemapNode, level: number) => {
-  const data = state.records;
     data.nodes.push(node);
     data.levels.set(node, level);
 
@@ -51,14 +57,16 @@ const addNode = (node: TreemapNode, level: number) => {
     if (level === 0 && !data.root) {
       data.root = node;
     }
-};
+  }
 
-const getRoot = (): TreemapNode | undefined => ({ name: '', children: state.records.outerNodes });
+  public getRoot() {
+    return { name: '', children: this.state.records.outerNodes };
+  }
 
-const addClass = (id: string, _style: string) => {
+  public addClass(id: string, _style: string) {
-  const classes = state.records.classes;
+    // const classes = this.state.records.classes;
-  const styleClass = classes.get(id) ?? { id, styles: [], textStyles: [] };
+    const styleClass = this.state.records.classes.get(id) ?? { id, styles: [], textStyles: [] };
-  classes.set(id, styleClass);
+    this.state.records.classes.set(id, styleClass);
 
     const styles = _style.replace(/\\,/g, '§§§').replace(/,/g, ';').replace(/§§§/g, ',').split(';');
 
@@ -79,34 +87,26 @@ const addClass = (id: string, _style: string) => {
       });
     }
 
-  classes.set(id, styleClass);
+    this.state.records.classes.set(id, styleClass);
-};
+  }
-const getClasses = (): Map<string, DiagramStyleClassDef> => {
-  return state.records.classes;
-};
 
-const getStylesForClass = (classSelector: string): string[] => {
+  public getClasses() {
-  return state.records.classes.get(classSelector)?.styles ?? [];
+    return this.state.records.classes;
-};
+  }
 
-const clear = () => {
+  public getStylesForClass(classSelector: string): string[] {
+    return this.state.records.classes.get(classSelector)?.styles ?? [];
+  }
+
+  public clear = () => {
     commonClear();
-  state.reset();
+    this.state.reset();
   };
 
-export const db: TreemapDB = {
+  public setAccTitle = setAccTitle;
-  getNodes,
+  public setAccDescription = setAccDescription;
-  addNode,
+  public setDiagramTitle = setDiagramTitle;
-  getRoot,
+  public getAccTitle = getAccTitle;
-  getConfig,
+  public getAccDescription = getAccDescription;
-  clear,
+  public getDiagramTitle = getDiagramTitle;
-  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,14 @@ 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';
 
 /**
  * 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 +83,7 @@ const getItemName = (item: { name?: string | number }): string => {
 };
 
 export const parser: ParserDefinition = {
+  parser: { yy: undefined },
   parse: async (text: string): Promise<void> => {
     try {
       // Use a generic parse that accepts any diagram type
@@ -91,7 +91,7 @@ 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);
+      populate(ast, parser.parser?.yy as TreemapDB);
     } catch (error) {
       log.error('Error parsing treemap:', error);
       throw error;