From d4f2f301d8a0659988cca31df1db9b62e23e04a8 Mon Sep 17 00:00:00 2001
From: Kalle Struik <kalle@kallestruik.nl>
Date: Wed, 16 Apr 2025 17:22:07 +0200
Subject: [PATCH] Make headers work properly when typed normally. Fixes #1

---
 frontend/src/editor/plugins/header_plugin.tsx | 113 +++++++++---------
 1 file changed, 59 insertions(+), 54 deletions(-)

diff --git a/frontend/src/editor/plugins/header_plugin.tsx b/frontend/src/editor/plugins/header_plugin.tsx
index d4afdc4..7da6733 100644
--- a/frontend/src/editor/plugins/header_plugin.tsx
+++ b/frontend/src/editor/plugins/header_plugin.tsx
@@ -1,8 +1,9 @@
 import { useLexicalComposerContext } from "@lexical/react/LexicalComposerContext";
-import { $createParagraphNode, $isParagraphNode, $isTextNode, SELECTION_CHANGE_COMMAND, TextNode } from "lexical";
+import { $createParagraphNode, $isParagraphNode, $isTextNode, SELECTION_CHANGE_COMMAND, TextNode, type LexicalEditor, type LexicalNode } from "lexical";
 import { useEffect } from "react";
 import { $createHeaderMarkerNode, $createHeaderNode, $isHeaderMarkerNode, $isHeaderNode, HeaderMarkerNode, HeaderNode } from "../nodes/header_node";
 import { mergeRegister } from "@lexical/utils";
+import { findBlockNode } from "../editor_utils";
 
 const HEADER_REGEX = /^#+ /;
 
@@ -11,68 +12,72 @@ export function HeaderPlugin() {
 
     useEffect(() => mergeRegister(
         // Create a header node if the text node matches the HEADER_REGEX
-        editor.registerNodeTransform(TextNode, (textNode) => {
-            const prevNode = textNode.getPreviousSibling();
-            if (prevNode) return;
-            const paragraphNode = textNode.getParent();
-            if (!$isParagraphNode(paragraphNode)) return;
-
-            const content = textNode.getTextContent();
-            const regexMatch = content.match(HEADER_REGEX);
-            if (!regexMatch) return;
-
-            const children = paragraphNode.getChildren();
-
-            const firstTextNode = children[0];
-            if (!$isTextNode(firstTextNode)) return;
-
-            const markerLength = regexMatch[0].length;
-            const textNodes = firstTextNode.splitText(markerLength);
-
-            const headerMarkerContent = textNodes[0];
-            if (!headerMarkerContent) return;
-
-            const headerNode = $createHeaderNode(markerLength - 1);
-            const headerMarkerNode = $createHeaderMarkerNode();
-
-            headerMarkerNode.append(headerMarkerContent);
-
-            headerNode.append(headerMarkerNode);
-            headerNode.append(...textNodes.slice(1));
-            headerNode.append(...children.slice(1));
-
-            paragraphNode.replace(headerNode, true);
-
-            editor.dispatchCommand(SELECTION_CHANGE_COMMAND, undefined);
-        }),
-        // Remove headers if they don't match the HEADER_REGEX
-        editor.registerNodeTransform(HeaderNode, (headerNode) => {
-            const content = headerNode.getTextContent();
-            if (content.match(HEADER_REGEX)) return;
-            headerNode.replace($createParagraphNode(), true);
-        }),
-        // Remove header markers if they don't match the HEADER_REGEX
+        editor.registerNodeTransform(TextNode, createHeaderTransform(editor)),
+        editor.registerNodeTransform(TextNode, ensureHeaderValidTransform),
+        editor.registerNodeTransform(HeaderNode, ensureHeaderValidTransform),
+        // Remove header markers if they aren't in a header
         editor.registerNodeTransform(HeaderMarkerNode, (node) => {
             const headerNode = node.getParent();
-            const content = node.getTextContent();
-            if ($isHeaderNode(headerNode) &&
-                content.match(HEADER_REGEX) &&
-                content.length - 1 == headerNode.getLevel()) {
+            if ($isHeaderNode(headerNode)) {
                 return;
             }
 
             node.getChildren().reverse().forEach(child => node.insertAfter(child));
             node.remove();
         }),
-        // Remove header nodes without a header marker
-        editor.registerNodeTransform(HeaderNode, (node) => {
-            const children = node.getChildren();
-            const headerMarker = children[0];
-            if ($isHeaderMarkerNode(headerMarker)) return;
-
-            node.replace($createParagraphNode(), true);
-        }),
     ));
 
     return null;
 }
+
+const createHeaderTransform = (editor: LexicalEditor) => (textNode: TextNode) => {
+    const prevNode = textNode.getPreviousSibling();
+    if (prevNode) return;
+    const paragraphNode = textNode.getParent();
+    if (!$isParagraphNode(paragraphNode)) return;
+
+    const content = textNode.getTextContent();
+    const regexMatch = content.match(HEADER_REGEX);
+    if (!regexMatch) return;
+
+    const children = paragraphNode.getChildren();
+
+    const firstTextNode = children[0];
+    if (!$isTextNode(firstTextNode)) return;
+
+    const markerLength = regexMatch[0].length;
+    const textNodes = firstTextNode.splitText(markerLength);
+
+    const headerMarkerContent = textNodes[0];
+    if (!headerMarkerContent) return;
+
+    const headerNode = $createHeaderNode(markerLength - 1);
+    const headerMarkerNode = $createHeaderMarkerNode();
+
+    headerMarkerNode.append(headerMarkerContent);
+
+    headerNode.append(headerMarkerNode);
+    headerNode.append(...textNodes.slice(1));
+    headerNode.append(...children.slice(1));
+
+    paragraphNode.replace(headerNode, true);
+
+    editor.dispatchCommand(SELECTION_CHANGE_COMMAND, undefined);
+}
+
+const ensureHeaderValidTransform = (node: LexicalNode) => {
+    const headerNode = findBlockNode(node);
+    if (!$isHeaderNode(headerNode)) {
+        return;
+    }
+
+    const markerNode = headerNode.getFirstChild();
+    const markerContent = markerNode?.getTextContent();
+    if ($isHeaderMarkerNode(markerNode)
+        && markerContent?.match(HEADER_REGEX)
+        && markerContent.length == headerNode.getLevel() + 1) {
+        return
+    }
+
+    headerNode.replace($createParagraphNode(), true);
+}