From cfe7cce41d3dd26757ba8acad669b39589b82916 Mon Sep 17 00:00:00 2001
From: Alois Klink <alois@aloisklink.com>
Date: Thu, 17 Oct 2024 18:59:43 +0900
Subject: [PATCH] refactor: improve types of `shapes`

Instead of being a `Record<string, ShapeHandler>`, the type now knows
every key in the variable.
---
 .../rendering-elements/nodes.spec.ts          |  2 +-
 .../rendering-elements/shapes.ts              | 36 +++++++++++--------
 2 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/packages/mermaid/src/rendering-util/rendering-elements/nodes.spec.ts b/packages/mermaid/src/rendering-util/rendering-elements/nodes.spec.ts
index a7a7691e5..c1f0e1437 100644
--- a/packages/mermaid/src/rendering-util/rendering-elements/nodes.spec.ts
+++ b/packages/mermaid/src/rendering-util/rendering-elements/nodes.spec.ts
@@ -84,7 +84,7 @@ describe('Test Alias for shapes', function () {
   });
 
   it('should support alias for shadedProcess shape ', function () {
-    const aliases = ['lined-process', 'lined-rectangle', 'lin-proc', 'lin-rect'];
+    const aliases = ['lined-process', 'lined-rectangle', 'lin-proc', 'lin-rect'] as const;
     for (const alias of aliases) {
       expect(shapes[alias]).toBe(shapes['shaded-process']);
     }
diff --git a/packages/mermaid/src/rendering-util/rendering-elements/shapes.ts b/packages/mermaid/src/rendering-util/rendering-elements/shapes.ts
index b336fc823..3ff98c4cd 100644
--- a/packages/mermaid/src/rendering-util/rendering-elements/shapes.ts
+++ b/packages/mermaid/src/rendering-util/rendering-elements/shapes.ts
@@ -1,3 +1,4 @@
+import type { Entries } from 'type-fest';
 import type { Node, ShapeRenderOptions } from '../types.js';
 import { anchor } from './shapes/anchor.js';
 import { bowTieRect } from './shapes/bowTieRect.js';
@@ -75,7 +76,7 @@ export interface ShapeDefinition {
   handler: ShapeHandler;
 }
 
-export const shapesDefs: ShapeDefinition[] = [
+export const shapesDefs = [
   {
     semanticName: 'Process',
     name: 'Rectangle',
@@ -442,11 +443,11 @@ export const shapesDefs: ShapeDefinition[] = [
     aliases: ['lined-document'],
     handler: linedWaveEdgedRect,
   },
-];
+] as const satisfies ShapeDefinition[];
 
 const generateShapeMap = () => {
   // These are the shapes that didn't have documentation present
-  const shapeMap: Record<string, ShapeHandler> = {
+  const undocumentedShapes = {
     // States
     state,
     choice,
@@ -464,18 +465,25 @@ const generateShapeMap = () => {
     imageSquare,
 
     anchor,
-  };
+  } as const;
 
-  for (const shape of shapesDefs) {
-    for (const alias of [
-      shape.shortName,
-      ...(shape.aliases ?? []),
-      ...(shape.internalAliases ?? []),
-    ]) {
-      shapeMap[alias] = shape.handler;
-    }
-  }
-  return shapeMap;
+  const entries = [
+    ...(Object.entries(undocumentedShapes) as Entries<typeof undocumentedShapes>),
+    ...shapesDefs.flatMap((shape) => {
+      const aliases = [
+        shape.shortName,
+        ...('aliases' in shape ? shape.aliases : []),
+        ...('internalAliases' in shape ? shape.internalAliases : []),
+      ];
+      return aliases.map((alias) => [alias, shape.handler] as const);
+    }),
+  ];
+  return Object.fromEntries(entries) as Record<
+    (typeof entries)[number][0],
+    (typeof entries)[number][1]
+  > satisfies Record<string, ShapeHandler>;
 };
 
 export const shapes = generateShapeMap();
+
+export type ShapeID = keyof typeof shapes;