From dfe36fcbf52b49e7e37239917286b32fafcdf363 Mon Sep 17 00:00:00 2001
From: Tom Moor <tom@getoutline.com>
Date: Fri, 14 Jun 2024 08:53:32 -0400
Subject: [PATCH] fix: Add a plugin to fix the last column in a table (#7036)

---
 shared/editor/lib/changedDescendants.ts | 37 ++++++++++++++++
 shared/editor/nodes/Table.ts            |  2 +
 shared/editor/plugins/FixTables.ts      | 58 +++++++++++++++++++++++++
 3 files changed, 97 insertions(+)
 create mode 100644 shared/editor/lib/changedDescendants.ts
 create mode 100644 shared/editor/plugins/FixTables.ts

diff --git a/shared/editor/lib/changedDescendants.ts b/shared/editor/lib/changedDescendants.ts
new file mode 100644
index 000000000..cd0ed6ee6
--- /dev/null
+++ b/shared/editor/lib/changedDescendants.ts
@@ -0,0 +1,37 @@
+import { Node } from "prosemirror-model";
+
+/**
+ * Helper for iterating through the nodes in a document that changed
+ * compared to the given previous document. Useful for avoiding
+ * duplicate work on each transaction.
+ */
+export function changedDescendants(
+  /** The previous node */
+  old: Node,
+  /** The current node */
+  cur: Node,
+  /** The offset of the current node */
+  offset: number,
+  /** The function to call for each changed node */
+  callback: (node: Node, pos: number) => void
+): void {
+  const oldSize = old.childCount,
+    curSize = cur.childCount;
+  outer: for (let i = 0, j = 0; i < curSize; i++) {
+    const child = cur.child(i);
+    for (let scan = j, e = Math.min(oldSize, i + 3); scan < e; scan++) {
+      if (old.child(scan) === child) {
+        j = scan + 1;
+        offset += child.nodeSize;
+        continue outer;
+      }
+    }
+    callback(child, offset);
+    if (j < oldSize && old.child(j).sameMarkup(child)) {
+      changedDescendants(old.child(j), child, offset + 1, callback);
+    } else {
+      child.nodesBetween(0, child.content.size, callback, offset + 1);
+    }
+    offset += child.nodeSize;
+  }
+}
diff --git a/shared/editor/nodes/Table.ts b/shared/editor/nodes/Table.ts
index af9bdcc3b..00283757d 100644
--- a/shared/editor/nodes/Table.ts
+++ b/shared/editor/nodes/Table.ts
@@ -23,6 +23,7 @@ import {
   deleteRowSelection,
 } from "../commands/table";
 import { MarkdownSerializerState } from "../lib/markdown/serializer";
+import { FixTablesPlugin } from "../plugins/FixTables";
 import tablesRule from "../rules/tables";
 import { EditorStyleHelper } from "../styles/EditorStyleHelper";
 import { TableLayout } from "../types";
@@ -112,6 +113,7 @@ export default class Table extends Node {
         lastColumnResizable: false,
       }),
       tableEditing(),
+      new FixTablesPlugin(),
     ];
   }
 }
diff --git a/shared/editor/plugins/FixTables.ts b/shared/editor/plugins/FixTables.ts
new file mode 100644
index 000000000..5e983e12b
--- /dev/null
+++ b/shared/editor/plugins/FixTables.ts
@@ -0,0 +1,58 @@
+import { Node } from "prosemirror-model";
+import { EditorState, Plugin, Transaction } from "prosemirror-state";
+import { TableMap } from "prosemirror-tables";
+import { changedDescendants } from "../lib/changedDescendants";
+import { getCellsInColumn } from "../queries/table";
+
+/**
+ * A ProseMirror plugin that fixes the last column in a table to ensure it fills the remaining width.
+ */
+export class FixTablesPlugin extends Plugin {
+  constructor() {
+    super({
+      appendTransaction: (_transactions, oldState, state) => {
+        let tr: Transaction | undefined;
+        const check = (node: Node) => {
+          if (node.type.spec.tableRole === "table") {
+            tr = this.fixTable(state, node, tr);
+          }
+        };
+        if (!oldState) {
+          state.doc.descendants(check);
+        } else if (oldState.doc !== state.doc) {
+          changedDescendants(oldState.doc, state.doc, 0, check);
+        }
+        return tr;
+      },
+    });
+  }
+
+  private fixTable(
+    state: EditorState,
+    table: Node,
+    tr: Transaction | undefined
+  ): Transaction | undefined {
+    let fixed = false;
+    const map = TableMap.get(table);
+    if (!tr) {
+      tr = state.tr;
+    }
+
+    // If the table has only one column, remove the colwidth attribute on all cells
+    if (map.width === 1) {
+      const cells = getCellsInColumn(0)(state);
+      cells.forEach((pos) => {
+        const node = state.doc.nodeAt(pos);
+        if (node?.attrs.colspan) {
+          fixed = true;
+          tr = tr!.setNodeMarkup(pos, undefined, {
+            ...node?.attrs,
+            colwidth: null,
+          });
+        }
+      });
+    }
+
+    return fixed ? tr : undefined;
+  }
+}