From 41d7a549b0d27122f37ec6350f5bb79a96bc26e3 Mon Sep 17 00:00:00 2001
From: yari-dewalt <yari.dewalt@gmail.com>
Date: Thu, 23 Jan 2025 09:36:49 -0800
Subject: [PATCH] Add support for classes

---
 .../parser/requirementDiagram.jison           | 39 +++++++--
 .../src/diagrams/requirement/requirementDb.ts | 86 ++++++++++++++++---
 .../mermaid/src/diagrams/requirement/types.ts |  8 ++
 .../shapes/requirementBox.ts                  |  5 +-
 4 files changed, 119 insertions(+), 19 deletions(-)

diff --git a/packages/mermaid/src/diagrams/requirement/parser/requirementDiagram.jison b/packages/mermaid/src/diagrams/requirement/parser/requirementDiagram.jison
index 16181ba98..16424f3b1 100644
--- a/packages/mermaid/src/diagrams/requirement/parser/requirementDiagram.jison
+++ b/packages/mermaid/src/diagrams/requirement/parser/requirementDiagram.jison
@@ -37,6 +37,7 @@ accDescr\s*"{"\s*                                { this.begin("acc_descr_multili
 
 "{"                         return 'STRUCT_START';
 "}"                         return 'STRUCT_STOP';
+":"{3}                       return 'STYLE_SEPARATOR';
 ":"                         return 'COLONSEP';
 
 "id"                        return 'ID';
@@ -77,13 +78,16 @@ accDescr\s*"{"\s*                                { this.begin("acc_descr_multili
 <style>\w+                          return 'ALPHA';
 <style>":" return 'COLON';
 <style>";" return 'SEMICOLON';
-<style>"," return 'COMMA';
 <style>"%" return 'PERCENT';
 <style>"-" return 'MINUS';
 <style>"#" return 'BRKT';
 <style>" "                             /* skip spaces */
+<style>["] { this.begin("string"); }
 <style>\n { this.popState(); }
 
+"classDef" { this.begin("style"); return 'CLASSDEF'; }
+"class" { this.begin("style"); return 'CLASS'; }
+
 "<-"        return 'END_ARROW_L';
 "->"        {return 'END_ARROW_R';}
 "-"         {return 'LINE';}
@@ -92,10 +96,11 @@ accDescr\s*"{"\s*                                { this.begin("acc_descr_multili
 <string>["]         { this.popState(); }
 <string>[^"]*       { return "qString"; }
 
-[\w][^\r\n\{\<\>\-\=]*                { yytext = yytext.trim(); return 'unqString';}
+[\w][^:,\r\n\{\<\>\-\=]*                { yytext = yytext.trim(); return 'unqString';}
 
 <*>\w+                          return 'ALPHA';
 <*>[0-9]+                       return 'NUM';
+<*>","                             return 'COMMA';
 
 /lex
 
@@ -122,6 +127,8 @@ diagram
   | directive diagram
   | direction diagram
   | styleStatement diagram
+  | classDefStatement diagram
+  | classStatement diagram
   | NEWLINE diagram
   ;
 
@@ -137,8 +144,9 @@ direction
     ;
 
 requirementDef
-  : requirementType requirementName STRUCT_START NEWLINE requirementBody
-    { yy.addRequirement($2, $1) };
+  : requirementType requirementName STRUCT_START NEWLINE requirementBody { yy.addRequirement($2, $1) }
+  | requirementType requirementName STYLE_SEPARATOR idList STRUCT_START NEWLINE requirementBody { yy.addRequirement($2, $1); yy.setClass([$2], $4); }
+  ;
 
 requirementBody
   : ID COLONSEP id NEWLINE requirementBody
@@ -182,8 +190,9 @@ verifyType
     { $$=yy.VerifyType.VERIFY_TEST;};
 
 elementDef
-  : ELEMENT elementName STRUCT_START NEWLINE elementBody
-    { yy.addElement($2) };
+  : ELEMENT elementName STRUCT_START NEWLINE elementBody { yy.addElement($2) }
+  | ELEMENT elementName STYLE_SEPARATOR idList STRUCT_START NEWLINE elementBody { yy.addElement($2); yy.setClass([$2], $4); }
+  ;
 
 elementBody
   : TYPE COLONSEP type NEWLINE elementBody
@@ -215,8 +224,24 @@ relationship
   | TRACES
       { $$=yy.Relationships.TRACES;};
 
+classDefStatement
+  : CLASSDEF idList stylesOpt {$$ = $CLASSDEF;yy.defineClass($idList,$stylesOpt);}
+  ;
+
+classStatement
+    : CLASS idList idList                            {yy.setClass($2, $3);}
+    | id STYLE_SEPARATOR idList {yy.setClass([$1], $3);}
+    ;
+
+idList
+    : ALPHA { $$ = [$ALPHA]; }
+    | idList COMMA ALPHA = { $$ = $idList.concat([$ALPHA]); }
+    | id { $$ = [$id]; }
+    | idList COMMA id = { $$ = $idList.concat([$id]); }
+    ;
+
 styleStatement
-  : STYLE ALPHA stylesOpt                              {$$ = $STYLE;yy.setCssStyle($2,$stylesOpt);}
+  : STYLE idList stylesOpt                              {$$ = $STYLE;yy.setCssStyle($2,$stylesOpt);}
   ;
 
 stylesOpt
diff --git a/packages/mermaid/src/diagrams/requirement/requirementDb.ts b/packages/mermaid/src/diagrams/requirement/requirementDb.ts
index ea12046d8..4fd9d0f17 100644
--- a/packages/mermaid/src/diagrams/requirement/requirementDb.ts
+++ b/packages/mermaid/src/diagrams/requirement/requirementDb.ts
@@ -16,6 +16,7 @@ import type {
   Relation,
   RelationshipType,
   Requirement,
+  RequirementClass,
   RequirementType,
   RiskLevel,
   VerifyType,
@@ -67,6 +68,7 @@ const getInitialRequirement = (): Requirement => ({
   name: '',
   type: RequirementType.REQUIREMENT as RequirementType,
   cssStyles: [],
+  classes: ['default'],
 });
 
 const getInitialElement = (): Element => ({
@@ -74,6 +76,7 @@ const getInitialElement = (): Element => ({
   type: '',
   docRef: '',
   cssStyles: [],
+  classes: ['default'],
 });
 
 // Update initial declarations
@@ -82,6 +85,7 @@ let latestRequirement: Requirement = getInitialRequirement();
 let requirements = new Map<string, Requirement>();
 let latestElement: Element = getInitialElement();
 let elements = new Map<string, Element>();
+let classes = new Map<string, RequirementClass>();
 
 // Add reset functions
 const resetLatestRequirement = () => {
@@ -102,6 +106,7 @@ const addRequirement = (name: string, type: RequirementType) => {
       risk: latestRequirement.risk,
       verifyMethod: latestRequirement.verifyMethod,
       cssStyles: [],
+      classes: ['default'],
     });
   }
   resetLatestRequirement();
@@ -142,6 +147,7 @@ const addElement = (name: string) => {
       type: latestElement.type,
       docRef: latestElement.docRef,
       cssStyles: [],
+      classes: ['default'],
     });
     log.info('Added new element: ', name);
   }
@@ -180,29 +186,84 @@ const clear = () => {
   requirements = new Map();
   resetLatestElement();
   elements = new Map();
+  classes = new Map();
   commonClear();
 };
 
-export const setCssStyle = function (id: string, styles: string[]) {
-  const node = requirements.get(id) ?? elements.get(id);
-  if (!styles || !node) {
-    return;
-  }
-  for (const s of styles) {
-    if (s.includes(',')) {
-      node.cssStyles.push(...s.split(','));
-    } else {
-      node.cssStyles.push(s);
+export const setCssStyle = function (ids: string[], styles: string[]) {
+  for (const id of ids) {
+    const node = requirements.get(id) ?? elements.get(id);
+    if (!styles || !node) {
+      return;
+    }
+    for (const s of styles) {
+      if (s.includes(',')) {
+        node.cssStyles.push(...s.split(','));
+      } else {
+        node.cssStyles.push(s);
+      }
     }
   }
 };
 
+export const setClass = function (ids: string[], classNames: string[]) {
+  for (const id of ids) {
+    const node = requirements.get(id) ?? elements.get(id);
+    if (node) {
+      for (const _class of classNames) {
+        node.classes.push(_class);
+        const styles = classes.get(_class)?.styles;
+        if (styles) {
+          node.cssStyles.push(...styles);
+        }
+      }
+    }
+  }
+};
+
+export const defineClass = function (ids: string[], style: string[]) {
+  for (const id of ids) {
+    let styleClass = classes.get(id);
+    if (styleClass === undefined) {
+      styleClass = { id, styles: [], textStyles: [] };
+      classes.set(id, styleClass);
+    }
+
+    if (style) {
+      style.forEach(function (s) {
+        if (/color/.exec(s)) {
+          const newStyle = s.replace('fill', 'bgFill'); // .replace('color', 'fill');
+          styleClass.textStyles.push(newStyle);
+        }
+        styleClass.styles.push(s);
+      });
+    }
+
+    requirements.forEach((value) => {
+      if (value.classes.includes(id)) {
+        value.cssStyles.push(...style.flatMap((s) => s.split(',')));
+      }
+    });
+    elements.forEach((value) => {
+      if (value.classes.includes(id)) {
+        value.cssStyles.push(...style.flatMap((s) => s.split(',')));
+      }
+    });
+  }
+};
+
+export const getClasses = () => {
+  return classes;
+};
+
 const getData = () => {
   const config = getConfig();
   const nodes: Node[] = [];
   const edges: Edge[] = [];
   for (const requirement of requirements.values()) {
     const node = requirement as unknown as Node;
+    node.cssStyles = requirement.cssStyles;
+    node.cssClasses = requirement.classes.join(' ');
     node.shape = 'requirementBox';
     node.look = config.look;
     nodes.push(node);
@@ -213,6 +274,8 @@ const getData = () => {
     node.shape = 'requirementBox';
     node.look = config.look;
     node.id = element.name;
+    node.cssStyles = element.cssStyles;
+    node.cssClasses = element.classes.join(' ');
 
     nodes.push(node);
   }
@@ -270,5 +333,8 @@ export default {
   getRelationships,
   clear,
   setCssStyle,
+  setClass,
+  defineClass,
+  getClasses,
   getData,
 };
diff --git a/packages/mermaid/src/diagrams/requirement/types.ts b/packages/mermaid/src/diagrams/requirement/types.ts
index 429fe9c91..ae71c65a3 100644
--- a/packages/mermaid/src/diagrams/requirement/types.ts
+++ b/packages/mermaid/src/diagrams/requirement/types.ts
@@ -18,6 +18,7 @@ export interface Requirement {
   risk: RiskLevel;
   verifyMethod: VerifyType;
   cssStyles: string[];
+  classes: string[];
 }
 
 export type RelationshipType =
@@ -40,4 +41,11 @@ export interface Element {
   type: string;
   docRef: string;
   cssStyles: string[];
+  classes: string[];
+}
+
+export interface RequirementClass {
+  id: string;
+  styles: string[];
+  textStyles: string[];
 }
diff --git a/packages/mermaid/src/rendering-util/rendering-elements/shapes/requirementBox.ts b/packages/mermaid/src/rendering-util/rendering-elements/shapes/requirementBox.ts
index c79209bad..3ce7b8ef1 100644
--- a/packages/mermaid/src/rendering-util/rendering-elements/shapes/requirementBox.ts
+++ b/packages/mermaid/src/rendering-util/rendering-elements/shapes/requirementBox.ts
@@ -1,4 +1,4 @@
-import { updateNodeBounds } from './util.js';
+import { getNodeClasses, updateNodeBounds } from './util.js';
 import intersect from '../intersect/index.js';
 import type { Node } from '../../types.js';
 import { styles2String, userNodeOverrides } from './handDrawnShapeStyles.js';
@@ -22,11 +22,12 @@ export async function requirementBox<T extends SVGGraphicsElement>(
   const PADDING = 20;
   const GAP = 20;
   const isRequirementNode = 'verifyMethod' in node;
+  const classes = getNodeClasses(node);
 
   // Add outer g element
   const shapeSvg = parent
     .insert('g')
-    .attr('class', '')
+    .attr('class', classes)
     .attr('id', node.domId ?? node.id);
 
   let typeHeight;