fix: Deleting last block in column

packages/core/src/api/blockManipulation/commands/replaceBlocks/replaceBlocks.ts

@@ -1,5 +1,5 @@
-import type { Node } from "prosemirror-model";
-import type { Transaction } from "prosemirror-state";
+import { type Node } from "prosemirror-model";
+import { type Transaction } from "prosemirror-state";
 import type { Block, PartialBlock } from "../../../../blocks/defaultBlocks.js";
 import type {
   BlockIdentifier,
@@ -10,6 +10,7 @@ import type {
 import { blockToNode } from "../../../nodeConversions/blockToNode.js";
 import { nodeToBlock } from "../../../nodeConversions/nodeToBlock.js";
 import { getPmSchema } from "../../../pmUtil.js";
+import { fixColumnList } from "./util/fixColumnList.js";
 
 export function removeAndInsertBlocks<
   BSchema extends BlockSchema,
@@ -70,26 +71,61 @@ export function removeAndInsertBlocks<
     }
 
     const oldDocSize = tr.doc.nodeSize;
-    // Checks if the block is the only child of its parent. In this case, we
-    // need to delete the parent `blockGroup` node instead of just the
-    // `blockContainer`.
+
     const $pos = tr.doc.resolve(pos - removedSize);
     if (
+      node.type.name === "column" &&
+      node.attrs.id !== $pos.nodeAfter?.attrs.id
+    ) {
+      // This is a hacky work around to handle removing all columns in a
+      // columnList. This is what happens when removing the last 2 columns:
+      //
+      // 1. The second-to-last `column` is removed.
+      // 2. `fixColumnList` runs, removing the `columnList` and inserting the
+      // contents of the last column in its place.
+      // 3. `removedSize` increases not just by the size of the second-to-last
+      // `column`, but also by the positions removed due to running
+      // `fixColumnList`. Some of these positions are after the contents of the
+      // last `column`, namely just after the `column` and `columnList`.
+      // 3. `tr.doc.descendants` traverses to the last `column`.
+      // 4. `removedSize` now includes positions that were removed after the
+      // last `column`. This causes `pos - removedSize` to point to an
+      // incorrect position, as it expects that the difference in document size
+      // accounted for by `removedSize` comes before the block being removed.
+      // 5. The deletion is offset by 3, because of those removed positions
+      // included in `removedSize` that occur after the last `column`.
+      //
+      // Hence why we have to shift the start of the deletion range back by 3.
+      // The offset for the end of the range is smaller as `node.nodeSize` is
+      // the size of the second `column`. Since it's been removed, we actually
+      // care about the size of its children - a difference of 2 positions.
+      tr.delete(pos - removedSize + 3, pos - removedSize + node.nodeSize + 1);
+    } else if (
       $pos.node().type.name === "blockGroup" &&
       $pos.node($pos.depth - 1).type.name !== "doc" &&
       $pos.node().childCount === 1
     ) {
+      // Checks if the block is the only child of a parent `blockGroup` node.
+      // In this case, we need to delete the parent `blockGroup` node instead
+      // of just the `blockContainer`.
       tr.delete($pos.before(), $pos.after());
     } else {
       tr.delete(pos - removedSize, pos - removedSize + node.nodeSize);
     }
+
+    if ($pos.node().type.name === "column") {
+      fixColumnList(tr, $pos.before(-1));
+    } else if ($pos.node().type.name === "columnList") {
+      fixColumnList(tr, $pos.before());
+    }
+
     const newDocSize = tr.doc.nodeSize;
     removedSize += oldDocSize - newDocSize;
 
     return false;
   });
 
