From 8e3f8b1a57b5fb4ab277ff235961a9ef768871ca Mon Sep 17 00:00:00 2001
From: Nikolay Rozhkov <nironame@gmail.com>
Date: Thu, 29 Jun 2023 15:39:50 +0300
Subject: [PATCH] Fixes after review

---
 .../mermaid/src/diagrams/sankey/sankeyDB.ts   |  9 ++++----
 .../src/diagrams/sankey/sankeyRenderer.ts     | 22 ++++++++++---------
 .../mermaid/src/docs/.vitepress/config.ts     |  1 +
 3 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/packages/mermaid/src/diagrams/sankey/sankeyDB.ts b/packages/mermaid/src/diagrams/sankey/sankeyDB.ts
index 1b6864af9..1921e1b85 100644
--- a/packages/mermaid/src/diagrams/sankey/sankeyDB.ts
+++ b/packages/mermaid/src/diagrams/sankey/sankeyDB.ts
@@ -43,12 +43,11 @@ class SankeyNode {
 
 const findOrCreateNode = (ID: string): SankeyNode => {
   ID = common.sanitizeText(ID, configApi.getConfig());
-  if (nodesMap[ID]) {
-    return nodesMap[ID];
-  }
 
-  nodesMap[ID] = new SankeyNode(ID);
-  nodes.push(nodesMap[ID]);
+  if (!nodesMap[ID]) {
+    nodesMap[ID] = new SankeyNode(ID);
+    nodes.push(nodesMap[ID]);
+  }
   return nodesMap[ID];
 };
 
diff --git a/packages/mermaid/src/diagrams/sankey/sankeyRenderer.ts b/packages/mermaid/src/diagrams/sankey/sankeyRenderer.ts
index cb3092f55..780144be0 100644
--- a/packages/mermaid/src/diagrams/sankey/sankeyRenderer.ts
+++ b/packages/mermaid/src/diagrams/sankey/sankeyRenderer.ts
@@ -20,6 +20,17 @@ import { configureSvgSize } from '../../setupGraphViewbox.js';
 import { Uid } from '../../rendering-util/uid.js';
 import { SankeyLinkColor, SankeyNodeAlignment } from '../../config.type.js';
 
+// Map config options to alignment functions
+const alignmentsMap: Record<
+  SankeyNodeAlignment,
+  (node: d3SankeyNode<object, object>, n: number) => number
+> = {
+  [SankeyNodeAlignment.left]: d3SankeyLeft,
+  [SankeyNodeAlignment.right]: d3SankeyRight,
+  [SankeyNodeAlignment.center]: d3SankeyCenter,
+  [SankeyNodeAlignment.justify]: d3SankeyJustify,
+};
+
 /**
  * Draws Sankey diagram.
  *
@@ -73,16 +84,7 @@ export const draw = function (text: string, id: string, _version: string, diagOb
   // @ts-ignore TODO: db should be coerced to sankey DB type
   const graph = diagObj.db.getGraph();
 
-  // Map config options to alignment functions
-  const alignmentsMap: Record<
-    SankeyNodeAlignment,
-    (node: d3SankeyNode<object, object>, n: number) => number
-  > = {
-    [SankeyNodeAlignment.left]: d3SankeyLeft,
-    [SankeyNodeAlignment.right]: d3SankeyRight,
-    [SankeyNodeAlignment.center]: d3SankeyCenter,
-    [SankeyNodeAlignment.justify]: d3SankeyJustify,
-  };
+  // Get alignment function
   const nodeAlign = alignmentsMap[nodeAlignment];
 
   // Construct and configure a Sankey generator
diff --git a/packages/mermaid/src/docs/.vitepress/config.ts b/packages/mermaid/src/docs/.vitepress/config.ts
index 8fceb810b..2b2914ca3 100644
--- a/packages/mermaid/src/docs/.vitepress/config.ts
+++ b/packages/mermaid/src/docs/.vitepress/config.ts
@@ -138,6 +138,7 @@ function sidebarSyntax() {
         { text: 'Mindmaps 🔥', link: '/syntax/mindmap' },
         { text: 'Timeline 🔥', link: '/syntax/timeline' },
         { text: 'Zenuml 🔥', link: '/syntax/zenuml' },
+        { text: 'Sankey 🔥', link: '/syntax/sankey' },
         { text: 'Other Examples', link: '/syntax/examples' },
       ],
     },