refactor: handle StateDB .extract() properly

Fix the StateDB properly for stateDiagram-v1 vs stateDiagram-v2.
mermaid-js · Feb 17, 2025 · cfe710f · cfe710f
1 parent 5482565
commit cfe710f
Show file tree

Hide file tree

Showing 9 changed files with 30 additions and 13 deletions.
diff --git a/packages/mermaid/src/diagrams/state/parser/state-parser.spec.js b/packages/mermaid/src/diagrams/state/parser/state-parser.spec.js
@@ -9,7 +9,7 @@ setConfig({
 describe('state parser can parse...', () => {
   let stateDb;
   beforeEach(function () {
-    stateDb = new StateDB();
+    stateDb = new StateDB(2);
     stateDiagram.parser.yy = stateDb;
     stateDiagram.parser.yy.clear();
   });

diff --git a/packages/mermaid/src/diagrams/state/parser/state-style.spec.js b/packages/mermaid/src/diagrams/state/parser/state-style.spec.js
@@ -9,7 +9,7 @@ setConfig({
 describe('ClassDefs and classes when parsing a State diagram', () => {
   let stateDb;
   beforeEach(function () {
-    stateDb = new StateDB();
+    stateDb = new StateDB(2);
     stateDiagram.parser.yy = stateDb;
     stateDiagram.parser.yy.clear();
   });
@@ -135,7 +135,6 @@ describe('ClassDefs and classes when parsing a State diagram', () => {
           diagram += '[*]:::exampleStyleClass --> b\n';
 
           stateDiagram.parser.parse(diagram);
-          stateDiagram.parser.yy.extract(stateDiagram.parser.yy.getRootDocV2());
 
           const states = stateDiagram.parser.yy.getStates();
           const classes = stateDiagram.parser.yy.getClasses();

diff --git a/packages/mermaid/src/diagrams/state/stateDb.js b/packages/mermaid/src/diagrams/state/stateDb.js
@@ -58,16 +58,27 @@ const newDoc = () => {
 const clone = (o) => JSON.parse(JSON.stringify(o));
 
 export class StateDB {
-  constructor() {
+  /**
+   * @param {1 | 2} version - v1 renderer or v2 renderer.
+   */
+  constructor(version) {
     this.clear();
 
+    this.version = version;
+
     // Needed for JISON since it only supports direct properties
     this.setRootDoc = this.setRootDoc.bind(this);
     this.getDividerId = this.getDividerId.bind(this);
     this.setDirection = this.setDirection.bind(this);
     this.trimColon = this.trimColon.bind(this);
   }
 
+  /**
+   * @private
+   * @type {1 | 2}
+   */
+  version;
+
   /**
    * @private
    * @type {Array}
@@ -130,7 +141,11 @@ export class StateDB {
     log.info('Setting root doc', o);
     // rootDoc = { id: 'root', doc: o };
     this.rootDoc = o;
-    this.extract(o);
+    if (this.version === 1) {
+      this.extract(o);
+    } else {
+      this.extract(this.getRootDocV2());
+    }
   }
 
   getRootDoc() {
@@ -190,6 +205,10 @@ export class StateDB {
       }
     }
   }
+
+  /**
+   * @private
+   */
   getRootDocV2() {
     this.docTranslator({ id: 'root' }, { id: 'root', doc: this.rootDoc }, true);
     return { id: 'root', doc: this.rootDoc };

diff --git a/packages/mermaid/src/diagrams/state/stateDb.spec.js b/packages/mermaid/src/diagrams/state/stateDb.spec.js
@@ -3,7 +3,7 @@ import { StateDB } from './stateDb.js';
 describe('State Diagram stateDb', () => {
   let stateDb;
   beforeEach(() => {
-    stateDb = new StateDB();
+    stateDb = new StateDB(1);
   });
 
   describe('addStyleClass', () => {
@@ -23,7 +23,7 @@ describe('State Diagram stateDb', () => {
   describe('addDescription to a state', () => {
     let stateDb;
     beforeEach(() => {
-      stateDb = new StateDB();
+      stateDb = new StateDB(1);
       stateDb.addState('state1');
     });
 
@@ -79,7 +79,7 @@ describe('State Diagram stateDb', () => {
 describe('state db class', () => {
   let stateDb;
   beforeEach(() => {
-    stateDb = new StateDB();
+    stateDb = new StateDB(1);
   });
   // This is to ensure that functions used in state JISON are exposed as function from StateDb
   it('should have functions used in flow JISON as own property', () => {

diff --git a/packages/mermaid/src/diagrams/state/stateDiagram-v2.spec.js b/packages/mermaid/src/diagrams/state/stateDiagram-v2.spec.js
@@ -6,7 +6,7 @@ describe('state diagram V2, ', function () {
   describe('when parsing an info graph it', function () {
     let stateDb;
     beforeEach(function () {
-      stateDb = new StateDB();
+      stateDb = new StateDB(2);
       parser.yy = stateDb;
       stateDiagram.parser.yy = stateDb;
       stateDiagram.parser.yy.clear();

diff --git a/packages/mermaid/src/diagrams/state/stateDiagram-v2.ts b/packages/mermaid/src/diagrams/state/stateDiagram-v2.ts
@@ -8,7 +8,7 @@ import renderer from './stateRenderer-v3-unified.js';
 export const diagram: DiagramDefinition = {
   parser,
   get db() {
-    return new StateDB();
+    return new StateDB(2);
   },
   renderer,
   styles,

diff --git a/packages/mermaid/src/diagrams/state/stateDiagram.spec.js b/packages/mermaid/src/diagrams/state/stateDiagram.spec.js
@@ -5,7 +5,7 @@ describe('state diagram, ', function () {
   describe('when parsing an info graph it', function () {
     let stateDb;
     beforeEach(function () {
-      stateDb = new StateDB();
+      stateDb = new StateDB(1);
       parser.yy = stateDb;
     });
 

diff --git a/packages/mermaid/src/diagrams/state/stateDiagram.ts b/packages/mermaid/src/diagrams/state/stateDiagram.ts
@@ -8,7 +8,7 @@ import renderer from './stateRenderer.js';
 export const diagram: DiagramDefinition = {
   parser,
   get db() {
-    return new StateDB();
+    return new StateDB(1);
   },
   renderer,
   styles,

diff --git a/packages/mermaid/src/diagrams/state/stateRenderer.js b/packages/mermaid/src/diagrams/state/stateRenderer.js
@@ -136,7 +136,6 @@ const renderDoc = (doc, diagram, parentId, altBkg, root, domDocument, diagObj) =
     return {};
   });
 
-  diagObj.db.extract(doc);
   const states = diagObj.db.getStates();
   const relations = diagObj.db.getRelations();