Update android restoreDOM to allow partial dom restoring (#4706)

* Update android restoreDOM to allow partial dom restoring

* Fix useContentKey re-render timing

.changeset/blue-chefs-chew.md

@@ -0,0 +1,5 @@
+---
+'slate-react': patch
+---
+
+Update android restoreDOM to use partial dom restoring

packages/slate-react/src/components/android/android-editable.tsx

@@ -14,11 +14,9 @@ import {
   isDOMNode,
   getDefaultView,
   getClipboardData,
-  isPlainTextOnlyPaste,
 } from '../../utils/dom'
 import {
   EDITOR_TO_ELEMENT,
-  EDITOR_TO_RESTORE_DOM,
   EDITOR_TO_WINDOW,
   ELEMENT_TO_NODE,
   IS_FOCUSED,
@@ -37,6 +35,7 @@ import {
 } from '../editable'
 
 import { useAndroidInputManager } from './use-android-input-manager'
+import { useContentKey } from '../../hooks/use-content-key'
 
 /**
  * Editable.
@@ -72,10 +71,7 @@ export const AndroidEditable = (props: EditableProps): JSX.Element => {
     []
   )
 
-  const [contentKey, setContentKey] = useState(0)
-  const onRestoreDOM = useCallback(() => {
-    setContentKey(prev => prev + 1)
-  }, [contentKey])
+  const contentKey = useContentKey(editor)
 
   // Whenever the editor updates...
   useIsomorphicLayoutEffect(() => {
@@ -87,10 +83,8 @@ export const AndroidEditable = (props: EditableProps): JSX.Element => {
       EDITOR_TO_ELEMENT.set(editor, ref.current)
       NODE_TO_ELEMENT.set(editor, ref.current)
       ELEMENT_TO_NODE.set(ref.current, editor)
-      EDITOR_TO_RESTORE_DOM.set(editor, onRestoreDOM)
     } else {
       NODE_TO_ELEMENT.delete(editor)
-      EDITOR_TO_RESTORE_DOM.delete(editor)
     }
 
     try {

packages/slate-react/src/components/android/android-input-manager.ts

@@ -16,7 +16,6 @@ import {
   isReplaceExpandedSelection,
   isTextInsertion,
 } from './mutation-detection'
-import { restoreDOM } from './restore-dom'
 
 // Replace with `const debug = console.log` to debug
 const debug = (...message: any[]) => {}
@@ -42,11 +41,13 @@ const debug = (...message: any[]) => {}
  * - Line breaks
  *
  * @param editor
+ * @param restoreDOM
  */
 
 export class AndroidInputManager {
-  constructor(private editor: ReactEditor) {
+  constructor(private editor: ReactEditor, private restoreDOM: () => void) {
     this.editor = editor
+    this.restoreDOM = restoreDOM
   }
 
   /**
@@ -65,7 +66,7 @@ export class AndroidInputManager {
       console.error(err)
 
       // Failed to reconcile mutations, restore DOM to its previous state
-      restoreDOM(this.editor)
+      this.restoreDOM()
     }
   }
 
@@ -135,9 +136,7 @@ export class AndroidInputManager {
 
     Editor.insertBreak(this.editor)
 
-    // To-do: Need a more granular solution to restoring only a specific portion
-    // of the document. Restoring the entire document is expensive.
-    restoreDOM(this.editor)
+    this.restoreDOM()
 
     if (selection) {
       // Compat: Move selection to the newly inserted block if it has not moved
@@ -167,7 +166,7 @@ export class AndroidInputManager {
       Editor.insertText(this.editor, text)
     }
 
-    restoreDOM(this.editor)
+    this.restoreDOM()
   }
 
   /**
@@ -180,7 +179,7 @@ export class AndroidInputManager {
     Editor.deleteBackward(this.editor)
     ReactEditor.focus(this.editor)
 
-    restoreDOM(this.editor)
+    this.restoreDOM()
   }
 
   /**
@@ -194,7 +193,7 @@ export class AndroidInputManager {
         const path = ReactEditor.findPath(this.editor, slateNode)
 
         Transforms.delete(this.editor, { at: path })
-        restoreDOM(this.editor)
+        this.restoreDOM()
       }
     }
   }

packages/slate-react/src/components/android/restore-dom.ts

@@ -1,14 +0,0 @@
-import { ReactEditor } from '../..'
-import { EDITOR_TO_RESTORE_DOM } from '../../utils/weak-maps'
-
-export function restoreDOM(editor: ReactEditor) {
-  try {
-    const onRestoreDOM = EDITOR_TO_RESTORE_DOM.get(editor)
-    if (onRestoreDOM) {
-      onRestoreDOM()
-    }
-  } catch (err) {
-    // eslint-disable-next-line no-console
-    console.error(err)
-  }
-}

packages/slate-react/src/components/android/use-android-input-manager.ts

@@ -1,8 +1,9 @@
-import { RefObject, useCallback, useRef, useState } from 'react'
+import { RefObject, useCallback, useMemo, useRef, useState } from 'react'
 
 import { useSlateStatic } from '../../hooks/use-slate-static'
 
 import { AndroidInputManager } from './android-input-manager'
+import { useRestoreDom } from './use-restore-dom'
 import { useMutationObserver } from './use-mutation-observer'
 import { useTrackUserInput } from './use-track-user-input'
 
@@ -15,8 +16,15 @@ const MUTATION_OBSERVER_CONFIG: MutationObserverInit = {
 
 export function useAndroidInputManager(node: RefObject<HTMLElement>) {
   const editor = useSlateStatic()
-  const [inputManager] = useState(() => new AndroidInputManager(editor))
+
   const { receivedUserInput, onUserInput } = useTrackUserInput()
+  const restoreDom = useRestoreDom(node, receivedUserInput)
+
+  const inputManager = useMemo(
+    () => new AndroidInputManager(editor, restoreDom),
+    [restoreDom, editor]
+  )
+
   const timeoutId = useRef<ReturnType<typeof setTimeout> | null>(null)
   const isReconciling = useRef(false)
   const flush = useCallback((mutations: MutationRecord[]) => {

packages/slate-react/src/components/android/use-restore-dom.tsx

@@ -0,0 +1,81 @@
+import React, { useCallback, useEffect, useRef } from 'react'
+import { Node as SlateNode, Path } from 'slate'
+import { ReactEditor, useSlateStatic } from '../..'
+import { DOMNode, isDOMElement } from '../../utils/dom'
+import { ELEMENT_TO_NODE, NODE_TO_RESTORE_DOM } from '../../utils/weak-maps'
+import { useMutationObserver } from './use-mutation-observer'
+
+const MUTATION_OBSERVER_CONFIG: MutationObserverInit = {
+  childList: true,
+  characterData: true,
+  subtree: true,
+}
+
+function findClosestKnowSlateNode(domNode: DOMNode): SlateNode | null {
+  let domEl = isDOMElement(domNode) ? domNode : domNode.parentElement
+
+  if (domEl && !domEl.hasAttribute('data-slate-node')) {
+    domEl = domEl.closest(`[data-slate-node]`)
+  }
+
+  const slateNode = domEl && ELEMENT_TO_NODE.get(domEl as HTMLElement)
+  if (slateNode) {
+    return slateNode
+  }
+
+  // Unknown dom element with a slate-slate-node attribute => the IME
+  // most likely duplicated the node so we have to restore the parent
+  return domEl?.parentElement
+    ? findClosestKnowSlateNode(domEl.parentElement)
+    : null
+}
+
+export function useRestoreDom(
+  node: React.RefObject<HTMLElement>,
+  receivedUserInput: React.RefObject<boolean>
+) {
+  const editor = useSlateStatic()
+  const mutatedNodes = useRef<Set<SlateNode>>(new Set())
+
+  const handleDOMMutation = useCallback((mutations: MutationRecord[]) => {
+    if (!receivedUserInput.current) {
+      return
+    }
+
+    mutations.forEach(({ target }) => {
+      const slateNode = findClosestKnowSlateNode(target)
+      if (!slateNode) {
+        return
+      }
+
+      return mutatedNodes.current.add(slateNode)
+    })
+  }, [])
+
+  useMutationObserver(node, handleDOMMutation, MUTATION_OBSERVER_CONFIG)
+
+  // Clear mutated nodes on every render
+  mutatedNodes.current.clear()
+  const restore = useCallback(() => {
+    const mutated = Array.from(mutatedNodes.current.values())
+
+    // Filter out child nodes of nodes that will be restored anyway
+    const nodesToRestore = mutated.filter(
+      n =>
+        !mutated.some(m =>
+          Path.isParent(
+            ReactEditor.findPath(editor, m),
+            ReactEditor.findPath(editor, n)
+          )
+        )
+    )
+
+    nodesToRestore.forEach(n => {
+      NODE_TO_RESTORE_DOM.get(n)?.()
+    })
+
+    mutatedNodes.current.clear()
+  }, [])
+
+  return restore
+}

packages/slate-react/src/components/element.tsx

@@ -1,6 +1,6 @@
-import React, { useRef } from 'react'
+import React, { Fragment, useRef } from 'react'
 import getDirection from 'direction'
-import { Editor, Node, Range, NodeEntry, Element as SlateElement } from 'slate'
+import { Editor, Node, Range, Element as SlateElement } from 'slate'
 
 import Text from './text'
 import useChildren from '../hooks/use-children'
@@ -19,6 +19,8 @@ import {
   RenderLeafProps,
   RenderPlaceholderProps,
 } from './editable'
+import { useContentKey } from '../hooks/use-content-key'
+import { IS_ANDROID } from '../utils/environment'
 
 /**
  * Element.
@@ -131,7 +133,14 @@ const Element = (props: {
     }
   })
 
-  return renderElement({ attributes, children, element })
+  const content = renderElement({ attributes, children, element })
+
+  if (IS_ANDROID) {
+    const contentKey = useContentKey(element)
+    return <Fragment key={contentKey}>{content}</Fragment>
+  }
+
+  return content
 }
 
 const MemoizedElement = React.memo(Element, (prev, next) => {

packages/slate-react/src/components/text.tsx

@@ -11,6 +11,8 @@ import {
   EDITOR_TO_KEY_TO_ELEMENT,
 } from '../utils/weak-maps'
 import { isDecoratorRangeListEqual } from '../utils/range-list'
+import { useContentKey } from '../hooks/use-content-key'
+import { IS_ANDROID } from '../utils/environment'
 
 /**
  * Text.
@@ -67,8 +69,10 @@ const Text = (props: {
     }
   })
 
+  const contentKey = IS_ANDROID ? useContentKey(text) : undefined
+
   return (
-    <span data-slate-node="text" ref={ref}>
+    <span data-slate-node="text" ref={ref} key={contentKey}>
       {children}
     </span>
   )

packages/slate-react/src/hooks/use-children.tsx

@@ -1,5 +1,5 @@
 import React from 'react'
-import { Editor, Range, Element, NodeEntry, Ancestor, Descendant } from 'slate'
+import { Editor, Range, Element, Ancestor, Descendant } from 'slate'
 
 import ElementComponent from '../components/element'
 import TextComponent from '../components/text'

packages/slate-react/src/hooks/use-content-key.ts

@@ -0,0 +1,38 @@
+import { useEffect, useRef, useState } from 'react'
+import { Node as SlateNode } from 'slate'
+import { NODE_TO_RESTORE_DOM } from '../utils/weak-maps'
+
+export function useContentKey(node: SlateNode) {
+  const contentKeyRef = useRef<number>(0)
+  const updateAnimationFrameRef = useRef<number | null>(null)
+
+  const [, setForceRerenderCounter] = useState(0)
+
+  useEffect(() => {
+    NODE_TO_RESTORE_DOM.set(node, () => {
+      // Update is already queued and node hasn't re-render yet
+      if (updateAnimationFrameRef.current) {
+        return
+      }
+
+      updateAnimationFrameRef.current = requestAnimationFrame(() => {
+        setForceRerenderCounter(state => state + 1)
+        updateAnimationFrameRef.current = null
+      })
+
+      contentKeyRef.current++
+    })
+
+    return () => {
+      NODE_TO_RESTORE_DOM.delete(node)
+    }
+  }, [node])
+
+  // Node was restored => clear scheduled update
+  if (updateAnimationFrameRef.current) {
+    cancelAnimationFrame(updateAnimationFrameRef.current)
+    updateAnimationFrameRef.current = null
+  }
+
+  return contentKeyRef.current
+}

packages/slate-react/src/utils/weak-maps.ts

@@ -1,5 +1,4 @@
-import { Node, Ancestor, Editor, Range } from 'slate'
-
+import { Ancestor, Editor, Node } from 'slate'
 import { Key } from './key'
 
 /**
@@ -40,7 +39,7 @@ export const IS_CLICKING: WeakMap<Editor, boolean> = new WeakMap()
 
 export const EDITOR_TO_ON_CHANGE = new WeakMap<Editor, () => void>()
 
-export const EDITOR_TO_RESTORE_DOM = new WeakMap<Editor, () => void>()
+export const NODE_TO_RESTORE_DOM = new WeakMap<Node, () => void>()
 
 /**
  * Symbols.