editor: batch update hidden attribute when toggling headings

Signed-off-by: 01zulfi <85733202+01zulfi@users.noreply.github.com>

packages/editor/src/extensions/heading/batch-attribute-step.ts

@@ -0,0 +1,145 @@
+/*
+This file is part of the Notesnook project (https://notesnook.com/)
+
+Copyright (C) 2023 Streetwriters (Private) Limited
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+import { Step, StepResult } from "@tiptap/pm/transform";
+import { Node, Schema, Fragment } from "@tiptap/pm/model";
+import { Mapping } from "@tiptap/pm/transform";
+
+interface AttributeUpdate {
+  pos: number;
+  attrName: string;
+  value: any;
+}
+
+export class BatchAttributeStep extends Step {
+  constructor(public updates: AttributeUpdate[]) {
+    super();
+  }
+
+  apply(doc: Node): StepResult {
+    const updateMap = new Map<number, Map<string, any>>();
+
+    for (const update of this.updates) {
+      if (!updateMap.has(update.pos)) {
+        updateMap.set(update.pos, new Map());
+      }
+      updateMap.get(update.pos)!.set(update.attrName, update.value);
+    }
+
+    const newDoc = doc.copy(this.updateContent(doc.content, 0, updateMap));
+
+    return StepResult.ok(newDoc);
+  }
+
+  private updateContent(
+    content: Fragment,
+    startPos: number,
+    updateMap: Map<number, Map<string, any>>
+  ): Fragment {
+    const nodes: Node[] = [];
+    let pos = startPos;
+
+    content.forEach((node: Node) => {
+      const nodePos = pos;
+      const attrs = updateMap.get(nodePos);
+
+      if (attrs) {
+        const newAttrs = { ...node.attrs };
+        for (const [attrName, value] of attrs.entries()) {
+          newAttrs[attrName] = value;
+        }
+
+        const newContent = node.isLeaf
+          ? node.content
+          : this.updateContent(node.content, pos + 1, updateMap);
+
+        nodes.push(node.type.create(newAttrs, newContent, node.marks));
+      } else {
+        if (!node.isLeaf && node.content.size > 0) {
+          const newContent = this.updateContent(
+            node.content,
+            pos + 1,
+            updateMap
+          );
+          nodes.push(node.copy(newContent));
+        } else {
+          nodes.push(node);
+        }
+      }
+
+      pos += node.nodeSize;
+    });
+
+    return Fragment.from(nodes);
+  }
+
+  invert(doc: Node): Step {
+    const inverseUpdates: AttributeUpdate[] = [];
+
+    for (const update of this.updates) {
+      const node = doc.nodeAt(update.pos);
+      if (node) {
+        inverseUpdates.push({
+          pos: update.pos,
+          attrName: update.attrName,
+          value: node.attrs[update.attrName]
+        });
+      }
+    }
+
+    return new BatchAttributeStep(inverseUpdates);
+  }
+
+  map(mapping: Mapping): Step | null {
+    const mappedUpdates: AttributeUpdate[] = [];
+
+    for (const update of this.updates) {
+      const mappedPos = mapping.map(update.pos);
+
+      if (mappedPos === null || mappedPos === undefined) {
+        continue;
+      }
+
+      mappedUpdates.push({
+        pos: mappedPos,
+        attrName: update.attrName,
+        value: update.value
+      });
+    }
+
+    if (mappedUpdates.length === 0) {
+      return null;
+    }
+
+    return new BatchAttributeStep(mappedUpdates);
+  }
+
+  toJSON() {
+    return {
+      stepType: "batchAttribute",
+      updates: this.updates
+    };
+  }
+
+  static fromJSON(_: Schema, json: any): BatchAttributeStep {
+    return new BatchAttributeStep(json.updates);
+  }
+}
+
+Step.jsonID("batchAttribute", BatchAttributeStep);

packages/editor/src/extensions/heading/heading.ts

@@ -27,6 +27,7 @@ import { Node } from "@tiptap/pm/model";
 import { Plugin, PluginKey, Selection, Transaction } from "@tiptap/pm/state";
 import { Callout } from "../callout/callout.js";
 import { changedDescendants } from "../../utils/prosemirror.js";
+import { BatchAttributeStep } from "./batch-attribute-step.js";
 
 const COLLAPSIBLE_BLOCK_TYPES = [
   "paragraph",
@@ -305,6 +306,8 @@ function toggleNodesUnderPos(
   headingLevel: number,
   isCollapsing: boolean
 ) {
+  console.time("heading-toggle - total");
+  console.time("heading-toggle - setup");
   const { doc } = tr;
   const node = doc.nodeAt(pos);
   if (!node) return;
@@ -315,6 +318,11 @@ function toggleNodesUnderPos(
   let insideCollapsedHeading = false;
   let nestedHeadingLevel: number | null = null;
 
+  const updates: { pos: number; attrName: string; value: any }[] = [];
+
+  console.timeEnd("heading-toggle - setup");
+  console.time("heading-toggle - traverse and collect");
+
   while (nextPos < doc.content.size) {
     const nextNode = doc.nodeAt(nextPos);
     if (!nextNode) break;
@@ -339,7 +347,7 @@ function toggleNodesUnderPos(
 
     if (COLLAPSIBLE_BLOCK_TYPES.includes(nextNode.type.name)) {
       if (isCollapsing) {
-        tr.setNodeAttribute(currentPos, "hidden", true);
+        updates.push({ pos: currentPos, attrName: "hidden", value: true });
       } else {
         if (insideCollapsedHeading) {
           if (
@@ -354,7 +362,7 @@ function toggleNodesUnderPos(
           }
         }
 
-        tr.setNodeAttribute(currentPos, "hidden", false);
+        updates.push({ pos: currentPos, attrName: "hidden", value: false });
         if (nextNode.type.name === "heading" && nextNode.attrs.collapsed) {
           insideCollapsedHeading = true;
           nestedHeadingLevel = nextNode.attrs.level;
@@ -363,10 +371,23 @@ function toggleNodesUnderPos(
     }
   }
 
+  console.timeEnd("heading-toggle - traverse and collect");
+  console.log(`heading-toggle - collected ${updates.length} updates`);
+
+  if (updates.length > 0) {
+    console.time("heading-toggle - apply batch step");
+    tr.step(new BatchAttributeStep(updates));
+    console.timeEnd("heading-toggle - apply batch step");
+  }
+
   if (shouldMoveCursor) {
+    console.time("heading-toggle - move cursor");
     const endPos = pos + node.nodeSize - 1;
     tr.setSelection(Selection.near(tr.doc.resolve(endPos)));
+    console.timeEnd("heading-toggle - move cursor");
   }
+
+  console.timeEnd("heading-toggle - total");
 }
 
 function findEndOfCollapsedSection(