From 70b948fd178d33ea9c3d2364fc4b8c94e1f39fe4 Mon Sep 17 00:00:00 2001
From: Sidharth Vinod <sidharthv96@gmail.com>
Date: Thu, 7 Dec 2023 00:36:21 +0530
Subject: [PATCH] refactor: Optimize renderKatex

---
 cSpell.json                                   |  1 +
 packages/mermaid/package.json                 |  1 +
 .../mermaid/src/diagrams/common/common.ts     | 50 ++++++++++---------
 pnpm-lock.yaml                                |  7 +++
 4 files changed, 36 insertions(+), 23 deletions(-)

diff --git a/cSpell.json b/cSpell.json
index d79d7950e..88d2553ee 100644
--- a/cSpell.json
+++ b/cSpell.json
@@ -80,6 +80,7 @@
     "logmsg",
     "lucida",
     "markdownish",
+    "mathml",
     "matthieu",
     "matthieumorel",
     "mdast",
diff --git a/packages/mermaid/package.json b/packages/mermaid/package.json
index 7a2194d77..db92e5793 100644
--- a/packages/mermaid/package.json
+++ b/packages/mermaid/package.json
@@ -91,6 +91,7 @@
     "@types/d3-shape": "^3.1.1",
     "@types/dompurify": "^3.0.2",
     "@types/jsdom": "^21.1.1",
+    "@types/katex": "^0.16.7",
     "@types/lodash-es": "^4.17.7",
     "@types/micromatch": "^4.0.2",
     "@types/prettier": "^2.7.2",
diff --git a/packages/mermaid/src/diagrams/common/common.ts b/packages/mermaid/src/diagrams/common/common.ts
index 84133f563..3fd5806c2 100644
--- a/packages/mermaid/src/diagrams/common/common.ts
+++ b/packages/mermaid/src/diagrams/common/common.ts
@@ -318,33 +318,37 @@ export const calculateMathMLDimensions = async (text: string, config: MermaidCon
  * @returns String containing MathML if KaTeX is supported, or an error message if it is not and stylesheets aren't present
  */
 export const renderKatex = async (text: string, config: MermaidConfig): Promise<string> => {
-  if (hasKatex(text) && (isMathMLSupported() || config.legacyMathML)) {
-    // @ts-ignore @types/katex does not work
-    const katex = (await import('katex')).default;
-    return text
-      .split(lineBreakRegex)
-      .map((line) =>
-        hasKatex(line)
-          ? `
+  if (!hasKatex(text)) {
+    return text;
+  }
+
+  if (!isMathMLSupported() && !config.legacyMathML) {
+    return text.replace(katexRegex, 'MathML is unsupported in this environment.');
+  }
+
+  const { default: katex } = await import('katex');
+  return text
+    .split(lineBreakRegex)
+    .map((line) =>
+      hasKatex(line)
+        ? `
             <div style="display: flex; align-items: center; justify-content: center; white-space: nowrap;">
               ${line}
             </div>
           `
-          : `<div>${line}</div>`
-      )
-      .join('')
-      .replace(katexRegex, (r, c) =>
-        katex
-          .renderToString(c, {
-            throwOnError: true,
-            displayMode: true,
-            output: isMathMLSupported() ? 'mathml' : 'htmlAndMathml',
-          })
-          .replace(/\n/g, ' ')
-          .replace(/<annotation.*<\/annotation>/g, '')
-      );
-  }
-  return text.replace(katexRegex, 'MathML is unsupported in this environment.');
+        : `<div>${line}</div>`
+    )
+    .join('')
+    .replace(katexRegex, (_, c) =>
+      katex
+        .renderToString(c, {
+          throwOnError: true,
+          displayMode: true,
+          output: isMathMLSupported() ? 'mathml' : 'htmlAndMathml',
+        })
+        .replace(/\n/g, ' ')
+        .replace(/<annotation.*<\/annotation>/g, '')
+    );
 };
 
 export default {
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index ff3ddf0b5..38ec0c9a7 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -279,6 +279,9 @@ importers:
       '@types/jsdom':
         specifier: ^21.1.1
         version: 21.1.1
+      '@types/katex':
+        specifier: ^0.16.7
+        version: 0.16.7
       '@types/lodash-es':
         specifier: ^4.17.7
         version: 4.17.7
@@ -4944,6 +4947,10 @@ packages:
     resolution: {integrity: sha512-Hr5Jfhc9eYOQNPYO5WLDq/n4jqijdHNlDXjuAQkkt+mWdQR+XJToOHrsD4cPaMXpn6KO7y2+wM8AZEs8VpBLVA==}
     dev: true
 
+  /@types/katex@0.16.7:
+    resolution: {integrity: sha512-HMwFiRujE5PjrgwHQ25+bsLJgowjGjm5Z8FVSf0N6PwgJrwxH0QxzHYDcKsTfV3wva0vzrpqMTJS2jXPr5BMEQ==}
+    dev: true
+
   /@types/keyv@3.1.4:
     resolution: {integrity: sha512-BQ5aZNSCpj7D6K2ksrRCTmKRLEpnPvWDiLPfoGyhZ++8YtiK9d/3DBKPJgry359X/P1PfruyYwvnvwFjuEiEIg==}
     dependencies: