Multiple update tags (#6507)

Co-authored-by: Bob Ippolito <[email protected]>
Co-authored-by: Ivaylo Pavlov <[email protected]>

Author: 3 people

packages/lexical/flow/Lexical.js.flow

@@ -215,7 +215,7 @@ type EditorReadOptions = {
 };
 type EditorUpdateOptions = {
   onUpdate?: () => void,
-  tag?: string,
+  tag?: string | Array<string>,
   skipTransforms?: true,
   discrete?: true,
 };

packages/lexical/src/LexicalEditor.ts

@@ -80,7 +80,7 @@ export type TextNodeThemeClasses = {
 export type EditorUpdateOptions = {
   onUpdate?: () => void;
   skipTransforms?: true;
-  tag?: string;
+  tag?: string | Array<string>;
   discrete?: true;
 };
 

packages/lexical/src/LexicalUpdates.ts

@@ -219,6 +219,20 @@ function $normalizeAllDirtyTextNodes(
   }
 }
 
+function addTags(editor: LexicalEditor, tags: undefined | string | string[]) {
+  if (!tags) {
+    return;
+  }
+  const updateTags = editor._updateTags;
+  let tags_ = tags;
+  if (!Array.isArray(tags)) {
+    tags_ = [tags];
+  }
+  for (const tag of tags_) {
+    updateTags.add(tag);
+  }
+}
+
 /**
  * Transform heuristic:
  * 1. We transform leaves first. If transforms generate additional dirty nodes we repeat step 1.
@@ -829,11 +843,9 @@ function processNestedUpdates(
       const [nextUpdateFn, options] = queuedUpdate;
 
       let onUpdate;
-      let tag;
 
       if (options !== undefined) {
         onUpdate = options.onUpdate;
-        tag = options.tag;
 
         if (options.skipTransforms) {
           skipTransforms = true;
@@ -851,9 +863,7 @@ function processNestedUpdates(
           editor._deferred.push(onUpdate);
         }
 
-        if (tag) {
-          editor._updateTags.add(tag);
-        }
+        addTags(editor, options.tag);
       }
 
       nextUpdateFn();
@@ -870,17 +880,12 @@ function $beginUpdate(
 ): void {
   const updateTags = editor._updateTags;
   let onUpdate;
-  let tag;
   let skipTransforms = false;
   let discrete = false;
 
   if (options !== undefined) {
     onUpdate = options.onUpdate;
-    tag = options.tag;
-
-    if (tag != null) {
-      updateTags.add(tag);
-    }
+    addTags(editor, options.tag);
 
     skipTransforms = options.skipTransforms || false;
     discrete = options.discrete || false;

packages/lexical/src/__tests__/unit/LexicalEditor.test.tsx

@@ -61,6 +61,7 @@ import {createRoot, Root} from 'react-dom/client';
 import invariant from 'shared/invariant';
 import * as ReactTestUtils from 'shared/react-test-utils';
 
+import {emptyFunction} from '../../LexicalUtils';
 import {
   $createTestDecoratorNode,
   $createTestElementNode,
@@ -2041,6 +2042,28 @@ describe('LexicalEditor tests', () => {
     ]);
   });
 
+  it('multiple update tags', async () => {
+    init();
+    const $mutateSomething = $createTextNode;
+
+    editor.update($mutateSomething, {
+      tag: ['a', 'b'],
+    });
+    expect(editor._updateTags).toEqual(new Set(['a', 'b']));
+    editor.update(
+      () => {
+        editor.update(emptyFunction, {tag: ['e', 'f']});
+      },
+      {
+        tag: ['c', 'd'],
+      },
+    );
+    expect(editor._updateTags).toEqual(new Set(['a', 'b', 'c', 'd', 'e', 'f']));
+
+    await Promise.resolve();
+    expect(editor._updateTags).toEqual(new Set([]));
+  });
+
   it('mutation listeners does not trigger when other node types are mutated', async () => {
     init();