From fa6bcd8b30456b8cc1bda249631911bf315ba5cb Mon Sep 17 00:00:00 2001
From: Knut Sveidqvist <knsv@sveido.com>
Date: Fri, 24 May 2024 13:28:33 +0200
Subject: [PATCH] #5237 Adding new rendering for flowcharts

---
 cypress/platform/knsv2.html                   | 12 ++-
 packages/mermaid/src/dagre-wrapper/nodes.js   |  2 -
 .../mermaid/src/diagrams/flowchart/flowDb.ts  | 43 ++++++++++
 .../src/diagrams/flowchart/flowDiagram-v2.ts  |  8 +-
 .../flowchart/flowDiagram-v3-unified.ts       | 25 ++++++
 .../flowchart/flowRenderer-v3-unified.ts      | 81 +++++++++++++++++++
 .../mermaid/src/diagrams/state/stateDb.js     |  1 +
 .../rendering-elements/nodes.js               |  9 ++-
 .../shapes/{rect.ts => drawRect.ts}           |  8 +-
 .../rendering-elements/shapes/roundedRect.ts  | 13 +++
 .../rendering-elements/shapes/squareRect.ts   | 11 +++
 .../rendering-elements/shapes/state.ts        | 11 +++
 .../mermaid/src/rendering-util/types.d.ts     |  8 +-
 13 files changed, 216 insertions(+), 16 deletions(-)
 create mode 100644 packages/mermaid/src/diagrams/flowchart/flowDiagram-v3-unified.ts
 create mode 100644 packages/mermaid/src/diagrams/flowchart/flowRenderer-v3-unified.ts
 rename packages/mermaid/src/rendering-util/rendering-elements/shapes/{rect.ts => drawRect.ts} (94%)
 create mode 100644 packages/mermaid/src/rendering-util/rendering-elements/shapes/roundedRect.ts
 create mode 100644 packages/mermaid/src/rendering-util/rendering-elements/shapes/squareRect.ts
 create mode 100644 packages/mermaid/src/rendering-util/rendering-elements/shapes/state.ts

diff --git a/cypress/platform/knsv2.html b/cypress/platform/knsv2.html
index 6ce618046..1c46866c5 100644
--- a/cypress/platform/knsv2.html
+++ b/cypress/platform/knsv2.html
@@ -136,15 +136,19 @@ sequenceDiagram
       </pre
     >
 
-    <pre id="diagram" class="mermaid">
+    <pre id="diagram" class="mermaid2">
       %%{init: {"layout": "elk", "mergeEdges": true} }%%
 stateDiagram
-  direction TB
-  T00 --> T0
-  T00 --> T1
+      A --> B
       </pre
     >
     <pre id="diagram" class="mermaid">
