Merge pull request #6296 from mermaid-js/sidv/stateDB-ts

chore: Convert StateDB to TS

.changeset/eleven-wolves-deny.md

@@ -0,0 +1,5 @@
+---
+'mermaid': patch
+---
+
+chore: Convert StateDB into TypeScript

packages/mermaid/src/diagrams/state/dataFetcher.ts

@@ -1,3 +1,4 @@
+import type { MermaidConfig } from '../../config.type.js';
 import { getConfig } from '../../diagram-api/diagramAPI.js';
 import { log } from '../../logger.js';
 import common from '../common/common.js';
@@ -33,9 +34,10 @@ import {
   STMT_RELATION,
   STMT_STATE,
 } from './stateCommon.js';
+import type { Edge, NodeData, StateStmt, Stmt, StyleClass } from './stateDb.js';
 
 // List of nodes created from the parsed diagram statement items
-let nodeDb = new Map();
+const nodeDb = new Map<string, NodeData>();
 
 let graphItemCount = 0; // used to construct ids, etc.
 
@@ -43,18 +45,27 @@ let graphItemCount = 0; // used to construct ids, etc.
  * Create a standard string for the dom ID of an item.
  * If a type is given, insert that before the counter, preceded by the type spacer
  *
- * @param itemId
- * @param counter
- * @param {string | null} type
- * @param typeSpacer
- * @returns {string}
  */
