From 3ceb3829a047417b35e1b159ce60253f7573d4e4 Mon Sep 17 00:00:00 2001
From: Matthew Lipski <matthewlipski@gmail.com>
Date: Mon, 3 Jun 2024 18:43:31 +0200
Subject: [PATCH] Added custom logic to better handle block un-nesting

---
 packages/core/src/editor/BlockNoteEditor.ts  |  7 +-
 packages/core/src/pm-nodes/BlockContainer.ts | 87 ++++++++++++++++++--
 2 files changed, 86 insertions(+), 8 deletions(-)

diff --git a/packages/core/src/editor/BlockNoteEditor.ts b/packages/core/src/editor/BlockNoteEditor.ts
index 53ec4d9f8..69816fb2c 100644
--- a/packages/core/src/editor/BlockNoteEditor.ts
+++ b/packages/core/src/editor/BlockNoteEditor.ts
@@ -917,7 +917,12 @@ export class BlockNoteEditor<
    * Lifts the block containing the text cursor out of its parent.
    */
   public unnestBlock() {
-    this._tiptapEditor.commands.liftListItem("blockContainer");
+    const { startPos } = getBlockInfoFromPos(
+      this._tiptapEditor.state.doc,
+      this._tiptapEditor.state.selection.from
+    )!;
+
+    this._tiptapEditor.commands.BNUnnestBlock(startPos);
   }
 
   // TODO: Fix when implementing HTML/Markdown import & export
diff --git a/packages/core/src/pm-nodes/BlockContainer.ts b/packages/core/src/pm-nodes/BlockContainer.ts
index 9050119f8..8cb23a164 100644
--- a/packages/core/src/pm-nodes/BlockContainer.ts
+++ b/packages/core/src/pm-nodes/BlockContainer.ts
@@ -57,6 +57,7 @@ declare module "@tiptap/core" {
         posInBlock: number,
         block: PartialBlock<BSchema, I, S>
       ) => ReturnType;
+      BNUnnestBlock: (posInBlock: number) => ReturnType;
     };
   }
 }
@@ -384,7 +385,6 @@ export const BlockContainer = Node.create<{
           }
 
           // Deletes next block and adds its text content to the nearest previous block.
-
           if (dispatch) {
             dispatch(
               state.tr
@@ -484,6 +484,79 @@ export const BlockContainer = Node.create<{
 
           return true;
         },
+      BNUnnestBlock:
+        (posInBlock) =>
+        ({ state, dispatch }) => {
+          const blockInfo = getBlockInfoFromPos(state.tr.doc, posInBlock)!;
+
+          if (blockInfo.depth > 2) {
+            if (dispatch) {
+              const parentBlockInfo = getBlockInfoFromPos(
+                state.tr.doc,
+                state.tr.doc.resolve(blockInfo.startPos).start(-2)
+              )!;
+
+              const blockChildren =
+                blockInfo.numChildBlocks > 0
+                  ? state.tr.doc
+                      .resolve(
+                        blockInfo.startPos + blockInfo.contentNode.nodeSize + 1
+                      )
+                      .node().content
+                  : Fragment.empty;
+              const siblingBlocksAfter = state.tr.doc.slice(
+                blockInfo.endPos + 1,
+                parentBlockInfo.endPos - 1
+              ).content;
+              // We need to move all siblings after the block into its children.
+              const children = blockChildren.append(siblingBlocksAfter);
+
+              // Checks if the block is the first child of its parent, as
+              // this means that we need to delete the entire `blockGroup`
+              // of the parent. Otherwise, we only need to delete the
+              // children after the block.
+              if (
+                parentBlockInfo.startPos +
+                  parentBlockInfo.contentNode.nodeSize +
+                  1 ===
+                blockInfo.startPos - 1
+              ) {
+                state.tr.delete(
+                  parentBlockInfo.startPos +
+                    parentBlockInfo.contentNode.nodeSize,
+                  parentBlockInfo.endPos
+                );
+              } else {
+                state.tr.delete(
+                  blockInfo.startPos - 1,
+                  parentBlockInfo.endPos - 1
+                );
+              }
+
+              const blockGroup = state.schema.nodes["blockGroup"].create(
+                null,
+                children
+              );
+              const block = state.schema.nodes["blockContainer"].create(
+                blockInfo.node.attrs,
+                [blockInfo.contentNode, blockGroup]
+              );
+
+              state.tr.insert(state.tr.selection.from - 2, block);
+              state.tr.setSelection(
+                new TextSelection(
+                  state.tr.doc.resolve(state.tr.selection.from - block.nodeSize)
+                )
+              );
+
+              dispatch(state.tr);
+            }
+
+            return true;
+          }
+
+          return false;
+        },
     };
   },
 
@@ -522,15 +595,15 @@ export const BlockContainer = Node.create<{
         // Removes a level of nesting if the block is indented if the selection is at the start of the block.
         () =>
           commands.command(({ state }) => {
-            const { startPos } = getBlockInfoFromPos(
-              state.doc,
+            const { startPos, depth } = getBlockInfoFromPos(
+              state.tr.doc,
               state.selection.from
             )!;
 
             const selectionAtBlockStart = state.selection.from === startPos + 1;
 
-            if (selectionAtBlockStart) {
-              return commands.liftListItem("blockContainer");
+            if (selectionAtBlockStart && depth > 2) {
+              return commands.BNUnnestBlock(startPos);
             }
 
             return false;
@@ -711,7 +784,7 @@ export const BlockContainer = Node.create<{
         this.editor.commands.sinkListItem("blockContainer");
         return true;
       },
-      "Shift-Tab": () => {
+      "Shift-Tab": ({ editor }) => {
         if (
           this.options.editor.formattingToolbar?.shown ||
           this.options.editor.linkToolbar?.shown ||
@@ -720,7 +793,7 @@ export const BlockContainer = Node.create<{
           // don't handle tabs if a toolbar is shown, so we can tab into / out of it
           return false;
         }
-        this.editor.commands.liftListItem("blockContainer");
+        this.editor.commands.BNUnnestBlock(editor.state.selection.from);
         return true;
       },
     };