feat: enhance use case diagram support with arrow types, class definitions and styles

2025-10-16 20:49:45 +02:00 · 2025-10-13 18:46:47 +05:30
parent b7ff1920a9
commit 5b2b3b8ae9
6 changed files with 527 additions and 27 deletions
--- a/packages/mermaid/src/diagrams/usecase/usecaseDb.ts
+++ b/packages/mermaid/src/diagrams/usecase/usecaseDb.ts
@@ -17,8 +17,9 @@ import type {
  Relationship,
  ActorMetadata,
  Direction,
+  ClassDef,
 } from './usecaseTypes.js';
-import { DEFAULT_DIRECTION } from './usecaseTypes.js';
+import { DEFAULT_DIRECTION, ARROW_TYPE } from './usecaseTypes.js';
 import type { RequiredDeep } from 'type-fest';
 import type { UsecaseDiagramConfig } from '../../config.type.js';
 import DEFAULT_CONFIG from '../../defaultConfig.js';
@@ -32,6 +33,7 @@ export const DEFAULT_USECASE_DB: RequiredDeep<UsecaseFields> = {
  useCases: new Map(),
  systemBoundaries: new Map(),
  relationships: [],
+  classDefs: new Map(),
  direction: DEFAULT_DIRECTION,
  config: DEFAULT_USECASE_CONFIG,
 } as const;
@@ -40,6 +42,7 @@ let actors = new Map<string, Actor>();
 let useCases = new Map<string, UseCase>();
 let systemBoundaries = new Map<string, SystemBoundary>();
 let relationships: Relationship[] = [];
+let classDefs = new Map<string, ClassDef>();
 let direction: Direction = DEFAULT_DIRECTION;
 const config: Required<UsecaseDiagramConfig> = structuredClone(DEFAULT_USECASE_CONFIG);

@@ -50,6 +53,7 @@ const clear = (): void => {
  useCases = new Map();
  systemBoundaries = new Map();
  relationships = [];
+  classDefs = new Map();
  direction = DEFAULT_DIRECTION;
  commonClear();
 };