-export function stateDomId(itemId = '', counter = 0, type = '', typeSpacer = DOMID_TYPE_SPACER) {
+export function stateDomId(
+  itemId = '',
+  counter = 0,
+  type: string | null = '',
+  typeSpacer = DOMID_TYPE_SPACER
+) {
   const typeStr = type !== null && type.length > 0 ? `${typeSpacer}${type}` : '';
   return `${DOMID_STATE}-${itemId}${typeStr}-${counter}`;
 }
 
-const setupDoc = (parentParsedItem, doc, diagramStates, nodes, edges, altFlag, look, classes) => {
+const setupDoc = (
+  parentParsedItem: StateStmt | undefined,
+  doc: Stmt[],
+  diagramStates: Map<string, StateStmt>,
+  nodes: NodeData[],
+  edges: Edge[],
+  altFlag: boolean,
+  look: MermaidConfig['look'],
+  classes: Map<string, StyleClass>
+) => {
   // graphItemCount = 0;
   log.trace('items', doc);
   doc.forEach((item) => {
@@ -95,7 +106,7 @@ const setupDoc = (parentParsedItem, doc, diagramStates, nodes, edges, altFlag, l
             arrowTypeEnd: 'arrow_barb',
             style: G_EDGE_STYLE,
             labelStyle: '',
-            label: common.sanitizeText(item.description, getConfig()),
+            label: common.sanitizeText(item.description ?? '', getConfig()),
             arrowheadStyle: G_EDGE_ARROWHEADSTYLE,
             labelpos: G_EDGE_LABELPOS,
             labelType: G_EDGE_LABELTYPE,
@@ -115,11 +126,10 @@ const setupDoc = (parentParsedItem, doc, diagramStates, nodes, edges, altFlag, l
  * Get the direction from the statement items.
  * Look through all of the documents (docs) in the parsedItems
  * Because is a _document_ direction, the default direction is not necessarily the same as the overall default _diagram_ direction.
- * @param {object[]} parsedItem - the parsed statement item to look through
+ * @param parsedItem - the parsed statement item to look through
- * @param [defaultDir] - the direction to use if none is found
+ * @param defaultDir - the direction to use if none is found
- * @returns {string}
  */
-const getDir = (parsedItem, defaultDir = DEFAULT_NESTED_DOC_DIR) => {
+const getDir = (parsedItem: { doc?: Stmt[] }, defaultDir = DEFAULT_NESTED_DOC_DIR) => {
   let dir = defaultDir;
   if (parsedItem.doc) {
     for (const parsedItemDoc of parsedItem.doc) {
@@ -131,7 +141,11 @@ const getDir = (parsedItem, defaultDir = DEFAULT_NESTED_DOC_DIR) => {
   return dir;
 };
 
-function insertOrUpdateNode(nodes, nodeData, classes) {
+function insertOrUpdateNode(
+  nodes: NodeData[],
+  nodeData: NodeData,
+  classes: Map<string, StyleClass>
+) {
   if (!nodeData.id || nodeData.id === '</join></fork>' || nodeData.id === '</choice>') {
     return;
   }
@@ -143,9 +157,9 @@ function insertOrUpdateNode(nodes, nodeData, classes) {
     }
 
     nodeData.cssClasses.split(' ').forEach((cssClass) => {
-      if (classes.get(cssClass)) {
+      const classDef = classes.get(cssClass);
-        const classDef = classes.get(cssClass);
+      if (classDef) {
-        nodeData.cssCompiledStyles = [...nodeData.cssCompiledStyles, ...classDef.styles];
+        nodeData.cssCompiledStyles = [...(nodeData.cssCompiledStyles ?? []), ...classDef.styles];
       }
     });
   }
@@ -162,31 +176,30 @@ function insertOrUpdateNode(nodes, nodeData, classes) {
  * If there aren't any or if dbInfoItem isn't defined, return an empty string.
  * Else create 1 string from the list of classes found
  *
- * @param {undefined | null | object} dbInfoItem
- * @returns {string}
  */
-function getClassesFromDbInfo(dbInfoItem) {
+function getClassesFromDbInfo(dbInfoItem?: StateStmt): string {
   return dbInfoItem?.classes?.join(' ') ?? '';
 }
 
-function getStylesFromDbInfo(dbInfoItem) {
+function getStylesFromDbInfo(dbInfoItem?: StateStmt): string[] {
   return dbInfoItem?.styles ?? [];
 }
 
 export const dataFetcher = (
-  parent,
+  parent: StateStmt | undefined,
-  parsedItem,
+  parsedItem: StateStmt,
-  diagramStates,
+  diagramStates: Map<string, StateStmt>,
-  nodes,
+  nodes: NodeData[],
-  edges,
+  edges: Edge[],
-  altFlag,
+  altFlag: boolean,
-  look,
+  look: MermaidConfig['look'],
-  classes
+  classes: Map<string, StyleClass>
 ) => {
   const itemId = parsedItem.id;
   const dbState = diagramStates.get(itemId);
   const classStr = getClassesFromDbInfo(dbState);
   const style = getStylesFromDbInfo(dbState);
+  const config = getConfig();
 
   log.info('dataFetcher parsedItem', parsedItem, dbState, style);
 
@@ -207,13 +220,13 @@ export const dataFetcher = (
       nodeDb.set(itemId, {
         id: itemId,
         shape,
-        description: common.sanitizeText(itemId, getConfig()),
+        description: common.sanitizeText(itemId, config),
         cssClasses: `${classStr} ${CSS_DIAGRAM_STATE}`,
         cssStyles: style,
       });
     }
 
-    const newNode = nodeDb.get(itemId);
+    const newNode = nodeDb.get(itemId)!;
 
     // Save data for description and group so that for instance a statement without description overwrites
     // one with description  @todo TODO What does this mean? If important, add a test for it
@@ -225,7 +238,7 @@ export const dataFetcher = (
         newNode.shape = SHAPE_STATE_WITH_DESC;
         newNode.description.push(parsedItem.description);
       } else {
-        if (newNode.description?.length > 0) {
+        if (newNode.description?.length && newNode.description.length > 0) {
           // if there is a description already transform it to an array
           newNode.shape = SHAPE_STATE_WITH_DESC;
           if (newNode.description === itemId) {
@@ -239,7 +252,7 @@ export const dataFetcher = (
           newNode.description = parsedItem.description;
         }
       }
-      newNode.description = common.sanitizeTextOrArray(newNode.description, getConfig());
+      newNode.description = common.sanitizeTextOrArray(newNode.description, config);
     }
 
     // If there's only 1 description entry, just use a regular state shape
@@ -262,7 +275,7 @@ export const dataFetcher = (
     }
 
     // This is what will be added to the graph
-    const nodeData = {
+    const nodeData: NodeData = {
       labelStyle: '',
       shape: newNode.shape,
       label: newNode.description,
@@ -294,19 +307,19 @@ export const dataFetcher = (
 
     if (parsedItem.note) {
       // Todo: set random id
-      const noteData = {
+      const noteData: NodeData = {
         labelStyle: '',
         shape: SHAPE_NOTE,
         label: parsedItem.note.text,
         cssClasses: CSS_DIAGRAM_NOTE,
         // useHtmlLabels: false,
         cssStyles: [],
-        cssCompilesStyles: [],
+        cssCompiledStyles: [],
         id: itemId + NOTE_ID + '-' + graphItemCount,
         domId: stateDomId(itemId, graphItemCount, NOTE),
         type: newNode.type,
         isGroup: newNode.type === 'group',
-        padding: getConfig().flowchart.padding,
+        padding: config.flowchart?.padding,
         look,
         position: parsedItem.note.position,
       };

packages/mermaid/src/diagrams/state/id-cache.js

@@ -1,16 +0,0 @@
-const idCache = {};
-
-export const set = (key, val) => {
-  idCache[key] = val;
-};
-
-export const get = (k) => idCache[k];
-export const keys = () => Object.keys(idCache);
-export const size = () => keys().length;
-
-export default {
-  get,
-  set,
-  keys,
-  size,
-};

packages/mermaid/src/diagrams/state/shapes.js

@@ -1,5 +1,4 @@
 import { line, curveBasis } from 'd3';
-import idCache from './id-cache.js';
 import { StateDB } from './stateDb.js';
 import utils from '../../utils.js';
 import common from '../common/common.js';
@@ -405,8 +404,6 @@ export const drawState = function (elem, stateDef) {
   stateInfo.width = stateBox.width + 2 * getConfig().state.padding;
   stateInfo.height = stateBox.height + 2 * getConfig().state.padding;
 
-  idCache.set(id, stateInfo);
-  // stateCnt++;
   return stateInfo;
 };
 

packages/mermaid/src/diagrams/state/stateCommon.ts

@@ -13,6 +13,10 @@ export const STMT_DIRECTION = 'dir';
 
 // parsed statement type for a state
 export const STMT_STATE = 'state';
+
+// parsed statement type for a root
+export const STMT_ROOT = 'root';
+
 // parsed statement type for a relation
 export const STMT_RELATION = 'relation';
 // parsed statement type for a classDef

packages/mermaid/src/diagrams/state/stateDb.js

@@ -1,706 +0,0 @@
-import { getConfig } from '../../diagram-api/diagramAPI.js';
-import { log } from '../../logger.js';
-import { generateId } from '../../utils.js';
-import common from '../common/common.js';
-import {
-  clear as commonClear,
-  getAccDescription,
-  getAccTitle,
-  getDiagramTitle,
-  setAccDescription,
-  setAccTitle,
-  setDiagramTitle,
-} from '../common/commonDb.js';
-import { dataFetcher, reset as resetDataFetching } from './dataFetcher.js';
-import { getDir } from './stateRenderer-v3-unified.js';
-
-import {
-  DEFAULT_DIAGRAM_DIRECTION,
-  DEFAULT_STATE_TYPE,
-  DIVIDER_TYPE,
-  STMT_APPLYCLASS,
-  STMT_CLASSDEF,
-  STMT_DIRECTION,
-  STMT_RELATION,
-  STMT_STATE,
-  STMT_STYLEDEF,
-} from './stateCommon.js';
-
-const START_NODE = '[*]';
-const START_TYPE = 'start';
-const END_NODE = START_NODE;
-const END_TYPE = 'end';
-
-const COLOR_KEYWORD = 'color';
-const FILL_KEYWORD = 'fill';
-const BG_FILL = 'bgFill';
-const STYLECLASS_SEP = ',';
-
-/**
- * Returns a new list of classes.
- * In the future, this can be replaced with a class common to all diagrams.
- * ClassDef information = { id: id, styles: [], textStyles: [] }
- *
- * @returns {Map<string, any>}
- */
-function newClassesList() {
-  return new Map();
-}
-
-const newDoc = () => {
-  return {
-    /** @type {{ id1: string, id2: string, relationTitle: string }[]} */
-    relations: [],
-    states: new Map(),
-    documents: {},
-  };
-};
-
-const clone = (o) => JSON.parse(JSON.stringify(o));
-
-export class StateDB {
-  /**
-   * @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}
-   */
-  nodes = [];
-  /**
-   * @private
-   * @type {Array}
-   */
-  edges = [];
-
-  /**
-   * @private
-   * @type {Array}
-   */
-  rootDoc = [];
-  /**
-   * @private
-   * @type {Map<string, any>}
-   */
-  classes = newClassesList(); // style classes defined by a classDef
-
-  /**
-   * @private
-   * @type {Object}
-   */
-  documents = {
-    root: newDoc(),
-  };
-
-  /**
-   * @private
-   * @type {Object}
-   */
-  currentDocument = this.documents.root;
-  /**
-   * @private
-   * @type {number}
-   */
-  startEndCount = 0;
-  /**
-   * @private
-   * @type {number}
-   */
-  dividerCnt = 0;
-
-  static relationType = {
-    AGGREGATION: 0,
-    EXTENSION: 1,
-    COMPOSITION: 2,
-    DEPENDENCY: 3,
-  };
-
-  setRootDoc(o) {
-    log.info('Setting root doc', o);
-    // rootDoc = { id: 'root', doc: o };
-    this.rootDoc = o;
-    if (this.version === 1) {
-      this.extract(o);
-    } else {
-      this.extract(this.getRootDocV2());
-    }
-  }
-
-  getRootDoc() {
-    return this.rootDoc;
-  }
-
-  /**
-   * @private
-   * @param {Object} parent
-   * @param {Object} node
-   * @param {boolean} first
-   */
-  docTranslator(parent, node, first) {
-    if (node.stmt === STMT_RELATION) {
-      this.docTranslator(parent, node.state1, true);
-      this.docTranslator(parent, node.state2, false);
-    } else {
-      if (node.stmt === STMT_STATE) {
-        if (node.id === '[*]') {
-          node.id = first ? parent.id + '_start' : parent.id + '_end';
-          node.start = first;
-        } else {
-          // This is just a plain state, not a start or end
-          node.id = node.id.trim();
-        }
-      }
-
-      if (node.doc) {
-        const doc = [];
-        // Check for concurrency
-        let currentDoc = [];
-        let i;
-        for (i = 0; i < node.doc.length; i++) {
-          if (node.doc[i].type === DIVIDER_TYPE) {
-            const newNode = clone(node.doc[i]);
-            newNode.doc = clone(currentDoc);
-            doc.push(newNode);
-            currentDoc = [];
-          } else {
-            currentDoc.push(node.doc[i]);
-          }
-        }
-
-        // If any divider was encountered
-        if (doc.length > 0 && currentDoc.length > 0) {
-          const newNode = {
-            stmt: STMT_STATE,
-            id: generateId(),
-            type: 'divider',
-            doc: clone(currentDoc),
-          };
-          doc.push(clone(newNode));
-          node.doc = doc;
-        }
-
-        node.doc.forEach((docNode) => this.docTranslator(node, docNode, true));
-      }
-    }
-  }
-
-  /**
-   * @private
-   */
-  getRootDocV2() {
-    this.docTranslator({ id: 'root' }, { id: 'root', doc: this.rootDoc }, true);
-    return { id: 'root', doc: this.rootDoc };
-    // Here
-  }
-
-  /**
-   * Convert all of the statements (stmts) that were parsed into states and relationships.
-   * This is done because a state diagram may have nested sections,
-   * where each section is a 'document' and has its own set of statements.
-   * Ex: the section within a fork has its own statements, and incoming and outgoing statements
-   * refer to the fork as a whole (document).
-   * See the parser grammar:  the definition of a document is a document then a 'line', where a line can be a statement.
-   * This will push the statement into the list of statements for the current document.
-   * @private
-   * @param _doc
-   */
-  extract(_doc) {
-    // const res = { states: [], relations: [] };
-    let doc;
-    if (_doc.doc) {
-      doc = _doc.doc;
-    } else {
-      doc = _doc;
-    }
-    // let doc = root.doc;
-    // if (!doc) {
-    //   doc = root;
-    // }
-    log.info(doc);
-    this.clear(true);
-
-    log.info('Extract initial document:', doc);
-
-    doc.forEach((item) => {
-      log.warn('Statement', item.stmt);
-      switch (item.stmt) {
-        case STMT_STATE:
-          this.addState(
-            item.id.trim(),
-            item.type,
-            item.doc,
-            item.description,
-            item.note,
-            item.classes,
-            item.styles,
-            item.textStyles
-          );
-          break;
-        case STMT_RELATION:
-          this.addRelation(item.state1, item.state2, item.description);
-          break;
-        case STMT_CLASSDEF:
-          this.addStyleClass(item.id.trim(), item.classes);
-          break;
-        case STMT_STYLEDEF:
-          {
-            const ids = item.id.trim().split(',');
-            const styles = item.styleClass.split(',');
-            ids.forEach((id) => {
-              let foundState = this.getState(id);
-              if (foundState === undefined) {
-                const trimmedId = id.trim();
-                this.addState(trimmedId);
-                foundState = this.getState(trimmedId);
-              }
-              foundState.styles = styles.map((s) => s.replace(/;/g, '')?.trim());
-            });
-          }
-          break;
-        case STMT_APPLYCLASS:
-          this.setCssClass(item.id.trim(), item.styleClass);
-          break;
-      }
-    });
-
-    const diagramStates = this.getStates();
-    const config = getConfig();
-    const look = config.look;
-
-    resetDataFetching();
-    dataFetcher(
-      undefined,
-      this.getRootDocV2(),
-      diagramStates,
-      this.nodes,
-      this.edges,
-      true,
-      look,
-      this.classes
-    );
-    this.nodes.forEach((node) => {
-      if (Array.isArray(node.label)) {
-        // add the rest as description
-        node.description = node.label.slice(1);
-        if (node.isGroup && node.description.length > 0) {
-          throw new Error(
-            'Group nodes can only have label. Remove the additional description for node [' +
-              node.id +
-              ']'
-          );
-        }
-        // add first description as label
-        node.label = node.label[0];
-      }
-    });
-  }
-
-  /**
-   * Function called by parser when a node definition has been found.
-   *
-   * @param {null | string} id
-   * @param {null | string} type
-   * @param {null | string} doc
-   * @param {null | string | string[]} descr - description for the state. Can be a string or a list or strings
-   * @param {null | string} note
-   * @param {null | string | string[]} classes - class styles to apply to this state. Can be a string (1 style) or an array of styles. If it's just 1 class, convert it to an array of that 1 class.
-   * @param {null | string | string[]} styles - styles to apply to this state. Can be a string (1 style) or an array of styles. If it's just 1 style, convert it to an array of that 1 style.
-   * @param {null | string | string[]} textStyles - text styles to apply to this state. Can be a string (1 text test) or an array of text styles. If it's just 1 text style, convert it to an array of that 1 text style.
-   */
-  addState(
-    id,
-    type = DEFAULT_STATE_TYPE,
-    doc = null,
-    descr = null,
-    note = null,
-    classes = null,
-    styles = null,
-    textStyles = null
-  ) {
-    const trimmedId = id?.trim();
-    // add the state if needed
-    if (!this.currentDocument.states.has(trimmedId)) {
-      log.info('Adding state ', trimmedId, descr);
-      this.currentDocument.states.set(trimmedId, {
-        id: trimmedId,
-        descriptions: [],
-        type,
-        doc,
-        note,
-        classes: [],
-        styles: [],
-        textStyles: [],
-      });
-    } else {
-      if (!this.currentDocument.states.get(trimmedId).doc) {
-        this.currentDocument.states.get(trimmedId).doc = doc;
-      }
-      if (!this.currentDocument.states.get(trimmedId).type) {
-        this.currentDocument.states.get(trimmedId).type = type;
-      }
-    }
-
-    if (descr) {
-      log.info('Setting state description', trimmedId, descr);
-      if (typeof descr === 'string') {
-        this.addDescription(trimmedId, descr.trim());
-      }
-
-      if (typeof descr === 'object') {
-        descr.forEach((des) => this.addDescription(trimmedId, des.trim()));
-      }
-    }
-
-    if (note) {
-      const doc2 = this.currentDocument.states.get(trimmedId);
-      doc2.note = note;
-      doc2.note.text = common.sanitizeText(doc2.note.text, getConfig());
-    }
-
-    if (classes) {
-      log.info('Setting state classes', trimmedId, classes);
-      const classesList = typeof classes === 'string' ? [classes] : classes;
-      classesList.forEach((cssClass) => this.setCssClass(trimmedId, cssClass.trim()));
-    }
-
-    if (styles) {
-      log.info('Setting state styles', trimmedId, styles);
-      const stylesList = typeof styles === 'string' ? [styles] : styles;
-      stylesList.forEach((style) => this.setStyle(trimmedId, style.trim()));
-    }
-
-    if (textStyles) {
-      log.info('Setting state styles', trimmedId, styles);
-      const textStylesList = typeof textStyles === 'string' ? [textStyles] : textStyles;
-      textStylesList.forEach((textStyle) => this.setTextStyle(trimmedId, textStyle.trim()));
-    }
-  }
-
-  clear(saveCommon) {
-    this.nodes = [];
-    this.edges = [];
-    this.documents = {
-      root: newDoc(),
-    };
-    this.currentDocument = this.documents.root;
-
-    // number of start and end nodes; used to construct ids
-    this.startEndCount = 0;
-    this.classes = newClassesList();
-    if (!saveCommon) {
-      commonClear();
-    }
-  }
-
-  getState(id) {
-    return this.currentDocument.states.get(id);
-  }
-  getStates() {
-    return this.currentDocument.states;
-  }
-  logDocuments() {
-    log.info('Documents = ', this.documents);
-  }
-  getRelations() {
-    return this.currentDocument.relations;
-  }
-
-  /**
-   * If the id is a start node ( [*] ), then return a new id constructed from
-   * the start node name and the current start node count.
-   * else return the given id
-   *
-   * @param {string} id
-   * @returns {string} - the id (original or constructed)
-   * @private
-   */
-  startIdIfNeeded(id = '') {
-    let fixedId = id;
-    if (id === START_NODE) {
-      this.startEndCount++;
-      fixedId = `${START_TYPE}${this.startEndCount}`;
-    }
-    return fixedId;
-  }
-
-  /**
-   * If the id is a start node ( [*] ), then return the start type ('start')
-   * else return the given type
-   *
-   * @param {string} id
-   * @param {string} type
-   * @returns {string} - the type that should be used
-   * @private
-   */
-  startTypeIfNeeded(id = '', type = DEFAULT_STATE_TYPE) {
-    return id === START_NODE ? START_TYPE : type;
-  }
-
-  /**
-   * If the id is an end node ( [*] ), then return a new id constructed from
-   * the end node name and the current start_end node count.
-   * else return the given id
-   *
-   * @param {string} id
-   * @returns {string} - the id (original or constructed)
-   * @private
-   */
-  endIdIfNeeded(id = '') {
-    let fixedId = id;
-    if (id === END_NODE) {
-      this.startEndCount++;
-      fixedId = `${END_TYPE}${this.startEndCount}`;
-    }
-    return fixedId;
-  }
-
-  /**
-   * If the id is an end node ( [*] ), then return the end type
-   * else return the given type
-   *
-   * @param {string} id
-   * @param {string} type
-   * @returns {string} - the type that should be used
-   * @private
-   */
-  endTypeIfNeeded(id = '', type = DEFAULT_STATE_TYPE) {
-    return id === END_NODE ? END_TYPE : type;
-  }
-
-  /**
-   *
-   * @param item1
-   * @param item2
-   * @param relationTitle
-   */
-  addRelationObjs(item1, item2, relationTitle) {
-    let id1 = this.startIdIfNeeded(item1.id.trim());
-    let type1 = this.startTypeIfNeeded(item1.id.trim(), item1.type);
-    let id2 = this.startIdIfNeeded(item2.id.trim());
-    let type2 = this.startTypeIfNeeded(item2.id.trim(), item2.type);
-
-    this.addState(
-      id1,
-      type1,
-      item1.doc,
-      item1.description,
-      item1.note,
-      item1.classes,
-      item1.styles,
-      item1.textStyles
-    );
-    this.addState(
-      id2,
-      type2,
-      item2.doc,
-      item2.description,
-      item2.note,
-      item2.classes,
-      item2.styles,
-      item2.textStyles
-    );
-
-    this.currentDocument.relations.push({
-      id1,
-      id2,
-      relationTitle: common.sanitizeText(relationTitle, getConfig()),
-    });
-  }
-
-  /**
-   * Add a relation between two items.  The items may be full objects or just the string id of a state.
-   *
-   * @param {string | object} item1
-   * @param {string | object} item2
-   * @param {string} title
-   */
-  addRelation(item1, item2, title) {
-    if (typeof item1 === 'object') {
-      this.addRelationObjs(item1, item2, title);
-    } else {
-      const id1 = this.startIdIfNeeded(item1.trim());
-      const type1 = this.startTypeIfNeeded(item1);
-      const id2 = this.endIdIfNeeded(item2.trim());
-      const type2 = this.endTypeIfNeeded(item2);
-
-      this.addState(id1, type1);
-      this.addState(id2, type2);
-      this.currentDocument.relations.push({
-        id1,
-        id2,
-        title: common.sanitizeText(title, getConfig()),
-      });
-    }
-  }
-
-  addDescription(id, descr) {
-    const theState = this.currentDocument.states.get(id);
-    const _descr = descr.startsWith(':') ? descr.replace(':', '').trim() : descr;
-    theState.descriptions.push(common.sanitizeText(_descr, getConfig()));
-  }
-
-  cleanupLabel(label) {
-    if (label.substring(0, 1) === ':') {
-      return label.substr(2).trim();
-    } else {
-      return label.trim();
-    }
-  }
-
-  getDividerId() {
-    this.dividerCnt++;
-    return 'divider-id-' + this.dividerCnt;
-  }
-
-  /**
-   * Called when the parser comes across a (style) class definition
-   * @example classDef my-style fill:#f96;
-   *
-   * @param {string} id - the id of this (style) class
-   * @param  {string | null} styleAttributes - the string with 1 or more style attributes (each separated by a comma)
-   */
-  addStyleClass(id, styleAttributes = '') {
-    // create a new style class object with this id
-    if (!this.classes.has(id)) {
-      this.classes.set(id, { id: id, styles: [], textStyles: [] }); // This is a classDef
-    }
-    const foundClass = this.classes.get(id);
-    if (styleAttributes !== undefined && styleAttributes !== null) {
-      styleAttributes.split(STYLECLASS_SEP).forEach((attrib) => {
-        // remove any trailing ;
-        const fixedAttrib = attrib.replace(/([^;]*);/, '$1').trim();
-
-        // replace some style keywords
-        if (RegExp(COLOR_KEYWORD).exec(attrib)) {
-          const newStyle1 = fixedAttrib.replace(FILL_KEYWORD, BG_FILL);
-          const newStyle2 = newStyle1.replace(COLOR_KEYWORD, FILL_KEYWORD);
-          foundClass.textStyles.push(newStyle2);
-        }
-        foundClass.styles.push(fixedAttrib);
-      });
-    }
-  }
-
-  /**
-   * Return all of the style classes
-   * @returns {{} | any | classes}
-   */
-  getClasses() {
-    return this.classes;
-  }
-
-  /**
-   * Add a (style) class or css class to a state with the given id.
-   * If the state isn't already in the list of known states, add it.
-   * Might be called by parser when a style class or CSS class should be applied to a state
-   *
-   * @param {string | string[]} itemIds The id or a list of ids of the item(s) to apply the css class to
-   * @param {string} cssClassName CSS class name
-   */
-  setCssClass(itemIds, cssClassName) {
-    itemIds.split(',').forEach((id) => {
-      let foundState = this.getState(id);
-      if (foundState === undefined) {
-        const trimmedId = id.trim();
-        this.addState(trimmedId);
-        foundState = this.getState(trimmedId);
-      }
-      foundState.classes.push(cssClassName);
-    });
-  }
-
-  /**
-   * Add a style to a state with the given id.
-   * @example style stateId fill:#f9f,stroke:#333,stroke-width:4px
-   *   where 'style' is the keyword
-   *   stateId is the id of a state
-   *   the rest of the string is the styleText (all of the attributes to be applied to the state)
-   *
-   * @param itemId The id of item to apply the style to
-   * @param styleText - the text of the attributes for the style
-   */
-  setStyle(itemId, styleText) {
-    const item = this.getState(itemId);
-    if (item !== undefined) {
-      item.styles.push(styleText);
-    }
-  }
-
-  /**
-   * Add a text style to a state with the given id
-   *
-   * @param itemId The id of item to apply the css class to
-   * @param cssClassName CSS class name
-   */
-  setTextStyle(itemId, cssClassName) {
-    const item = this.getState(itemId);
-    if (item !== undefined) {
-      item.textStyles.push(cssClassName);
-    }
-  }
-
-  /**
-   * Finds the direction statement in the root document.
-   * @private
-   * @returns {{ value: string } | undefined} - the direction statement if present
-   */
-  getDirectionStatement() {
-    return this.rootDoc.find((doc) => doc.stmt === STMT_DIRECTION);
-  }
-
-  getDirection() {
-    return this.getDirectionStatement()?.value ?? DEFAULT_DIAGRAM_DIRECTION;
-  }
-
-  setDirection(dir) {
-    const doc = this.getDirectionStatement();
-    if (doc) {
-      doc.value = dir;
-    } else {
-      this.rootDoc.unshift({ stmt: STMT_DIRECTION, value: dir });
-    }
-  }
-
-  trimColon(str) {
-    return str && str[0] === ':' ? str.substr(1).trim() : str.trim();
-  }
-
-  getData() {
-    const config = getConfig();
-    return {
-      nodes: this.nodes,
-      edges: this.edges,
-      other: {},
-      config,
-      direction: getDir(this.getRootDocV2()),
-    };
-  }
-
-  getConfig() {
-    return getConfig().state;
-  }
-  getAccTitle = getAccTitle;
-  setAccTitle = setAccTitle;
-  getAccDescription = getAccDescription;
-  setAccDescription = setAccDescription;
-  setDiagramTitle = setDiagramTitle;
-  getDiagramTitle = getDiagramTitle;
-}

packages/mermaid/src/diagrams/state/stateDb.ts

@@ -0,0 +1,693 @@
+import { getConfig } from '../../diagram-api/diagramAPI.js';
+import { log } from '../../logger.js';
+import { generateId } from '../../utils.js';
+import common from '../common/common.js';
+import {
+  clear as commonClear,
+  getAccDescription,
+  getAccTitle,
+  getDiagramTitle,
+  setAccDescription,
+  setAccTitle,
+  setDiagramTitle,
+} from '../common/commonDb.js';
+import { dataFetcher, reset as resetDataFetcher } from './dataFetcher.js';
+import { getDir } from './stateRenderer-v3-unified.js';
+import {
+  DEFAULT_DIAGRAM_DIRECTION,
+  DEFAULT_STATE_TYPE,
+  DIVIDER_TYPE,
+  STMT_APPLYCLASS,
+  STMT_CLASSDEF,
+  STMT_RELATION,
+  STMT_ROOT,
+  STMT_DIRECTION,
+  STMT_STATE,
+  STMT_STYLEDEF,
+} from './stateCommon.js';
+import type { MermaidConfig } from '../../config.type.js';
+
+const CONSTANTS = {
+  START_NODE: '[*]',
+  START_TYPE: 'start',
+  END_NODE: '[*]',
+  END_TYPE: 'end',
+  COLOR_KEYWORD: 'color',
+  FILL_KEYWORD: 'fill',
+  BG_FILL: 'bgFill',
+  STYLECLASS_SEP: ',',
+} as const;
+
+interface BaseStmt {
+  stmt: 'applyClass' | 'classDef' | 'dir' | 'relation' | 'state' | 'style' | 'root' | 'default';
+}
+
+interface ApplyClassStmt extends BaseStmt {
+  stmt: 'applyClass';
+  id: string;
+  styleClass: string;
+}
+
+interface ClassDefStmt extends BaseStmt {
+  stmt: 'classDef';
+  id: string;
+  classes: string;
+}
+
+interface DirectionStmt extends BaseStmt {
+  stmt: 'dir';
+  value: 'TB' | 'BT' | 'RL' | 'LR';
+}
+
+interface RelationStmt extends BaseStmt {
+  stmt: 'relation';
+  state1: StateStmt;
+  state2: StateStmt;
+  description?: string;
+}
+
+export interface StateStmt extends BaseStmt {
+  stmt: 'state' | 'default';
+  id: string;
+  type: 'default' | 'fork' | 'join' | 'choice' | 'divider' | 'start' | 'end';
+  description?: string;
+  descriptions?: string[];
+  doc?: Stmt[];
+  note?: Note;
+  start?: boolean;
+  classes?: string[];
+  styles?: string[];
+  textStyles?: string[];
+}
+
+interface StyleStmt extends BaseStmt {
+  stmt: 'style';
+  id: string;
+  styleClass: string;
+}
+
+export interface RootStmt {
+  id: 'root';
+  stmt: 'root';
+  doc?: Stmt[];
+}
+
+interface Note {
+  position?: 'left of' | 'right of';
+  text: string;
+}
+
+export type Stmt =
+  | ApplyClassStmt
+  | ClassDefStmt
+  | DirectionStmt
+  | RelationStmt
+  | StateStmt
+  | StyleStmt
+  | RootStmt;
+
+interface DiagramEdge {
+  id1: string;
+  id2: string;
+  relationTitle?: string;
+}
+
+interface Document {
+  relations: DiagramEdge[];
+  states: Map<string, StateStmt>;
+  documents: Record<string, Document>;
+}
+
+export interface StyleClass {
+  id: string;
+  styles: string[];
+  textStyles: string[];
+}
+
+export interface NodeData {
+  labelStyle?: string;
+  shape: string;
+  label?: string | string[];
+  cssClasses: string;
+  cssCompiledStyles?: string[];
+  cssStyles: string[];
+  id: string;
+  dir?: string;
+  domId?: string;
+  type?: string;
+  isGroup?: boolean;
+  padding?: number;
+  rx?: number;
+  ry?: number;
+  look?: MermaidConfig['look'];
+  parentId?: string;
+  centerLabel?: boolean;
+  position?: string;
+  description?: string | string[];
+}
+
+export interface Edge {
+  id: string;
+  start: string;
+  end: string;
+  arrowhead: string;
+  arrowTypeEnd: string;
+  style: string;
+  labelStyle: string;
+  label?: string;
+  arrowheadStyle: string;
+  labelpos: string;
+  labelType: string;
+  thickness: string;
+  classes: string;
+  look: MermaidConfig['look'];
+}
+
+/**
+ * Returns a new list of classes.
+ * In the future, this can be replaced with a class common to all diagrams.
+ * ClassDef information = \{ id: id, styles: [], textStyles: [] \}
+ */
+const newClassesList = (): Map<string, StyleClass> => new Map();
+const newDoc = (): Document => ({
+  relations: [],
+  states: new Map(),
+  documents: {},
+});
+const clone = <T>(o: T): T => JSON.parse(JSON.stringify(o));
+
+export class StateDB {
+  private nodes: NodeData[] = [];
+  private edges: Edge[] = [];
+  private rootDoc: Stmt[] = [];
+  private classes = newClassesList();
+  private documents = { root: newDoc() };
+  private currentDocument = this.documents.root;
+  private startEndCount = 0;
+  private dividerCnt = 0;
+
+  static readonly relationType = {
+    AGGREGATION: 0,
+    EXTENSION: 1,
+    COMPOSITION: 2,
+    DEPENDENCY: 3,
+  } as const;
+
+  constructor(private version: 1 | 2) {
+    this.clear();
+    // Bind methods used by JISON
+    this.setRootDoc = this.setRootDoc.bind(this);
+    this.getDividerId = this.getDividerId.bind(this);
+    this.setDirection = this.setDirection.bind(this);
+    this.trimColon = this.trimColon.bind(this);
+  }
+
+  /**
+   * Convert all of the statements (stmts) that were parsed into states and relationships.
+   * This is done because a state diagram may have nested sections,
+   * where each section is a 'document' and has its own set of statements.
+   * Ex: the section within a fork has its own statements, and incoming and outgoing statements
+   * refer to the fork as a whole (document).
+   * See the parser grammar:  the definition of a document is a document then a 'line', where a line can be a statement.
+   * This will push the statement into the list of statements for the current document.
+   */
+  extract(statements: Stmt[] | { doc: Stmt[] }) {
+    this.clear(true);
+    for (const item of Array.isArray(statements) ? statements : statements.doc) {
+      switch (item.stmt) {
+        case STMT_STATE:
+          this.addState(item.id.trim(), item.type, item.doc, item.description, item.note);
+          break;
+        case STMT_RELATION:
+          this.addRelation(item.state1, item.state2, item.description);
+          break;
+        case STMT_CLASSDEF:
+          this.addStyleClass(item.id.trim(), item.classes);
+          break;
+        case STMT_STYLEDEF:
+          this.handleStyleDef(item);
+          break;
+        case STMT_APPLYCLASS:
+          this.setCssClass(item.id.trim(), item.styleClass);
+          break;
+      }
+    }
+    const diagramStates = this.getStates();
+    const config = getConfig();
+
+    resetDataFetcher();
+    dataFetcher(
+      undefined,
+      this.getRootDocV2() as StateStmt,
+      diagramStates,
+      this.nodes,
+      this.edges,
+      true,
+      config.look,
+      this.classes
+    );
+
+    // Process node labels
+    for (const node of this.nodes) {
+      if (!Array.isArray(node.label)) {
+        continue;
+      }
+
+      node.description = node.label.slice(1);
+      if (node.isGroup && node.description.length > 0) {
+        throw new Error(
+          `Group nodes can only have label. Remove the additional description for node [${node.id}]`
+        );
+      }
+      node.label = node.label[0];
+    }
+  }
+
+  private handleStyleDef(item: StyleStmt) {
+    const ids = item.id.trim().split(',');
+    const styles = item.styleClass.split(',');
+
+    for (const id of ids) {
+      let state = this.getState(id);
+      if (!state) {
+        const trimmedId = id.trim();
+        this.addState(trimmedId);
+        state = this.getState(trimmedId);
+      }
+      if (state) {
+        state.styles = styles.map((s) => s.replace(/;/g, '')?.trim());
+      }
+    }
+  }
+
+  setRootDoc(o: Stmt[]) {
+    log.info('Setting root doc', o);
+    this.rootDoc = o;
+    if (this.version === 1) {
+      this.extract(o);
+    } else {
+      this.extract(this.getRootDocV2());
+    }
+  }
+
+  docTranslator(parent: RootStmt | StateStmt, node: Stmt, first: boolean) {
+    if (node.stmt === STMT_RELATION) {
+      this.docTranslator(parent, node.state1, true);
+      this.docTranslator(parent, node.state2, false);
+      return;
+    }
+
+    if (node.stmt === STMT_STATE) {
+      if (node.id === CONSTANTS.START_NODE) {
+        node.id = parent.id + (first ? '_start' : '_end');
+        node.start = first;
+      } else {
+        // This is just a plain state, not a start or end
+        node.id = node.id.trim();
+      }
+    }
+
+    if ((node.stmt !== STMT_ROOT && node.stmt !== STMT_STATE) || !node.doc) {
+      return;
+    }
+
+    const doc = [];
+    // Check for concurrency
+    let currentDoc = [];
+    for (const stmt of node.doc) {
+      if ((stmt as StateStmt).type === DIVIDER_TYPE) {
+        const newNode = clone(stmt as StateStmt);
+        newNode.doc = clone(currentDoc);
+        doc.push(newNode);
+        currentDoc = [];
+      } else {
+        currentDoc.push(stmt);
+      }
+    }
+
+    // If any divider was encountered
+    if (doc.length > 0 && currentDoc.length > 0) {
+      const newNode = {
+        stmt: STMT_STATE,
+        id: generateId(),
+        type: 'divider',
+        doc: clone(currentDoc),
+      } satisfies StateStmt;
+      doc.push(clone(newNode));
+      node.doc = doc;
+    }
+
+    node.doc.forEach((docNode) => this.docTranslator(node, docNode, true));
+  }
+
+  private getRootDocV2() {
+    this.docTranslator(
+      { id: STMT_ROOT, stmt: STMT_ROOT },
+      { id: STMT_ROOT, stmt: STMT_ROOT, doc: this.rootDoc },
+      true
+    );
+    return { id: STMT_ROOT, doc: this.rootDoc };
+  }
+
+  /**
+   * Function called by parser when a node definition has been found.
+   *
+   * @param descr - description for the state. Can be a string or a list or strings
+   * @param classes - class styles to apply to this state. Can be a string (1 style) or an array of styles. If it's just 1 class, convert it to an array of that 1 class.
+   * @param styles - styles to apply to this state. Can be a string (1 style) or an array of styles. If it's just 1 style, convert it to an array of that 1 style.
+   * @param textStyles - text styles to apply to this state. Can be a string (1 text test) or an array of text styles. If it's just 1 text style, convert it to an array of that 1 text style.
+   */
+  addState(
+    id: string,
+    type: StateStmt['type'] = DEFAULT_STATE_TYPE,
+    doc: Stmt[] | undefined = undefined,
+    descr: string | string[] | undefined = undefined,
+    note: Note | undefined = undefined,
+    classes: string | string[] | undefined = undefined,
+    styles: string | string[] | undefined = undefined,
+    textStyles: string | string[] | undefined = undefined
+  ) {
+    const trimmedId = id?.trim();
+    if (!this.currentDocument.states.has(trimmedId)) {
+      log.info('Adding state ', trimmedId, descr);
+      this.currentDocument.states.set(trimmedId, {
+        stmt: STMT_STATE,
+        id: trimmedId,
+        descriptions: [],
+        type,
+        doc,
+        note,
+        classes: [],
+        styles: [],
+        textStyles: [],
+      });
+    } else {
+      const state = this.currentDocument.states.get(trimmedId);
+      if (!state) {
+        throw new Error(`State not found: ${trimmedId}`);
+      }
+      if (!state.doc) {
+        state.doc = doc;
+      }
+      if (!state.type) {
+        state.type = type;
+      }
+    }
+
+    if (descr) {
+      log.info('Setting state description', trimmedId, descr);
+      const descriptions = Array.isArray(descr) ? descr : [descr];
+      descriptions.forEach((des) => this.addDescription(trimmedId, des.trim()));
+    }
+
+    if (note) {
+      const doc2 = this.currentDocument.states.get(trimmedId);
+      if (!doc2) {
+        throw new Error(`State not found: ${trimmedId}`);
+      }
+      doc2.note = note;
+      doc2.note.text = common.sanitizeText(doc2.note.text, getConfig());
+    }
+
+    if (classes) {
+      log.info('Setting state classes', trimmedId, classes);
+      const classesList = Array.isArray(classes) ? classes : [classes];
+      classesList.forEach((cssClass) => this.setCssClass(trimmedId, cssClass.trim()));
+    }
+
+    if (styles) {
+      log.info('Setting state styles', trimmedId, styles);
+      const stylesList = Array.isArray(styles) ? styles : [styles];
+      stylesList.forEach((style) => this.setStyle(trimmedId, style.trim()));
+    }
+
+    if (textStyles) {
+      log.info('Setting state styles', trimmedId, styles);
+      const textStylesList = Array.isArray(textStyles) ? textStyles : [textStyles];
+      textStylesList.forEach((textStyle) => this.setTextStyle(trimmedId, textStyle.trim()));
+    }
+  }
+
+  clear(saveCommon?: boolean) {
+    this.nodes = [];
+    this.edges = [];
+    this.documents = { root: newDoc() };
+    this.currentDocument = this.documents.root;
+
+    // number of start and end nodes; used to construct ids
+    this.startEndCount = 0;
+    this.classes = newClassesList();
+    if (!saveCommon) {
+      commonClear();
+    }
+  }
+
+  getState(id: string) {
+    return this.currentDocument.states.get(id);
+  }
+
+  getStates() {
+    return this.currentDocument.states;
+  }
+
+  logDocuments() {
+    log.info('Documents = ', this.documents);
+  }
+
+  getRelations() {
+    return this.currentDocument.relations;
+  }
+
+  /**
+   * If the id is a start node ( [*] ), then return a new id constructed from
+   * the start node name and the current start node count.
+   * else return the given id
+   */
+  startIdIfNeeded(id = '') {
+    if (id === CONSTANTS.START_NODE) {
+      this.startEndCount++;
+      return `${CONSTANTS.START_TYPE}${this.startEndCount}`;
+    }
+    return id;
+  }
+
+  /**
+   * If the id is a start node ( [*] ), then return the start type ('start')
+   * else return the given type
+   */
+  startTypeIfNeeded(id = '', type: StateStmt['type'] = DEFAULT_STATE_TYPE) {
+    return id === CONSTANTS.START_NODE ? CONSTANTS.START_TYPE : type;
+  }
+
+  /**
+   * If the id is an end node ( [*] ), then return a new id constructed from
+   * the end node name and the current start_end node count.
+   * else return the given id
+   */
+  endIdIfNeeded(id = '') {
+    if (id === CONSTANTS.END_NODE) {
+      this.startEndCount++;
+      return `${CONSTANTS.END_TYPE}${this.startEndCount}`;
+    }
+    return id;
+  }
+
+  /**
+   * If the id is an end node ( [*] ), then return the end type
+   * else return the given type
+   *
+   */
+  endTypeIfNeeded(id = '', type: StateStmt['type'] = DEFAULT_STATE_TYPE) {
+    return id === CONSTANTS.END_NODE ? CONSTANTS.END_TYPE : type;
+  }
+
+  addRelationObjs(item1: StateStmt, item2: StateStmt, relationTitle = '') {
+    const id1 = this.startIdIfNeeded(item1.id.trim());
+    const type1 = this.startTypeIfNeeded(item1.id.trim(), item1.type);
+    const id2 = this.startIdIfNeeded(item2.id.trim());
+    const type2 = this.startTypeIfNeeded(item2.id.trim(), item2.type);
+    this.addState(
+      id1,
+      type1,
+      item1.doc,
+      item1.description,
+      item1.note,
+      item1.classes,
+      item1.styles,
+      item1.textStyles
+    );
+    this.addState(
+      id2,
+      type2,
+      item2.doc,
+      item2.description,
+      item2.note,
+      item2.classes,
+      item2.styles,
+      item2.textStyles
+    );
+    this.currentDocument.relations.push({
+      id1,
+      id2,
+      relationTitle: common.sanitizeText(relationTitle, getConfig()),
+    });
+  }
+
+  /**
+   * Add a relation between two items.  The items may be full objects or just the string id of a state.
+   */
+  addRelation(item1: string | StateStmt, item2: string | StateStmt, title?: string) {
+    if (typeof item1 === 'object' && typeof item2 === 'object') {
+      this.addRelationObjs(item1, item2, title);
+    } else if (typeof item1 === 'string' && typeof item2 === 'string') {
+      const id1 = this.startIdIfNeeded(item1.trim());
+      const type1 = this.startTypeIfNeeded(item1);
+      const id2 = this.endIdIfNeeded(item2.trim());
+      const type2 = this.endTypeIfNeeded(item2);
+
+      this.addState(id1, type1);
+      this.addState(id2, type2);
+      this.currentDocument.relations.push({
+        id1,
+        id2,
+        relationTitle: title ? common.sanitizeText(title, getConfig()) : undefined,
+      });
+    }
+  }
+
+  addDescription(id: string, descr: string) {
+    const theState = this.currentDocument.states.get(id);
+    const _descr = descr.startsWith(':') ? descr.replace(':', '').trim() : descr;
+    theState?.descriptions?.push(common.sanitizeText(_descr, getConfig()));
+  }
+
+  cleanupLabel(label: string) {
+    return label.startsWith(':') ? label.slice(2).trim() : label.trim();
+  }
+
+  getDividerId() {
+    this.dividerCnt++;
+    return `divider-id-${this.dividerCnt}`;
+  }
+
+  /**
+   * Called when the parser comes across a (style) class definition
+   * @example classDef my-style fill:#f96;
+   *
+   * @param id - the id of this (style) class
+   * @param styleAttributes - the string with 1 or more style attributes (each separated by a comma)
+   */
+  addStyleClass(id: string, styleAttributes = '') {
+    // create a new style class object with this id
+    if (!this.classes.has(id)) {
+      this.classes.set(id, { id, styles: [], textStyles: [] });
+    }
+    const foundClass = this.classes.get(id);
+    if (styleAttributes && foundClass) {
+      styleAttributes.split(CONSTANTS.STYLECLASS_SEP).forEach((attrib) => {
+        const fixedAttrib = attrib.replace(/([^;]*);/, '$1').trim();
+        if (RegExp(CONSTANTS.COLOR_KEYWORD).exec(attrib)) {
+          const newStyle1 = fixedAttrib.replace(CONSTANTS.FILL_KEYWORD, CONSTANTS.BG_FILL);
+          const newStyle2 = newStyle1.replace(CONSTANTS.COLOR_KEYWORD, CONSTANTS.FILL_KEYWORD);
+          foundClass.textStyles.push(newStyle2);
+        }
+        foundClass.styles.push(fixedAttrib);
+      });
+    }
+  }
+
+  getClasses() {
+    return this.classes;
+  }
+
+  /**
+   * Add a (style) class or css class to a state with the given id.
+   * If the state isn't already in the list of known states, add it.
+   * Might be called by parser when a style class or CSS class should be applied to a state
+   *
+   * @param itemIds - The id or a list of ids of the item(s) to apply the css class to
+   * @param cssClassName - CSS class name
+   */
+  setCssClass(itemIds: string, cssClassName: string) {
+    itemIds.split(',').forEach((id) => {
+      let foundState = this.getState(id);
+      if (!foundState) {
+        const trimmedId = id.trim();
+        this.addState(trimmedId);
+        foundState = this.getState(trimmedId);
+      }
+      foundState?.classes?.push(cssClassName);
+    });
+  }
+
+  /**
+   * Add a style to a state with the given id.
+   * @example style stateId fill:#f9f,stroke:#333,stroke-width:4px
+   *   where 'style' is the keyword
+   *   stateId is the id of a state
+   *   the rest of the string is the styleText (all of the attributes to be applied to the state)
+   *
+   * @param itemId - The id of item to apply the style to
+   * @param styleText - the text of the attributes for the style
+   */
+  setStyle(itemId: string, styleText: string) {
+    this.getState(itemId)?.styles?.push(styleText);
+  }
+
+  /**
+   * Add a text style to a state with the given id
+   *
+   * @param itemId - The id of item to apply the css class to
+   * @param cssClassName - CSS class name
+   */
+  setTextStyle(itemId: string, cssClassName: string) {
+    this.getState(itemId)?.textStyles?.push(cssClassName);
+  }
+
+  /**
+   * Finds the direction statement in the root document.
+   * @returns the direction statement if present
+   */
+  private getDirectionStatement() {
+    return this.rootDoc.find((doc): doc is DirectionStmt => doc.stmt === STMT_DIRECTION);
+  }
+
+  getDirection() {
+    return this.getDirectionStatement()?.value ?? DEFAULT_DIAGRAM_DIRECTION;
+  }
+
+  setDirection(dir: DirectionStmt['value']) {
+    const doc = this.getDirectionStatement();
+    if (doc) {
+      doc.value = dir;
+    } else {
+      this.rootDoc.unshift({ stmt: STMT_DIRECTION, value: dir });
+    }
+  }
+
+  trimColon(str: string) {
+    return str.startsWith(':') ? str.slice(1).trim() : str.trim();
+  }
+
+  getData() {
+    const config = getConfig();
+    return {
+      nodes: this.nodes,
+      edges: this.edges,
+      other: {},
+      config,
+      direction: getDir(this.getRootDocV2()),
+    };
+  }
+
+  getConfig() {
+    return getConfig().state;
+  }
+
+  getAccTitle = getAccTitle;
+  setAccTitle = setAccTitle;
+  getAccDescription = getAccDescription;
+  setAccDescription = setAccDescription;
+  setDiagramTitle = setDiagramTitle;
+  getDiagramTitle = getDiagramTitle;
+}

packages/mermaid/src/diagrams/state/stateDiagram-v2.spec.js

@@ -5,6 +5,7 @@ import { StateDB } from './stateDb.js';
 describe('state diagram V2, ', function () {
   // TODO - these examples should be put into ./parser/stateDiagram.spec.js
   describe('when parsing an info graph it', function () {
+    /** @type {StateDB} */
     let stateDb;
     beforeEach(function () {
       stateDb = new StateDB(2);
@@ -347,6 +348,20 @@ describe('state diagram V2, ', function () {
       `;
 
       parser.parse(str);
+      expect(stateDb.getState('Active').note).toMatchInlineSnapshot(`
+        {
+          "position": "left of",
+          "text": "this is a short<br>note",
+        }
+      `);
+      expect(stateDb.getState('Inactive').note).toMatchInlineSnapshot(`
+        {
+          "position": "right of",
+          "text": "A note can also
+              be defined on
+              several lines",
+        }
+      `);
     });
     it('should handle multiline notes with different line breaks', function () {
       const str = `stateDiagram-v2
@@ -357,6 +372,12 @@ describe('state diagram V2, ', function () {
       `;
 
       parser.parse(str);
+      expect(stateDb.getStates().get('State1').note).toMatchInlineSnapshot(`
+        {
+          "position": "right of",
+          "text": "Line1<br>Line2<br>Line3<br>Line4<br>Line5",
+        }
+      `);
     });
     it('should handle floating notes', function () {
       const str = `stateDiagram-v2
@@ -367,15 +388,14 @@ describe('state diagram V2, ', function () {
       parser.parse(str);
     });
     it('should handle floating notes', function () {
-      const str = `stateDiagram-v2\n
+      const str = `stateDiagram-v2
       state foo
       note "This is a floating note" as N1
       `;
-
       parser.parse(str);
     });
     it('should handle notes for composite (nested) states', function () {
-      const str = `stateDiagram-v2\n
+      const str = `stateDiagram-v2
       [*] --> NotShooting
 
       state "Not Shooting State" as NotShooting {
@@ -390,6 +410,12 @@ describe('state diagram V2, ', function () {
       `;
 
       parser.parse(str);
+      expect(stateDb.getState('NotShooting').note).toMatchInlineSnapshot(`
+        {
+          "position": "right of",
+          "text": "This is a note on a composite state",
+        }
+      `);
     });
 
     it('A composite state should be able to link to itself', () => {