-  // Throws an error if now all blocks could be found.
+  // Throws an error if not all blocks could be found.
   if (idsOfBlocksToRemove.size > 0) {
     const notFoundIds = [...idsOfBlocksToRemove].join("\n");
 

packages/core/src/api/blockManipulation/commands/replaceBlocks/util/fixColumnList.ts

@@ -0,0 +1,163 @@
+import { Slice, type Node } from "prosemirror-model";
+import { type Transaction } from "prosemirror-state";
+import { ReplaceAroundStep } from "prosemirror-transform";
+
+/**
+ * Checks if a `column` node is empty, i.e. if it has only a single empty
+ * block.
+ * @param column The column to check.
+ * @returns Whether the column is empty.
+ */
+export function isEmptyColumn(column: Node) {
+  if (!column || column.type.name !== "column") {
+    throw new Error("Invalid columnPos: does not point to column node.");
+  }
+
+  const blockContainer = column.firstChild;
+  if (!blockContainer) {
+    throw new Error("Invalid column: does not have child node.");
+  }
+
+  const blockContent = blockContainer.firstChild;
+  if (!blockContent) {
+    throw new Error("Invalid blockContainer: does not have child node.");
+  }
+
+  return (
+    column.childCount === 1 &&
+    blockContainer.childCount === 1 &&
+    blockContent.type.spec.content === "inline*" &&
+    blockContent.content.content.length === 0
+  );
+}
+
+/**
+ * Removes all empty `column` nodes in a `columnList`. A `column` node is empty
+ * if it has only a single empty block. If, however, removing the `column`s
+ * leaves the `columnList` that has fewer than two, ProseMirror will re-add
+ * empty columns.
+ * @param tr The `Transaction` to add the changes to.
+ * @param columnListPos The position just before the `columnList` node.
+ */
+export function removeEmptyColumns(tr: Transaction, columnListPos: number) {
+  const $columnListPos = tr.doc.resolve(columnListPos);
+  const columnList = $columnListPos.nodeAfter;
+  if (!columnList || columnList.type.name !== "columnList") {
+    throw new Error(
+      "Invalid columnListPos: does not point to columnList node.",
+    );
+  }
+
+  for (
+    let columnIndex = columnList.childCount - 1;
+    columnIndex >= 0;
+    columnIndex--
+  ) {
+    const columnPos = tr.doc
+      .resolve($columnListPos.pos + 1)
+      .posAtIndex(columnIndex);
+    const $columnPos = tr.doc.resolve(columnPos);
+    const column = $columnPos.nodeAfter;
+    if (!column || column.type.name !== "column") {
+      throw new Error("Invalid columnPos: does not point to column node.");
+    }
+
+    if (isEmptyColumn(column)) {
+      tr.delete(columnPos, columnPos + column?.nodeSize);
+    }
+  }
+}
+
+/**
+ * Fixes potential issues in a `columnList` node after a
+ * `blockContainer`/`column` node is (re)moved from it:
+ *
+ * - Removes all empty `column` nodes. A `column` node is empty if it has only
+ * a single empty block.
+ * - If all but one `column` nodes are empty, replaces the `columnList` with
+ * the content of the non-empty `column`.
+ * - If all `column` nodes are empty, removes the `columnList` entirely.
+ * @param tr The `Transaction` to add the changes to.
+ * @param columnListPos
+ * @returns The position just before the `columnList` node.
+ */
+export function fixColumnList(tr: Transaction, columnListPos: number) {
+  removeEmptyColumns(tr, columnListPos);
+
+  const $columnListPos = tr.doc.resolve(columnListPos);
+  const columnList = $columnListPos.nodeAfter;
+  if (!columnList || columnList.type.name !== "columnList") {
+    throw new Error(
+      "Invalid columnListPos: does not point to columnList node.",
+    );
+  }
+
+  if (columnList.childCount > 2) {
+    // Do nothing if the `columnList` has at least two non-empty `column`s.
+    return;
+  }
+
+  if (columnList.childCount < 2) {
+    throw new Error("Invalid columnList: contains fewer than two children.");
+  }
+
+  const firstColumnBeforePos = columnListPos + 1;
+  const $firstColumnBeforePos = tr.doc.resolve(firstColumnBeforePos);
+  const firstColumn = $firstColumnBeforePos.nodeAfter;
+
+  const lastColumnAfterPos = columnListPos + columnList.nodeSize - 1;
+  const $lastColumnAfterPos = tr.doc.resolve(lastColumnAfterPos);
+  const lastColumn = $lastColumnAfterPos.nodeBefore;
+
+  if (!firstColumn || !lastColumn) {
+    throw new Error("Invalid columnList: does not contain children.");
+  }
+
+  const firstColumnEmpty = isEmptyColumn(firstColumn);
+  const lastColumnEmpty = isEmptyColumn(lastColumn);
+
+  if (firstColumnEmpty && lastColumnEmpty) {
+    // Removes `columnList`
+    tr.delete(columnListPos, columnListPos + columnList.nodeSize);
+
+    return;
+  }
+
+  if (firstColumnEmpty) {
+    tr.step(
+      new ReplaceAroundStep(
+        // Replaces `columnList`.
+        columnListPos,
+        columnListPos + columnList.nodeSize,
+        // Replaces with content of last `column`.
+        lastColumnAfterPos - lastColumn.nodeSize + 1,
+        lastColumnAfterPos - 1,
+        // Doesn't append anything.
+        Slice.empty,
+        0,
+        false,
+      ),
+    );
+
+    return;
+  }
+
+  if (lastColumnEmpty) {
+    tr.step(
+      new ReplaceAroundStep(
+        // Replaces `columnList`.
+        columnListPos,
+        columnListPos + columnList.nodeSize,
+        // Replaces with content of first `column`.
+        firstColumnBeforePos + 1,
+        firstColumnBeforePos + firstColumn.nodeSize - 1,
+        // Doesn't append anything.
+        Slice.empty,
+        0,
+        false,
+      ),
+    );
+
+    return;
+  }
+}