@@ -139,7 +143,7 @@ const addRelationship = (relationship: Relationship): void => {

  // Validate arrow type if present
  if (relationship.arrowType !== undefined) {
-    const validArrowTypes = [0, 1, 2]; // SOLID_ARROW, BACK_ARROW, LINE_SOLID
+    const validArrowTypes = [0, 1, 2, 3, 4, 5, 6]; // SOLID_ARROW, BACK_ARROW, LINE_SOLID, CIRCLE_ARROW, CROSS_ARROW
    if (!validArrowTypes.includes(relationship.arrowType)) {
      throw new Error(
        `Invalid arrow type: ${relationship.arrowType}. Valid arrow types are: ${validArrowTypes.join(', ')}`
@@ -155,6 +159,37 @@ const addRelationship = (relationship: Relationship): void => {

 const getRelationships = (): Relationship[] => relationships;

+// ClassDef management
+const addClassDef = (classDef: ClassDef): void => {
+  if (!classDef.id) {
+    throw new Error(
+      `Invalid classDef: ClassDef must have an id. Received: ${JSON.stringify(classDef)}`
+    );
+  }
+
+  classDefs.set(classDef.id, classDef);
+  log.debug(`Added classDef: ${classDef.id}`);
+};
+
+const getClassDefs = (): Map<string, ClassDef> => classDefs;
+
+const getClassDef = (id: string): ClassDef | undefined => classDefs.get(id);
+
+/**
+ * Get compiled styles from class definitions
+ * Similar to flowchart's getCompiledStyles method
+ */
+const getCompiledStyles = (classNames: string[]): string[] => {
+  let compiledStyles: string[] = [];
+  for (const className of classNames) {
+    const cssClass = classDefs.get(className);
+    if (cssClass?.styles) {
+      compiledStyles = [...compiledStyles, ...(cssClass.styles ?? [])].map((s) => s.trim());
+    }
+  }
+  return compiledStyles;
+};
+
 // Direction management
 const setDirection = (dir: Direction): void => {
  // Normalize TD to TB (same as flowchart)
@@ -176,11 +211,17 @@ const getData = (): LayoutData => {

  // Convert actors to nodes
  for (const actor of actors.values()) {
+    const classesArray = ['default', 'usecase-actor'];
+    const cssCompiledStyles = getCompiledStyles(classesArray);
+
+    // Determine which shape to use based on whether actor has an icon
+    const actorShape = actor.metadata?.icon ? 'usecaseActorIcon' : 'usecaseActor';
+
    const node: Node = {
      id: actor.id,
      label: actor.name,
      description: actor.description ? [actor.description] : undefined,
-      shape: 'usecaseActor', // Use custom actor shape
+      shape: actorShape, // Use icon shape if icon is present, otherwise stick figure
      isGroup: false,
      padding: 10,
      look: globalConfig.look,
@@ -192,6 +233,8 @@ const getData = (): LayoutData => {
              .join(' ')
          : ''
      }`.trim(),
+      cssStyles: actor.styles ?? [], // Direct styles
+      cssCompiledStyles, // Compiled styles from class definitions
      // Pass actor metadata to the shape handler
      metadata: actor.metadata,
    } as Node & { metadata?: ActorMetadata };
@@ -200,6 +243,16 @@ const getData = (): LayoutData => {

  // Convert use cases to nodes
  for (const useCase of useCases.values()) {
+    // Build CSS classes string
+    let cssClasses = 'usecase-element';
+    const classesArray = ['default', 'usecase-element'];
+    if (useCase.classes && useCase.classes.length > 0) {
+      cssClasses += ' ' + useCase.classes.join(' ');
+      classesArray.push(...useCase.classes);
+    }
+
+    // Get compiled styles from class definitions
+    const cssCompiledStyles = getCompiledStyles(classesArray);
    const node: Node = {
      id: useCase.id,
      label: useCase.name,
@@ -208,7 +261,9 @@ const getData = (): LayoutData => {
      isGroup: false,
      padding: 10,
      look: globalConfig.look,
-      cssClasses: 'usecase-element',
+      cssClasses,
+      cssStyles: useCase.styles ?? [], // Direct styles
+      cssCompiledStyles, // Compiled styles from class definitions
      // If use case belongs to a system boundary, set parentId
      ...(useCase.systemBoundary && { parentId: useCase.systemBoundary }),
    };
@@ -217,6 +272,13 @@ const getData = (): LayoutData => {

  // Convert system boundaries to group nodes
  for (const boundary of systemBoundaries.values()) {
+    const classesArray = [
+      'default',
+      'system-boundary',
+      `system-boundary-${boundary.type ?? 'rect'}`,
+    ];
+    const cssCompiledStyles = getCompiledStyles(classesArray);
+
    const node: ClusterNode & { boundaryType?: string } = {
      id: boundary.id,
      label: boundary.name,
@@ -225,6 +287,8 @@ const getData = (): LayoutData => {
      padding: 20,
      look: globalConfig.look,
      cssClasses: `system-boundary system-boundary-${boundary.type ?? 'rect'}`,
+      cssStyles: boundary.styles ?? [], // Direct styles
+      cssCompiledStyles, // Compiled styles from class definitions
      // Pass boundary type to the shape handler
      boundaryType: boundary.type,
    };
@@ -233,6 +297,34 @@ const getData = (): LayoutData => {

  // Convert relationships to edges
  relationships.forEach((relationship, index) => {
+    // Determine arrow types based on relationship.arrowType
+    let arrowTypeEnd = 'none';
+    let arrowTypeStart = 'none';
+
+    switch (relationship.arrowType) {
+      case ARROW_TYPE.SOLID_ARROW: // -->
+        arrowTypeEnd = 'arrow_point';
+        break;
+      case ARROW_TYPE.BACK_ARROW: // <--
+        arrowTypeStart = 'arrow_point';
+        break;
+      case ARROW_TYPE.CIRCLE_ARROW: // --o
+        arrowTypeEnd = 'arrow_circle';
+        break;
+      case ARROW_TYPE.CROSS_ARROW: // --x
+        arrowTypeEnd = 'arrow_cross';
+        break;
+      case ARROW_TYPE.CIRCLE_ARROW_REVERSED: // o--
+        arrowTypeStart = 'arrow_circle';
+        break;
+      case ARROW_TYPE.CROSS_ARROW_REVERSED: // x--
+        arrowTypeStart = 'arrow_cross';
+        break;
+      case ARROW_TYPE.LINE_SOLID: // --
+        // Both remain 'none'
+        break;
+    }
+
    const edge: Edge = {
      id: relationship.id || `edge-${index}`,
      start: relationship.from,
@@ -240,15 +332,10 @@ const getData = (): LayoutData => {
      source: relationship.from,
      target: relationship.to,
      label: relationship.label,
+      labelpos: 'c', // Center label position for proper dagre layout
      type: relationship.type,
-      arrowTypeEnd:
-        relationship.arrowType === 0
-          ? 'arrow_point' // Forward arrow (-->)
-          : 'none', // No end arrow for back arrow or line
-      arrowTypeStart:
-        relationship.arrowType === 1
-          ? 'arrow_point' // Back arrow (<--)
-          : 'none', // No start arrow for forward arrow or line
+      arrowTypeEnd,
+      arrowTypeStart,
      classes: `relationship relationship-${relationship.type}`,
      look: globalConfig.look,
      thickness: 'normal',
@@ -297,6 +384,10 @@ export const db: UsecaseDB = {
  addRelationship,
  getRelationships,

+  addClassDef,
+  getClassDefs,
+  getClassDef,
+
  // Direction management
  setDirection,
  getDirection,
--- a/packages/mermaid/src/diagrams/usecase/usecaseParser.ts
+++ b/packages/mermaid/src/diagrams/usecase/usecaseParser.ts
@@ -10,14 +10,28 @@ import type {
  SystemBoundary,
  Relationship,
  ArrowType,
+  ClassDef,
 } from './usecaseTypes.js';
 import { db } from './usecaseDb.js';

 // ANTLR parser result interface
 interface UsecaseParseResult {
-  actors: { id: string; name: string; metadata?: Record<string, string> }[];
-  useCases: { id: string; name: string; nodeId?: string; systemBoundary?: string }[];
-  systemBoundaries: { id: string; name: string; useCases: string[]; type?: 'package' | 'rect' }[];
+  actors: { id: string; name: string; metadata?: Record<string, string>; styles?: string[] }[];
+  useCases: {
+    id: string;
+    name: string;
+    nodeId?: string;
+    systemBoundary?: string;
+    classes?: string[];
+    styles?: string[];
+  }[];
+  systemBoundaries: {
+    id: string;
+    name: string;
+    useCases: string[];
+    type?: 'package' | 'rect';
+    styles?: string[];
+  }[];
  relationships: {
    id: string;
    from: string;
@@ -26,6 +40,7 @@ interface UsecaseParseResult {
    arrowType: number;
    label?: string;
  }[];
+  classDefs?: Map<string, { id: string; styles: string[] }>;
  direction?: string;
  accDescr?: string;
  accTitle?: string;
@@ -54,17 +69,20 @@ const populateDb = (ast: UsecaseParseResult, db: UsecaseDB) => {
      id: actorData.id,
      name: actorData.name,
      metadata: actorData.metadata,
+      styles: actorData.styles,
    };
    db.addActor(actor);
  });

-  // Add use cases (ANTLR result already has id, name, nodeId, and systemBoundary)
+  // Add use cases (ANTLR result already has id, name, nodeId, systemBoundary, and classes)
  ast.useCases.forEach((useCaseData) => {
    const useCase: UseCase = {
      id: useCaseData.id,
      name: useCaseData.name,
      nodeId: useCaseData.nodeId,
      systemBoundary: useCaseData.systemBoundary,
+      classes: useCaseData.classes,
+      styles: useCaseData.styles,
    };
    db.addUseCase(useCase);
  });
@@ -77,6 +95,7 @@ const populateDb = (ast: UsecaseParseResult, db: UsecaseDB) => {
        name: boundaryData.name,
        useCases: boundaryData.useCases,
        type: boundaryData.type || 'rect', // default to 'rect' if not specified
+        styles: boundaryData.styles,
      };
      db.addSystemBoundary(systemBoundary);
    });
@@ -95,6 +114,17 @@ const populateDb = (ast: UsecaseParseResult, db: UsecaseDB) => {
    db.addRelationship(relationship);
  });

+  // Add class definitions
+  if (ast.classDefs) {
+    ast.classDefs.forEach((classDefData) => {
+      const classDef: ClassDef = {
+        id: classDefData.id,
+        styles: classDefData.styles,
+      };
+      db.addClassDef(classDef);
+    });
+  }
+
  // Set direction if provided
  if (ast.direction) {
    db.setDirection(ast.direction as any);
@@ -104,6 +134,7 @@ const populateDb = (ast: UsecaseParseResult, db: UsecaseDB) => {
    actors: ast.actors.length,
    useCases: ast.useCases.length,
    relationships: ast.relationships.length,
+    classDefs: ast.classDefs?.size ?? 0,
    direction: ast.direction,
  });
 };
--- a/packages/mermaid/src/diagrams/usecase/usecaseRenderer.ts
+++ b/packages/mermaid/src/diagrams/usecase/usecaseRenderer.ts
@@ -27,7 +27,7 @@ const draw: DrawDefinition = async (_text, id, _version, diag) => {

  data4Layout.nodeSpacing = 50; // Default node spacing
  data4Layout.rankSpacing = 50; // Default rank spacing
-  data4Layout.markers = ['point']; // Use point markers for usecase diagrams
+  data4Layout.markers = ['point', 'circle', 'cross']; // Support point, circle, and cross markers
  data4Layout.diagramId = id;

  log.debug('Usecase layout data:', data4Layout);
--- a/packages/mermaid/src/diagrams/usecase/usecaseTypes.ts
+++ b/packages/mermaid/src/diagrams/usecase/usecaseTypes.ts
@@ -9,6 +9,7 @@ export interface Actor {
  name: string;
  description?: string;
  metadata?: ActorMetadata;
+  styles?: string[]; // Direct CSS styles applied to this actor
 }

 export interface UseCase {
@@ -17,6 +18,8 @@ export interface UseCase {
  description?: string;
  nodeId?: string; // Optional node ID (e.g., 'a' in 'a(Go through code)')
  systemBoundary?: string; // Optional reference to system boundary
+  classes?: string[]; // CSS classes applied to this use case
+  styles?: string[]; // Direct CSS styles applied to this use case
 }

 export type SystemBoundaryType = 'package' | 'rect';
@@ -26,6 +29,7 @@ export interface SystemBoundary {
  name: string;
  useCases: string[]; // Array of use case IDs within this boundary
  type?: SystemBoundaryType; // Type of boundary rendering (default: 'rect')
+  styles?: string[]; // Direct CSS styles applied to this system boundary
 }

 // Arrow types for usecase diagrams (matching parser types)
@@ -33,6 +37,10 @@ export const ARROW_TYPE = {
  SOLID_ARROW: 0, // -->
  BACK_ARROW: 1, // <--
  LINE_SOLID: 2, // --
+  CIRCLE_ARROW: 3, // --o
+  CROSS_ARROW: 4, // --x
+  CIRCLE_ARROW_REVERSED: 5, // o--
+  CROSS_ARROW_REVERSED: 6, // x--
 } as const;

 export type ArrowType = (typeof ARROW_TYPE)[keyof typeof ARROW_TYPE];
@@ -49,13 +57,19 @@ export interface Relationship {
 // Direction types for usecase diagrams
 export type Direction = 'TB' | 'TD' | 'BT' | 'RL' | 'LR';

-export const DEFAULT_DIRECTION: Direction = 'TB';
+export const DEFAULT_DIRECTION: Direction = 'LR';
+
+export interface ClassDef {
+  id: string;
+  styles: string[];
+}

 export interface UsecaseFields {
  actors: Map<string, Actor>;
  useCases: Map<string, UseCase>;
  systemBoundaries: Map<string, SystemBoundary>;
  relationships: Relationship[];
+  classDefs: Map<string, ClassDef>;
  direction: Direction;
  config: Required<UsecaseDiagramConfig>;
 }
@@ -82,6 +96,11 @@ export interface UsecaseDB extends DiagramDB {
  addRelationship: (relationship: Relationship) => void;
  getRelationships: () => Relationship[];

+  // ClassDef management
+  addClassDef: (classDef: ClassDef) => void;
+  getClassDefs: () => Map<string, ClassDef>;
+  getClassDef: (id: string) => ClassDef | undefined;
+
  // Direction management
  setDirection: (direction: Direction) => void;
  getDirection: () => Direction;
--- a/packages/parser/src/language/usecase/Usecase.g4
+++ b/packages/parser/src/language/usecase/Usecase.g4
@@ -11,9 +11,17 @@ statement
    | systemBoundaryStatement
    | systemBoundaryTypeStatement
    | directionStatement
+    | classDefStatement
+    | classStatement
+    | styleStatement
+    | usecaseStatement
    | NEWLINE
    ;

+usecaseStatement
+    : entityName NEWLINE*
+    ;
+
 actorStatement
    : 'actor' actorList NEWLINE*
    ;
@@ -60,10 +68,16 @@ systemBoundaryContent
    ;

 usecaseInBoundary
-    : IDENTIFIER
+    : usecaseWithClass
+    | IDENTIFIER
    | STRING
    ;

+usecaseWithClass
+    : IDENTIFIER CLASS_SEPARATOR IDENTIFIER
+    | STRING CLASS_SEPARATOR IDENTIFIER
+    ;
+
 systemBoundaryTypeStatement
    : systemBoundaryName '@' '{' systemBoundaryTypeContent '}' NEWLINE*
    ;
@@ -82,7 +96,9 @@ systemBoundaryType
    ;

 entityName
-    : IDENTIFIER
+    : IDENTIFIER CLASS_SEPARATOR IDENTIFIER
+    | STRING CLASS_SEPARATOR IDENTIFIER
+    | IDENTIFIER
    | STRING
    | nodeIdWithLabel
    ;
@@ -106,6 +122,10 @@ arrow
    : SOLID_ARROW
    | BACK_ARROW
    | LINE_SOLID
+    | CIRCLE_ARROW
+    | CROSS_ARROW
+    | CIRCLE_ARROW_REVERSED
+    | CROSS_ARROW_REVERSED
    | labeledArrow
    ;

@@ -113,6 +133,10 @@ labeledArrow
    : LINE_SOLID edgeLabel SOLID_ARROW
    | BACK_ARROW edgeLabel LINE_SOLID
    | LINE_SOLID edgeLabel LINE_SOLID
+    | LINE_SOLID edgeLabel CIRCLE_ARROW
+    | LINE_SOLID edgeLabel CROSS_ARROW
+    | CIRCLE_ARROW_REVERSED edgeLabel LINE_SOLID
+    | CROSS_ARROW_REVERSED edgeLabel LINE_SOLID
    ;

 edgeLabel
@@ -132,6 +156,43 @@ direction
    | 'LR'
    ;

+classDefStatement
+    : 'classDef' IDENTIFIER stylesOpt NEWLINE*
+    ;
+
+stylesOpt
+    : style
+    | stylesOpt COMMA style
+    ;
+
+style
+    : styleComponent
+    | style styleComponent
+    ;
+
+styleComponent
+    : IDENTIFIER
+    | NUMBER
+    | HASH_COLOR
+    | COLON
+    | STRING
+    | DASH
+    | DOT
+    | PERCENT
+    ;
+
+classStatement
+    : 'class' nodeList IDENTIFIER NEWLINE*
+    ;
+
+styleStatement
+    : 'style' IDENTIFIER stylesOpt NEWLINE*
+    ;
+
+nodeList
+    : IDENTIFIER (',' IDENTIFIER)*
+    ;
+
 // Lexer rules
 SOLID_ARROW
    : '-->'
@@ -141,6 +202,21 @@ BACK_ARROW
    : '<--'
    ;

+CIRCLE_ARROW
+    : '--o'
+    ;
+CIRCLE_ARROW_REVERSED
+    : 'o--'
+    ;
+
+CROSS_ARROW
+    : '--x'
+    ;
+    
+CROSS_ARROW_REVERSED
+    : 'x--'
+    ;
+
 LINE_SOLID
    : '--'
    ;
@@ -165,6 +241,10 @@ COLON
    : ':'
    ;

+CLASS_SEPARATOR
+    : ':::'
+    ;
+
 IDENTIFIER
    : [a-zA-Z_][a-zA-Z0-9_]*
    ;
@@ -174,6 +254,28 @@ STRING
    | '\'' (~['\r\n])* '\''
    ;

+HASH_COLOR
+    : '#' [a-fA-F0-9]+
+    ;
+
+NUMBER
+    : [0-9]+ ('.' [0-9]+)? ([a-zA-Z]+)?
+    ;
+
+// These tokens are defined last so they have lowest priority
+// This ensures arrow tokens like '-->' are matched before DASH
+DASH
+    : '-'
+    ;
+
+DOT
+    : '.'
+    ;
+
+PERCENT
+    : '%'
+    ;
+
 NEWLINE
    : [\r\n]+
    ;
--- a/packages/parser/src/language/usecase/visitor.ts
+++ b/packages/parser/src/language/usecase/visitor.ts
@@ -19,6 +19,8 @@ import type {
  SystemBoundaryTypePropertyContext,
  SystemBoundaryTypeContext,
  UsecaseInBoundaryContext,
+  UsecaseWithClassContext,
+  UsecaseStatementContext,
  ActorNameContext,
  ActorDeclarationContext,
  NodeIdWithLabelContext,
@@ -32,6 +34,9 @@ import type {
  EdgeLabelContext,
  DirectionStatementContext,
  DirectionContext,
+  ClassDefStatementContext,
+  ClassStatementContext,
+  NodeListContext,
 } from './generated/UsecaseParser.js';
 import { ARROW_TYPE } from './types.js';
 import type {
@@ -41,6 +46,7 @@ import type {
  Relationship,
  UsecaseParseResult,
  ArrowType,
+  ClassDef,
 } from './types.js';

 export class UsecaseAntlrVisitor extends UsecaseVisitor<void> {
@@ -50,6 +56,7 @@ export class UsecaseAntlrVisitor extends UsecaseVisitor<void> {
  private relationships: Relationship[] = [];
  private relationshipCounter = 0;
  private direction = 'TB'; // Default direction
+  private classDefs = new Map<string, ClassDef>();

  constructor() {
    super();
@@ -62,6 +69,10 @@ export class UsecaseAntlrVisitor extends UsecaseVisitor<void> {
    this.visitSystemBoundaryStatement = this.visitSystemBoundaryStatementImpl.bind(this);
    this.visitSystemBoundaryTypeStatement = this.visitSystemBoundaryTypeStatementImpl.bind(this);
    this.visitDirectionStatement = this.visitDirectionStatementImpl.bind(this);
+    this.visitClassDefStatement = this.visitClassDefStatementImpl.bind(this);
+    this.visitClassStatement = this.visitClassStatementImpl.bind(this);
+    this.visitStyleStatement = this.visitStyleStatementImpl.bind(this);
+    this.visitUsecaseStatement = this.visitUsecaseStatementImpl.bind(this);
    this.visitActorName = this.visitActorNameImpl.bind(this);
    this.visitArrow = this.visitArrowImpl.bind(this);
  }
@@ -95,7 +106,7 @@ export class UsecaseAntlrVisitor extends UsecaseVisitor<void> {

  /**
   * Visit statement rule
-   * Grammar: statement : actorStatement | relationshipStatement | systemBoundaryStatement | systemBoundaryTypeStatement | directionStatement | NEWLINE ;
+   * Grammar: statement : actorStatement | relationshipStatement | systemBoundaryStatement | systemBoundaryTypeStatement | directionStatement | classDefStatement | classStatement | usecaseStatement | NEWLINE ;
   */
  private visitStatementImpl(ctx: StatementContext): void {
    if (ctx.actorStatement?.()) {
@@ -108,6 +119,23 @@ export class UsecaseAntlrVisitor extends UsecaseVisitor<void> {
      this.visitSystemBoundaryTypeStatementImpl(ctx.systemBoundaryTypeStatement()!);
    } else if (ctx.directionStatement?.()) {
      this.visitDirectionStatementImpl(ctx.directionStatement()!);
+    } else if (ctx.classDefStatement?.()) {
+      const classDefStmt = ctx.classDefStatement();
+      if (classDefStmt) {
+        this.visitClassDefStatementImpl(classDefStmt);
+      }
+    } else if (ctx.classStatement?.()) {
+      const classStmt = ctx.classStatement();
+      if (classStmt) {
+        this.visitClassStatementImpl(classStmt);
+      }
+    } else if (ctx.styleStatement?.()) {
+      this.visitStyleStatementImpl(ctx.styleStatement());
+    } else if (ctx.usecaseStatement?.()) {
+      const usecaseStmt = ctx.usecaseStatement();
+      if (usecaseStmt) {
+        this.visitUsecaseStatementImpl(usecaseStmt);
+      }
    }
    // NEWLINE is ignored
  }
@@ -269,9 +297,15 @@ export class UsecaseAntlrVisitor extends UsecaseVisitor<void> {

  /**
   * Visit usecaseInBoundary rule
-   * Grammar: usecaseInBoundary : IDENTIFIER | STRING ;
+   * Grammar: usecaseInBoundary : usecaseWithClass | IDENTIFIER | STRING ;
   */
  private visitUsecaseInBoundaryImpl(ctx: UsecaseInBoundaryContext): string {
+    // Check for usecaseWithClass (e.g., "debugging:::case1")
+    const usecaseWithClass = ctx.usecaseWithClass?.();
+    if (usecaseWithClass) {
+      return this.visitUsecaseWithClassImpl(usecaseWithClass);
+    }
+
    const identifier = ctx.IDENTIFIER?.();
    if (identifier) {
      return identifier.getText();
@@ -287,6 +321,37 @@ export class UsecaseAntlrVisitor extends UsecaseVisitor<void> {
    return '';
  }

+  /**
+   * Visit usecaseWithClass rule
+   * Grammar: usecaseWithClass : IDENTIFIER CLASS_SEPARATOR IDENTIFIER | STRING CLASS_SEPARATOR IDENTIFIER ;
+   */
+  private visitUsecaseWithClassImpl(ctx: UsecaseWithClassContext): string {
+    let usecaseName = '';
+    let className = '';
+
+    const identifier0 = ctx.IDENTIFIER(0);
+    const identifier1 = ctx.IDENTIFIER(1);
+    const string = ctx.STRING();
+
+    if (identifier0 && identifier1) {
+      // IDENTIFIER:::IDENTIFIER
+      usecaseName = identifier0.getText();
+      className = identifier1.getText();
+    } else if (string && identifier0) {
+      // STRING:::IDENTIFIER
+      const text = string.getText();
+      usecaseName = text.slice(1, -1); // Remove quotes
+      className = identifier0.getText();
+    }
+
+    // Apply class to the use case
+    if (usecaseName && className) {
+      this.applyClassToEntity(usecaseName, className);
+    }
+
+    return usecaseName;
+  }
+
  /**
   * Visit systemBoundaryTypeStatement rule
   * Grammar: systemBoundaryTypeStatement : systemBoundaryName '\@' '\{' systemBoundaryTypeContent '\}' NEWLINE* ;
@@ -432,15 +497,46 @@ export class UsecaseAntlrVisitor extends UsecaseVisitor<void> {

  /**
   * Visit entityName rule
-   * Grammar: entityName : IDENTIFIER | STRING | nodeIdWithLabel ;
+   * Grammar: entityName : IDENTIFIER CLASS_SEPARATOR IDENTIFIER | STRING CLASS_SEPARATOR IDENTIFIER | IDENTIFIER | STRING | nodeIdWithLabel ;
   */
  private visitEntityNameImpl(ctx: EntityNameContext): string {
-    const identifier = ctx.IDENTIFIER?.();
+    const classSeparator = ctx.CLASS_SEPARATOR?.();
+
+    // Check for class application syntax (e.g., "debugging:::case1")
+    if (classSeparator) {
+      let entityName = '';
+      let className = '';
+
+      const identifier0 = ctx.IDENTIFIER(0);
+      const identifier1 = ctx.IDENTIFIER(1);
+      const string0 = ctx.STRING();
+
+      if (identifier0 && identifier1) {
+        // IDENTIFIER:::IDENTIFIER
+        entityName = identifier0.getText();
+        className = identifier1.getText();
+      } else if (string0 && identifier0) {
+        // STRING:::IDENTIFIER
+        const text = string0.getText();
+        entityName = text.slice(1, -1); // Remove quotes
+        className = identifier0.getText();
+      }
+
+      // Apply class to the entity
+      if (entityName && className) {
+        this.applyClassToEntity(entityName, className);
+      }
+
+      return entityName;
+    }
+
+    // Regular entity name without class
+    const identifier = ctx.IDENTIFIER(0);
    if (identifier) {
      return identifier.getText();
    }

-    const string = ctx.STRING?.();
+    const string = ctx.STRING();
    if (string) {
      const text = string.getText();
      // Remove quotes from string
@@ -455,6 +551,30 @@ export class UsecaseAntlrVisitor extends UsecaseVisitor<void> {
    return '';
  }

+  /**
+   * Apply a class to an entity (use case)
+   */
+  private applyClassToEntity(entityName: string, className: string): void {
+    // Find or create the use case
+    let useCase = this.useCases.find((uc) => uc.id === entityName);
+    if (!useCase) {
+      useCase = {
+        id: entityName,
+        name: entityName,
+        classes: [],
+      };
+      this.useCases.push(useCase);
+    }
+
+    // Add the class if not already present
+    if (!useCase.classes) {
+      useCase.classes = [];
+    }
+    if (!useCase.classes.includes(className)) {
+      useCase.classes.push(className);
+    }
+  }
+
  /**
   * Visit actorDeclaration rule
   * Grammar: actorDeclaration : 'actor' actorName ;
@@ -544,7 +664,7 @@ export class UsecaseAntlrVisitor extends UsecaseVisitor<void> {

  /**
   * Visit arrow rule
-   * Grammar: arrow : SOLID_ARROW | BACK_ARROW | LINE_SOLID | labeledArrow ;
+   * Grammar: arrow : SOLID_ARROW | BACK_ARROW | LINE_SOLID | CIRCLE_ARROW | CROSS_ARROW | CIRCLE_ARROW_REVERSED | CROSS_ARROW_REVERSED | labeledArrow ;
   */
  private visitArrowImpl(ctx: ArrowContext): { arrowType: ArrowType; label?: string } {
    // Check if this is a labeled arrow
@@ -559,6 +679,14 @@ export class UsecaseAntlrVisitor extends UsecaseVisitor<void> {
      return { arrowType: ARROW_TYPE.BACK_ARROW };
    } else if (ctx.LINE_SOLID()) {
      return { arrowType: ARROW_TYPE.LINE_SOLID };
+    } else if (ctx.CIRCLE_ARROW()) {
+      return { arrowType: ARROW_TYPE.CIRCLE_ARROW };
+    } else if (ctx.CROSS_ARROW()) {
+      return { arrowType: ARROW_TYPE.CROSS_ARROW };
+    } else if (ctx.CIRCLE_ARROW_REVERSED()) {
+      return { arrowType: ARROW_TYPE.CIRCLE_ARROW_REVERSED };
+    } else if (ctx.CROSS_ARROW_REVERSED()) {
+      return { arrowType: ARROW_TYPE.CROSS_ARROW_REVERSED };
    }

    // Fallback (should not happen with proper grammar)
@@ -567,7 +695,7 @@ export class UsecaseAntlrVisitor extends UsecaseVisitor<void> {

  /**
   * Visit labeled arrow rule
-   * Grammar: labeledArrow : LINE_SOLID edgeLabel SOLID_ARROW | BACK_ARROW edgeLabel LINE_SOLID | LINE_SOLID edgeLabel LINE_SOLID ;
+   * Grammar: labeledArrow : LINE_SOLID edgeLabel SOLID_ARROW | BACK_ARROW edgeLabel LINE_SOLID | LINE_SOLID edgeLabel LINE_SOLID | LINE_SOLID edgeLabel CIRCLE_ARROW | LINE_SOLID edgeLabel CROSS_ARROW | CIRCLE_ARROW_REVERSED edgeLabel LINE_SOLID | CROSS_ARROW_REVERSED edgeLabel LINE_SOLID ;
   */
  private visitLabeledArrowImpl(ctx: LabeledArrowContext): { arrowType: ArrowType; label: string } {
    const label = this.visitEdgeLabelImpl(ctx.edgeLabel());
@@ -577,6 +705,14 @@ export class UsecaseAntlrVisitor extends UsecaseVisitor<void> {
      return { arrowType: ARROW_TYPE.SOLID_ARROW, label };
    } else if (ctx.BACK_ARROW()) {
      return { arrowType: ARROW_TYPE.BACK_ARROW, label };
+    } else if (ctx.CIRCLE_ARROW()) {
+      return { arrowType: ARROW_TYPE.CIRCLE_ARROW, label };
+    } else if (ctx.CROSS_ARROW()) {
+      return { arrowType: ARROW_TYPE.CROSS_ARROW, label };
+    } else if (ctx.CIRCLE_ARROW_REVERSED()) {
+      return { arrowType: ARROW_TYPE.CIRCLE_ARROW_REVERSED, label };
+    } else if (ctx.CROSS_ARROW_REVERSED()) {
+      return { arrowType: ARROW_TYPE.CROSS_ARROW_REVERSED, label };
    } else {
      return { arrowType: ARROW_TYPE.LINE_SOLID, label };
    }
@@ -622,6 +758,126 @@ export class UsecaseAntlrVisitor extends UsecaseVisitor<void> {
    return text;
  }

+  /**
+   * Visit classDefStatement rule
+   * Grammar: classDefStatement : 'classDef' IDENTIFIER stylesOpt NEWLINE* ;
+   */
+  visitClassDefStatementImpl(ctx: ClassDefStatementContext): void {
+    const className = ctx.IDENTIFIER().getText();
+    const stylesOptCtx = ctx.stylesOpt();
+
+    // Get all style properties as an array of strings
+    const styles = this.visitStylesOptImpl(stylesOptCtx);
+
+    this.classDefs.set(className, {
+      id: className,
+      styles,
+    });
+  }
+
+  /**
+   * Visit stylesOpt rule
+   * Grammar: stylesOpt : style | stylesOpt COMMA style ;
+   * Returns an array of style strings like ['stroke:#f00', 'fill:#ff0']
+   */
+  private visitStylesOptImpl(ctx: any): string[] {
+    const styles: string[] = [];
+
+    // Check if this is a recursive stylesOpt (stylesOpt COMMA style)
+    const stylesOptCtx = ctx.stylesOpt?.();
+    if (stylesOptCtx) {
+      styles.push(...this.visitStylesOptImpl(stylesOptCtx));
+    }
+
+    // Get the style context
+    const styleCtx = ctx.style();
+    if (styleCtx) {
+      const styleText = this.visitStyleImpl(styleCtx);
+      styles.push(styleText);
+    }
+
+    return styles;
+  }
+
+  /**
+   * Visit style rule
+   * Grammar: style : styleComponent | style styleComponent ;
+   * Returns a single style string like 'stroke:#f00'
+   */
+  private visitStyleImpl(ctx: any): string {
+    // Get all text from the style context
+    return ctx.getText();
+  }
+
+  /**
+   * Visit classStatement rule
+   * Grammar: classStatement : 'class' nodeList IDENTIFIER NEWLINE* ;
+   */
+  visitClassStatementImpl(ctx: ClassStatementContext): void {
+    const nodeIds = this.visitNodeListImpl(ctx.nodeList());
+    const className = ctx.IDENTIFIER().getText();
+
+    // Apply class to each node
+    nodeIds.forEach((nodeId) => {
+      this.applyClassToEntity(nodeId, className);
+    });
+  }
+
+  /**
+   * Visit styleStatement rule
+   * Grammar: styleStatement : 'style' IDENTIFIER stylesOpt NEWLINE* ;
+   */
+  visitStyleStatementImpl(ctx: any): void {
+    const nodeId = ctx.IDENTIFIER().getText();
+    const stylesOptCtx = ctx.stylesOpt();
+
+    // Get all style properties as an array of strings
+    const styles = this.visitStylesOptImpl(stylesOptCtx);
+
+    // Apply styles directly to the entity
+    let entity = this.useCases.find((uc) => uc.id === nodeId);
+    if (!entity) {
+      entity = this.actors.find((a) => a.id === nodeId);
+    }
+    if (!entity) {
+      entity = this.systemBoundaries.find((sb) => sb.id === nodeId);
+    }
+
+    if (entity) {
+      // Initialize styles array if it doesn't exist
+      if (!entity.styles) {
+        entity.styles = [];
+      }
+      // Add the new styles
+      entity.styles.push(...styles);
+    }
+  }
+
+  /**
+   * Visit nodeList rule
+   * Grammar: nodeList : IDENTIFIER (',' IDENTIFIER)* ;
+   */
+  private visitNodeListImpl(ctx: NodeListContext): string[] {
+    const identifiers = ctx.IDENTIFIER();
+    return identifiers.map((id) => id.getText());
+  }
+
+  /**
+   * Visit usecaseStatement rule
+   * Grammar: usecaseStatement : entityName NEWLINE* ;
+   */
+  visitUsecaseStatementImpl(ctx: UsecaseStatementContext): void {
+    const entityName = this.visitEntityNameImpl(ctx.entityName());
+
+    // Create a standalone use case if it doesn't already exist
+    if (!this.useCases.some((uc) => uc.id === entityName)) {
+      this.useCases.push({
+        id: entityName,
+        name: entityName,
+      });
+    }
+  }
+
  /**
   * Get the parse result after visiting the diagram
   */
@@ -631,6 +887,7 @@ export class UsecaseAntlrVisitor extends UsecaseVisitor<void> {
      useCases: this.useCases,
      systemBoundaries: this.systemBoundaries,
      relationships: this.relationships,
+      classDefs: this.classDefs,
      direction: this.direction,
    };
  }