From 53bb5eb02499bee481bfd45949b885942b7e6e89 Mon Sep 17 00:00:00 2001
From: Yash Singh <saiansh2525@gmail.com>
Date: Wed, 1 May 2024 10:42:55 -0700
Subject: [PATCH] fix: some changes

---
 packages/mermaid/src/diagrams/er/erRenderer.js          | 2 +-
 packages/mermaid/src/diagrams/state/stateRenderer-v2.js | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/packages/mermaid/src/diagrams/er/erRenderer.js b/packages/mermaid/src/diagrams/er/erRenderer.js
index 7f0a0e914a..33fb5bd4ae 100644
--- a/packages/mermaid/src/diagrams/er/erRenderer.js
+++ b/packages/mermaid/src/diagrams/er/erRenderer.js
@@ -296,7 +296,7 @@ const drawAttributes = (groupNode, entityTextNode, attributes) => {
  * Use D3 to construct the svg elements for the entities
  *
  * @param svgNode The svg node that contains the diagram
- * @param entities The entities to be drawn
+ * @param {Map<string, object>} entities The entities to be drawn
  * @param graph The graph that contains the vertex and edge definitions post-layout
  * @returns {object} The first entity that was inserted
  */
diff --git a/packages/mermaid/src/diagrams/state/stateRenderer-v2.js b/packages/mermaid/src/diagrams/state/stateRenderer-v2.js
index 4d702bf3c1..264d4501f9 100644
--- a/packages/mermaid/src/diagrams/state/stateRenderer-v2.js
+++ b/packages/mermaid/src/diagrams/state/stateRenderer-v2.js
@@ -129,13 +129,13 @@ export function stateDomId(itemId = '', counter = 0, type = '', typeSpacer = DOM
  * @param g - graph
  * @param {object} parent
  * @param {object} parsedItem - parsed statement item
- * @param {object[]} diagramStates - the list of all known  states for the diagram
+ * @param {Map<string, object>} diagramStates - the list of all known  states for the diagram
  * @param {object} diagramDb
  * @param {boolean} altFlag - for clusters, add the "statediagram-cluster-alt" CSS class
  */
 const setupNode = (g, parent, parsedItem, diagramStates, diagramDb, altFlag) => {
   const itemId = parsedItem.id;
-  const classStr = getClassesFromDbInfo(diagramStates[itemId]);
+  const classStr = getClassesFromDbInfo(diagramStates.get(itemId));
 
   if (itemId !== 'root') {
     let shape = SHAPE_STATE;
@@ -303,7 +303,7 @@ const setupNode = (g, parent, parsedItem, diagramStates, diagramDb, altFlag) =>
  * @param g
  * @param parentParsedItem - parsed Item that is the parent of this document (doc)
  * @param doc - the document to set up; it is a list of parsed statements
- * @param {object[]} diagramStates - the list of all known states for the diagram
+ * @param {Map<string, object>} diagramStates - the list of all known states for the diagram
  * @param diagramDb
  * @param {boolean} altFlag
  * @todo This duplicates some of what is done in stateDb.js extract method