fix(core): keep comments when dragging table rows and columns

packages/core/src/extensions/TableHandles/TableHandles.test.ts

@@ -0,0 +1,127 @@
+import { describe, expect, it } from "vite-plus/test";
+
+import { CommentMark } from "../../comments/mark.js";
+import { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
+import { TableHandlesView } from "./TableHandles.js";
+
+/**
+ * @vitest-environment jsdom
+ */
+
+function createEditor() {
+  const editor = BlockNoteEditor.create({
+    extensions: [
+      createExtension({ key: "commentMark", tiptapExtensions: [CommentMark] }),
+    ],
+  });
+  editor.mount(document.body.appendChild(document.createElement("div")));
+  editor.replaceBlocks(editor.document, [
+    {
+      id: "table",
+      type: "table",
+      content: {
+        type: "tableContent",
+        columnWidths: [100, 200],
+        rows: [{ cells: ["a", "b"] }, { cells: ["c", "d"] }],
+      },
+    },
+  ]);
+
+  // Comment on the "a" cell.
+  const { state } = editor.prosemirrorView;
+  let from = -1;
+  state.doc.descendants((node, pos) => {
+    if (node.isText && node.text === "a") {
+      from = pos;
+    }
+  });
+  editor.prosemirrorView.dispatch(
+    state.tr.addMark(
+      from,
+      from + 1,
+      state.schema.marks.comment.create({ threadId: "thread" }),
+    ),
+  );
+
+  return editor;
+}
+
+// Simulates dropping a row/column handle that was dragged from `from` to `to`.
+function drop(
+  editor: BlockNoteEditor<any, any, any>,
+  orientation: "row" | "col",
+  from: number,
+  to: number,
+) {
+  const view = (editor.prosemirrorView as any).pluginViews.find(
+    (v: unknown) => v instanceof TableHandlesView,
+  ) as TableHandlesView;
+  // What hovering the table sets up.
+  view.tableId = "table";
+  view.tableElement =
+    editor.prosemirrorView.dom.querySelector<HTMLElement>('[data-id="table"]')!;
+  view.state = {
+    show: true,
+    showAddOrRemoveRowsButton: false,
+    showAddOrRemoveColumnsButton: false,
+    referencePosCell: undefined,
+    referencePosTable: new DOMRect(),
+    block: editor.getBlock("table") as any,
+    rowIndex: orientation === "row" ? to : 0,
+    colIndex: orientation === "col" ? to : 0,
+    draggingState: {
+      draggedCellOrientation: orientation,
+      originalIndex: from,
+      mousePos: 0,
+    },
+    widgetContainer: undefined,
+  };
+  expect(view.dropHandler(new Event("drop") as DragEvent)).toBe(true);
+}
+
+// Text of each cell, with the thread ID of any comment on it.
+function cells(editor: BlockNoteEditor<any, any, any>) {
+  const rows: string[][] = [];
+  editor.prosemirrorState.doc.descendants((node) => {
+    if (node.type.name === "tableRow") {
+      rows.push([]);
+    }
+    if (node.isText) {
+      const comment = node.marks.find((m) => m.type.name === "comment");
+      rows[rows.length - 1].push(
+        node.text! + (comment ? `@${comment.attrs.threadId}` : ""),
+      );
+    }
+  });
+  return rows;
+}
+
+describe("Table handles drag & drop", () => {
+  it("keeps comments when moving a row", () => {
+    const editor = createEditor();
+
+    drop(editor, "row", 0, 1);
+
+    expect(cells(editor)).toEqual([
+      ["c", "d"],
+      ["a@thread", "b"],
+    ]);
+    editor._tiptapEditor.destroy();
+  });
+
+  it("keeps comments and column widths when moving a column", () => {
+    const editor = createEditor();
+
+    drop(editor, "col", 0, 1);
+
+    expect(cells(editor)).toEqual([
+      ["b", "a@thread"],
+      ["d", "c"],
+    ]);
+    expect((editor.getBlock("table") as any).content.columnWidths).toEqual([
+      200, 100,
+    ]);
+    editor._tiptapEditor.destroy();
+  });
+});

packages/core/src/extensions/TableHandles/TableHandles.ts

@@ -1,5 +1,11 @@
 import type { Node } from "prosemirror-model";
-import { EditorState, Plugin, PluginKey, PluginView } from "prosemirror-state";
+import {
+  EditorState,
+  Plugin,
+  PluginKey,
+  PluginView,
+  TextSelection,
+} from "prosemirror-state";
 import {
   CellSelection,
   addColumnAfter,
@@ -9,6 +15,8 @@ import {
   deleteColumn,
   deleteRow,
   mergeCells,
+  moveTableColumn,
+  moveTableRow,
   splitCell,
 } from "prosemirror-tables";
 import { Decoration, DecorationSet, EditorView } from "prosemirror-view";
@@ -19,11 +27,10 @@ import {
   canColumnBeDraggedInto,
   canRowBeDraggedInto,
   cropEmptyRowsOrColumns,
+  getAbsoluteTableCells,
   getCellsAtColumnHandle,
   getCellsAtRowHandle,
   getDimensionsOfTable,
-  moveColumn,
-  moveRow,
 } from "../../api/blockManipulation/tables/tables.js";
 import { nodeToBlock } from "../../api/nodeConversions/nodeToBlock.js";
 import { getNodeById } from "../../api/nodeUtil.js";
@@ -472,62 +479,52 @@ export class TableHandlesView implements PluginView {
     // Clear so a re-dispatched drop short-circuits above (issue #2691).
     this.state.draggingState = undefined;
 
-    const columnWidths = this.state.block.content.columnWidths;
+    const block = this.state.block;
+    const isRow = draggingState.draggedCellOrientation === "row";
+    const targetIndex = isRow ? rowIndex : colIndex;
+    const canBeDraggedInto = isRow
+      ? canRowBeDraggedInto
+      : canColumnBeDraggedInto;
 
-    if (draggingState.draggedCellOrientation === "row") {
-      if (
-        !canRowBeDraggedInto(
-          this.state.block,
-          draggingState.originalIndex,
-          rowIndex,
-        )
-      ) {
-        // If the target row is invalid, don't move the row
-        return false;
-      }
-      const newTable = moveRow(
-        this.state.block,
-        draggingState.originalIndex,
-        rowIndex,
-      );
-      this.editor.updateBlock(this.state.block, {
-        type: "table",
-        content: {
-          ...this.state.block.content,
-          rows: newTable as any,
-        },
-      });
-    } else {
-      if (
-        !canColumnBeDraggedInto(
-          this.state.block,
-          draggingState.originalIndex,
-          colIndex,
-        )
-      ) {
-        // If the target column is invalid, don't move the column
+    if (!canBeDraggedInto(block, draggingState.originalIndex, targetIndex)) {
+      // If the target row/column is invalid, don't move it
+      return false;
+    }
+
+    // The handles use relative indices, prosemirror-tables absolute ones.
+    function toAbsoluteIndex(index: number) {
+      return isRow
+        ? getAbsoluteTableCells({ row: index, col: 0 }, block).row
+        : getAbsoluteTableCells({ row: 0, col: index }, block).col;
+    }
+    const from = toAbsoluteIndex(draggingState.originalIndex);
+    const to = toAbsoluteIndex(targetIndex);
+
+    // Moves the cells in place instead of rebuilding the table from block
+    // JSON, which would drop marks that aren't styles, like comments (issue
+    // #2904). Cell attributes, including column widths, move with the cells.
+    const moved = this.editor.exec((state, dispatch) => {
+      const tablePos = this.getTablePos(state.doc);
+      if (tablePos === undefined) {
         return false;
       }
-      const newTable = moveColumn(
-        this.state.block,
-        draggingState.originalIndex,
-        colIndex,
+      // prosemirror-tables finds the table from the selection.
+      const stateInTable = state.apply(
+        state.tr.setSelection(
+          TextSelection.near(state.doc.resolve(tablePos + 1)),
+        ),
       );
-      const [columnWidth] = columnWidths.splice(draggingState.originalIndex, 1);
-      columnWidths.splice(colIndex, 0, columnWidth);
-      this.editor.updateBlock(this.state.block, {
-        type: "table",
-        content: {
-          ...this.state.block.content,
-          columnWidths,
-          rows: newTable as any,
-        },
-      });
+      const move = isRow ? moveTableRow : moveTableColumn;
+      return move({ from, to, select: false })(stateInTable, dispatch);
+    });
+
+    if (!moved) {
+      return false;
     }
 
-    // Have to reset text cursor position to the block as `updateBlock` moves
-    // the existing selection out of the block.
-    this.editor.setTextCursorPosition(this.state.block.id);
+    // Have to reset text cursor position to the block as replacing the table
+    // moves the existing selection out of it.
+    this.editor.setTextCursorPosition(block.id);
 
     return true;
   };