[lexical] Feature: Add onUpdate function during update with $onUpdate…

… (correct baselline) (#6773)

packages/lexical/flow/Lexical.js.flow

@@ -844,6 +844,7 @@ declare export function $isParagraphNode(
 export type EventHandler = (event: Event, editor: LexicalEditor) => void;
 declare export function $hasUpdateTag(tag: string): boolean;
 declare export function $addUpdateTag(tag: string): void;
+declare export function $onUpdate(updateFn: () => void): void;
 declare export function $getNearestNodeFromDOMNode(
   startingDOM: Node,
 ): LexicalNode | null;

packages/lexical/src/LexicalUtils.ts

@@ -1314,6 +1314,19 @@ export function $addUpdateTag(tag: string): void {
   editor._updateTags.add(tag);
 }
 
+/**
+ * Add a function to run after the current update. This will run after any
+ * `onUpdate` function already supplied to `editor.update()`, as well as any
+ * functions added with previous calls to `$onUpdate`.
+ *
+ * @param updateFn The function to run after the current update.
+ */
+export function $onUpdate(updateFn: () => void): void {
+  errorOnReadOnly();
+  const editor = getActiveEditor();
+  editor._deferred.push(updateFn);
+}
+
 export function $maybeMoveChildrenSelectionToParent(
   parentNode: LexicalNode,
 ): BaseSelection | null {

packages/lexical/src/__tests__/unit/LexicalUtils.test.ts

@@ -23,6 +23,7 @@ import {
 } from 'lexical';
 
 import {
+  $onUpdate,
   emptyFunction,
   generateRandomKey,
   getCachedTypeToNodeMap,
@@ -242,6 +243,43 @@ describe('LexicalUtils tests', () => {
       });
     });
 
+    describe('$onUpdate', () => {
+      test('added fn runs after update, original onUpdate, and prior calls to $onUpdate', () => {
+        const {editor} = testEnv;
+        const runs: string[] = [];
+
+        editor.update(
+          () => {
+            $getRoot().append(
+              $createParagraphNode().append($createTextNode('foo')),
+            );
+            $onUpdate(() => {
+              runs.push('second');
+            });
+            $onUpdate(() => {
+              runs.push('third');
+            });
+          },
+          {
+            onUpdate: () => {
+              runs.push('first');
+            },
+          },
+        );
+
+        // Flush pending updates
+        editor.read(() => {});
+
+        expect(runs).toEqual(['first', 'second', 'third']);
+      });
+
+      test('adding fn throws outside update', () => {
+        expect(() => {
+          $onUpdate(() => {});
+        }).toThrow();
+      });
+    });
+
     test('getCachedTypeToNodeMap', async () => {
       const {editor} = testEnv;
       const paragraphKeys: string[] = [];

packages/lexical/src/index.ts

@@ -172,6 +172,7 @@ export {
   $isRootOrShadowRoot,
   $isTokenOrSegmented,
   $nodesOfType,
+  $onUpdate,
   $selectAll,
   $setCompositionKey,
   $setSelection,