+      %%{init: {"layout": "elk", "mergeEdges": true} }%%
+flowchart
+      A --> B(This is B)
+      </pre
+    >
+    <pre id="diagram" class="mermaid2">
       %%{init: {"layout": "elk", "mergeEdges": false, "elk.nodePlacement.strategy": "NETWORK_SIMPLEX"} }%%
 stateDiagram
   State T0 {
diff --git a/packages/mermaid/src/dagre-wrapper/nodes.js b/packages/mermaid/src/dagre-wrapper/nodes.js
index 66117369c..cb5b35be6 100644
--- a/packages/mermaid/src/dagre-wrapper/nodes.js
+++ b/packages/mermaid/src/dagre-wrapper/nodes.js
@@ -1135,8 +1135,6 @@ export const insertNode = async (elem, node, dir) => {
   let newEl;
   let el;
 
-  console.log('insertNode element', elem, elem.node());
-  // debugger;
   // Add link when appropriate
   if (node.link) {
     let target;
diff --git a/packages/mermaid/src/diagrams/flowchart/flowDb.ts b/packages/mermaid/src/diagrams/flowchart/flowDb.ts
index 797130e71..ef4ae0a76 100644
--- a/packages/mermaid/src/diagrams/flowchart/flowDb.ts
+++ b/packages/mermaid/src/diagrams/flowchart/flowDb.ts
@@ -2,6 +2,7 @@ import { select } from 'd3';
 import utils from '../../utils.js';
 import { getConfig, defaultConfig } from '../../diagram-api/diagramAPI.js';
 import common from '../common/common.js';
+import type { LayoutData, LayoutMethod, Node, Edge } from '../../rendering-util/types.js';
 import { log } from '../../logger.js';
 import {
   setAccTitle,
@@ -755,11 +756,53 @@ export const lex = {
   firstGraph,
 };
 
+const getTypeFromVertex = (vertex: FlowVertex) => {
+  if (vertex.type === 'square') {
+    return 'squareRect';
+  }
+  if (vertex.type === 'round') {
+    return 'roundedRect';
+  }
+
+  return vertex.type || 'squareRect';
+};
+
+export const getData = () => {
+  const config = getConfig();
+  const nodes: Node[] = [];
+  const edges: Edge[] = [];
+
+  // extract(getRootDocV2());
+  // const diagramStates = getStates();
+  const n = getVertices();
+  n.forEach((vertex) => {
+    const node: Node = {
+      id: vertex.id,
+      label: vertex.text,
+      labelStyle: '',
+      padding: config.flowchart?.padding || 8,
+      cssStyles: vertex.styles.join(' '),
+      cssClasses: vertex.classes.join(' '),
+      shape: getTypeFromVertex(vertex),
+      dir: vertex.dir,
+      domId: vertex.domId,
+      type: undefined,
+      isGroup: false,
+    };
+    nodes.push(node);
+  });
+
+  const useRough = config.look === 'handdrawn';
+
+  return { nodes, edges, other: {}, config };
+};
+
 export default {
   defaultConfig: () => defaultConfig.flowchart,
   setAccTitle,
   getAccTitle,
   getAccDescription,
+  getData,
   setAccDescription,
   addVertex,
   lookUpDomId,
diff --git a/packages/mermaid/src/diagrams/flowchart/flowDiagram-v2.ts b/packages/mermaid/src/diagrams/flowchart/flowDiagram-v2.ts
index 368a98cca..e3db21540 100644
--- a/packages/mermaid/src/diagrams/flowchart/flowDiagram-v2.ts
+++ b/packages/mermaid/src/diagrams/flowchart/flowDiagram-v2.ts
@@ -1,7 +1,8 @@
 // @ts-ignore: JISON doesn't support types
 import flowParser from './parser/flow.jison';
 import flowDb from './flowDb.js';
-import flowRendererV2 from './flowRenderer-v2.js';
+// import flowRendererV2 from './flowRenderer-v2.js';
+import flowRendererV3 from './flowRenderer-v3-unified.js';
 import flowStyles from './styles.js';
 import type { MermaidConfig } from '../../config.type.js';
 import { setConfig } from '../../diagram-api/diagramAPI.js';
@@ -9,7 +10,8 @@ import { setConfig } from '../../diagram-api/diagramAPI.js';
 export const diagram = {
   parser: flowParser,
   db: flowDb,
-  renderer: flowRendererV2,
+  // renderer: flowRendererV2,
+  renderer: flowRendererV3,
   styles: flowStyles,
   init: (cnf: MermaidConfig) => {
     if (!cnf.flowchart) {
@@ -18,7 +20,7 @@ export const diagram = {
     cnf.flowchart.arrowMarkerAbsolute = cnf.arrowMarkerAbsolute;
     // flowchart-v2 uses dagre-wrapper, which doesn't have access to flowchart cnf
     setConfig({ flowchart: { arrowMarkerAbsolute: cnf.arrowMarkerAbsolute } });
-    flowRendererV2.setConf(cnf.flowchart);
+    flowRendererV3.setConf(cnf.flowchart);
     flowDb.clear();
     flowDb.setGen('gen-2');
   },
diff --git a/packages/mermaid/src/diagrams/flowchart/flowDiagram-v3-unified.ts b/packages/mermaid/src/diagrams/flowchart/flowDiagram-v3-unified.ts
new file mode 100644
index 000000000..368a98cca
--- /dev/null
+++ b/packages/mermaid/src/diagrams/flowchart/flowDiagram-v3-unified.ts
@@ -0,0 +1,25 @@
+// @ts-ignore: JISON doesn't support types
+import flowParser from './parser/flow.jison';
+import flowDb from './flowDb.js';
+import flowRendererV2 from './flowRenderer-v2.js';
+import flowStyles from './styles.js';
+import type { MermaidConfig } from '../../config.type.js';
+import { setConfig } from '../../diagram-api/diagramAPI.js';
+
+export const diagram = {
+  parser: flowParser,
+  db: flowDb,
+  renderer: flowRendererV2,
+  styles: flowStyles,
+  init: (cnf: MermaidConfig) => {
+    if (!cnf.flowchart) {
+      cnf.flowchart = {};
+    }
+    cnf.flowchart.arrowMarkerAbsolute = cnf.arrowMarkerAbsolute;
+    // flowchart-v2 uses dagre-wrapper, which doesn't have access to flowchart cnf
+    setConfig({ flowchart: { arrowMarkerAbsolute: cnf.arrowMarkerAbsolute } });
+    flowRendererV2.setConf(cnf.flowchart);
+    flowDb.clear();
+    flowDb.setGen('gen-2');
+  },
+};
diff --git a/packages/mermaid/src/diagrams/flowchart/flowRenderer-v3-unified.ts b/packages/mermaid/src/diagrams/flowchart/flowRenderer-v3-unified.ts
new file mode 100644
index 000000000..bf2d92176
--- /dev/null
+++ b/packages/mermaid/src/diagrams/flowchart/flowRenderer-v3-unified.ts
@@ -0,0 +1,81 @@
+import { log } from '../../logger.js';
+import type { DiagramStyleClassDef } from '../../diagram-api/types.js';
+import type { LayoutData, LayoutMethod } from '../../rendering-util/types.js';
+import { getConfig } from '../../diagram-api/diagramAPI.js';
+import { render } from '../../rendering-util/render.js';
+import { getDiagramElements } from '../../rendering-util/inserElementsForSize.js';
+import { setupViewPortForSVG } from '../../rendering-util/setupViewPortForSVG.js';
+import { getDirection } from './flowDb.js';
+
+import utils from '../../utils.js';
+
+// Configuration
+const conf: Record<string, any> = {};
+
+export const setConf = function (cnf: Record<string, any>) {
+  const keys = Object.keys(cnf);
+  for (const key of keys) {
+    conf[key] = cnf[key];
+  }
+};
+
+export const getClasses = function (
+  text: string,
+  diagramObj: any
+): Record<string, DiagramStyleClassDef> {
+  // diagramObj.db.extract(diagramObj.db.getRootDocV2());
+  return diagramObj.db.getClasses();
+};
+
+export const draw = async function (text: string, id: string, _version: string, diag: any) {
+  log.info('REF0:');
+  log.info('Drawing state diagram (v2)', id);
+  const { securityLevel, state: conf, layout } = getConfig();
+
+  const DIR = getDirection();
+
+  // The getData method provided in all supported diagrams is used to extract the data from the parsed structure
+  // into the Layout data format
+  console.log('Before getData: ');
+  const data4Layout = diag.db.getData() as LayoutData;
+  console.log('Data: ', data4Layout);
+  // Create the root SVG - the element is the div containing the SVG element
+  const { element, svg } = getDiagramElements(id, securityLevel);
+
+  // // For some diagrams this call is not needed, but in the state diagram it is
+  // await insertElementsForSize(element, data4Layout);
+
+  // console.log('data4Layout:', data4Layout);
+
+  // // Now we have layout data with real sizes, we can perform the layout
+  // const data4Rendering = doLayout(data4Layout, id, _version, 'dagre-wrapper');
+
+  // // The performRender method provided in all supported diagrams is used to render the data
+  // performRender(data4Rendering);
+
+  data4Layout.type = diag.type;
+  // data4Layout.layoutAlgorithm = 'dagre-wrapper';
+  // data4Layout.layoutAlgorithm = 'elk';
+  data4Layout.layoutAlgorithm = layout;
+  data4Layout.direction = DIR;
+  data4Layout.nodeSpacing = conf?.nodeSpacing || 50;
+  data4Layout.rankSpacing = conf?.rankSpacing || 50;
+  data4Layout.markers = ['barb'];
+  data4Layout.diagramId = id;
+  console.log('REF1:', data4Layout);
+  await render(data4Layout, svg, element);
+  const padding = 8;
+  utils.insertTitle(
+    element,
+    'statediagramTitleText',
+    conf?.titleTopMargin || 0,
+    diag.db.getDiagramTitle()
+  );
+  setupViewPortForSVG(svg, padding, 'flowchart', conf?.useMaxWidth || false);
+};
+
+export default {
+  setConf,
+  getClasses,
+  draw,
+};
diff --git a/packages/mermaid/src/diagrams/state/stateDb.js b/packages/mermaid/src/diagrams/state/stateDb.js
index cda36f2ad..7193820c0 100644
--- a/packages/mermaid/src/diagrams/state/stateDb.js
+++ b/packages/mermaid/src/diagrams/state/stateDb.js
@@ -585,6 +585,7 @@ export const getData = () => {
   const useRough = config.look === 'handdrawn';
   dataFetcher(undefined, getRootDocV2(), diagramStates, nodes, edges, true, useRough);
 
+  console.log('State Nodes XDX:', nodes);
   return { nodes, edges, other: {}, config };
 };
 
diff --git a/packages/mermaid/src/rendering-util/rendering-elements/nodes.js b/packages/mermaid/src/rendering-util/rendering-elements/nodes.js
index 7216749c8..52219ac08 100644
--- a/packages/mermaid/src/rendering-util/rendering-elements/nodes.js
+++ b/packages/mermaid/src/rendering-util/rendering-elements/nodes.js
@@ -1,5 +1,7 @@
 import { log } from '$root/logger.js';
-import { rect } from './shapes/rect.ts';
+import { state } from './shapes/state.ts';
+import { roundedRect } from './shapes/roundedRect.ts';
+import { squareRect } from './shapes/squareRect.ts';
 import { stateStart } from './shapes/stateStart.ts';
 import { stateEnd } from './shapes/stateEnd.ts';
 import { forkJoin } from './shapes/forkJoin.ts';
@@ -15,13 +17,15 @@ const formatClass = (str) => {
 };
 
 const shapes = {
-  rect,
+  state,
   stateStart,
   stateEnd,
   fork: forkJoin,
   join: forkJoin,
   choice,
   note,
+  roundedRect,
+  squareRect,
 };
 
 let nodeElems = {};
@@ -30,7 +34,6 @@ export const insertNode = async (elem, node, dir) => {
   let newEl;
   let el;
 
-  // debugger;
   // Add link when appropriate
   if (node.link) {
     let target;
diff --git a/packages/mermaid/src/rendering-util/rendering-elements/shapes/rect.ts b/packages/mermaid/src/rendering-util/rendering-elements/shapes/drawRect.ts
similarity index 94%
rename from packages/mermaid/src/rendering-util/rendering-elements/shapes/rect.ts
rename to packages/mermaid/src/rendering-util/rendering-elements/shapes/drawRect.ts
index 9ced7ee1f..675ccee09 100644
--- a/packages/mermaid/src/rendering-util/rendering-elements/shapes/rect.ts
+++ b/packages/mermaid/src/rendering-util/rendering-elements/shapes/drawRect.ts
@@ -1,7 +1,7 @@
 import { log } from '$root/logger.js';
 import { labelHelper, updateNodeBounds } from './util.js';
 import intersect from '../intersect/index.js';
-import type { Node } from '$root/rendering-util/types.d.ts';
+import type { Node, RectOptions } from '$root/rendering-util/types.d.ts';
 import { createRoundedRectPathD } from './roundedRectPath.js';
 import { getConfig } from '$root/diagram-api/diagramAPI.js';
 import { userNodeOverrides } from '$root/rendering-util/rendering-elements/shapes/handdrawnStyles.js';
@@ -58,7 +58,7 @@ function applyNodePropertyBorders(
   rect.attr('stroke-dasharray', strokeDashArray.join(' '));
 }
 
-export const rect = async (parent: SVGAElement, node: Node) => {
+export const drawRect = async (parent: SVGAElement, node: Node, options: RectOptions) => {
   const { themeVariables, handdrawnSeed } = getConfig();
   const { nodeBorder, mainBkg } = themeVariables;
 
@@ -75,7 +75,7 @@ export const rect = async (parent: SVGAElement, node: Node) => {
   const y = -bbox.height / 2 - halfPadding;
 
   let rect;
-  const { rx, ry, style: cssStyles, useRough } = node;
+  const { rx, ry, cssStyles, useRough } = node;
   if (useRough) {
     const rc = rough.svg(shapeSvg);
     const options = userNodeOverrides(node, {
@@ -101,6 +101,8 @@ export const rect = async (parent: SVGAElement, node: Node) => {
       .attr('class', 'basic label-container')
       .attr('style', cssStyles)
       .attr('rx', rx)
+      .attr('data-id', 'abc')
+      .attr('data-et', 'node')
       .attr('ry', ry)
       .attr('x', x)
       .attr('y', y)
diff --git a/packages/mermaid/src/rendering-util/rendering-elements/shapes/roundedRect.ts b/packages/mermaid/src/rendering-util/rendering-elements/shapes/roundedRect.ts
new file mode 100644
index 000000000..aa254c000
--- /dev/null
+++ b/packages/mermaid/src/rendering-util/rendering-elements/shapes/roundedRect.ts
@@ -0,0 +1,13 @@
+import type { Node, RectOptions } from '$root/rendering-util/types.d.ts';
+import { drawRect } from './drawRect.js';
+
+export const roundedRect = async (parent: SVGAElement, node: Node) => {
+  const options = {
+    rx: 5,
+    ry: 5,
+    classes: '',
+  } as RectOptions;
+
+  console.log('roundedRect XDX');
+  return drawRect(parent, node, options);
+};
diff --git a/packages/mermaid/src/rendering-util/rendering-elements/shapes/squareRect.ts b/packages/mermaid/src/rendering-util/rendering-elements/shapes/squareRect.ts
new file mode 100644
index 000000000..8946b141c
--- /dev/null
+++ b/packages/mermaid/src/rendering-util/rendering-elements/shapes/squareRect.ts
@@ -0,0 +1,11 @@
+import type { Node, RectOptions } from '$root/rendering-util/types.d.ts';
+import { drawRect } from './drawRect.js';
+
+export const squareRect = async (parent: SVGAElement, node: Node) => {
+  const options = {
+    rx: 0,
+    ry: 0,
+    classes: '',
+  } as RectOptions;
+  return drawRect(parent, node, options);
+};
diff --git a/packages/mermaid/src/rendering-util/rendering-elements/shapes/state.ts b/packages/mermaid/src/rendering-util/rendering-elements/shapes/state.ts
new file mode 100644
index 000000000..73d119045
--- /dev/null
+++ b/packages/mermaid/src/rendering-util/rendering-elements/shapes/state.ts
@@ -0,0 +1,11 @@
+import type { Node } from '$root/rendering-util/types.d.ts';
+import { drawRect } from './drawRect.js';
+
+export const state = async (parent: SVGAElement, node: Node) => {
+  const options = {
+    rx: 5,
+    ry: 5,
+    classes: 'flowchart-node',
+  };
+  return drawRect(parent, node, options);
+};
diff --git a/packages/mermaid/src/rendering-util/types.d.ts b/packages/mermaid/src/rendering-util/types.d.ts
index 1276c41d8..997d93d62 100644
--- a/packages/mermaid/src/rendering-util/types.d.ts
+++ b/packages/mermaid/src/rendering-util/types.d.ts
@@ -1,5 +1,5 @@
 import config from '../../dist/defaultConfig';
-import { MermaidConfig } from '../../dist/config.type';
+import type { MermaidConfig } from '../../dist/config.type';
 export type MarkdownWordType = 'normal' | 'strong' | 'emphasis';
 export interface MarkdownWord {
   content: string;
@@ -90,6 +90,12 @@ interface Edge {
   useRough?: boolean;
 }
 
+interface RectOptions {
+  rx: number;
+  ry: number;
+  classes: string;
+}
+
 // Extending the Node interface for specific types if needed
 interface ClassDiagramNode extends Node {
   memberData: any; // Specific property for class diagram nodes