fix: remove memory leaks in Text and Editable

ianstormtaylor · Feb 20, 2023 · 22c4e40 · 22c4e40
1 parent be564ce
commit 22c4e40
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 17 deletions.
diff --git a/packages/slate-react/src/components/editable.tsx b/packages/slate-react/src/components/editable.tsx
@@ -147,7 +147,7 @@ export const Editable = (props: EditableProps) => {
   const editor = useSlate()
   // Rerender editor when composition status changed
   const [isComposing, setIsComposing] = useState(false)
-  const ref = useRef<HTMLDivElement>(null)
+  const ref = useRef<HTMLDivElement | null>(null)
   const deferredOperations = useRef<DeferredOperation[]>([])
 
   const { onUserInput, receivedUserInput } = useTrackUserInput()
@@ -698,6 +698,23 @@ export const Editable = (props: EditableProps) => {
     [readOnly, propsOnDOMBeforeInput]
   )
 
+  const callbackRef = useCallback(
+    node => {
+      if (node == null) {
+        EDITOR_TO_ELEMENT.delete(editor)
+        NODE_TO_ELEMENT.delete(editor)
+
+        if (HAS_BEFORE_INPUT_SUPPORT) {
+          // @ts-ignore The `beforeinput` event isn't recognized.
+          ref.current.removeEventListener('beforeinput', onDOMBeforeInput)
+        }
+      }
+
+      ref.current = node
+    },
+    [ref, onDOMBeforeInput]
+  )
+
   // Attach a native DOM event handler for `beforeinput` events, because React's
   // built-in `onBeforeInput` is actually a leaky polyfill that doesn't expose
   // real `beforeinput` events sadly... (2019/11/04)
@@ -845,7 +862,7 @@ export const Editable = (props: EditableProps) => {
             // this magic zIndex="-1" will fix it
             zindex={-1}
             suppressContentEditableWarning
-            ref={ref}
+            ref={callbackRef}
             style={{
               ...(disableDefaultStyles
                 ? {}

diff --git a/packages/slate-react/src/components/text.tsx b/packages/slate-react/src/components/text.tsx
@@ -1,4 +1,4 @@
-import React, { useRef } from 'react'
+import React, { useRef, useCallback } from 'react'
 import { Element, Range, Text as SlateText } from 'slate'
 import { ReactEditor, useSlateStatic } from '..'
 import { useIsomorphicLayoutEffect } from '../hooks/use-isomorphic-layout-effect'
@@ -32,7 +32,7 @@ const Text = (props: {
     text,
   } = props
   const editor = useSlateStatic()
-  const ref = useRef<HTMLSpanElement>(null)
+  const ref = useRef<HTMLSpanElement | null>(null)
   const leaves = SlateText.decorations(text, decorations)
   const key = ReactEditor.findKey(editor, text)
   const children = []
@@ -54,20 +54,26 @@ const Text = (props: {
   }
 
   // Update element-related weak maps with the DOM element ref.
-  useIsomorphicLayoutEffect(() => {
-    const KEY_TO_ELEMENT = EDITOR_TO_KEY_TO_ELEMENT.get(editor)
-    if (ref.current) {
-      KEY_TO_ELEMENT?.set(key, ref.current)
-      NODE_TO_ELEMENT.set(text, ref.current)
-      ELEMENT_TO_NODE.set(ref.current, text)
-    } else {
-      KEY_TO_ELEMENT?.delete(key)
-      NODE_TO_ELEMENT.delete(text)
-    }
-  })
-
+  const callbackRef = useCallback(
+    (span: HTMLSpanElement | null) => {
+      const KEY_TO_ELEMENT = EDITOR_TO_KEY_TO_ELEMENT.get(editor)
+      if (span) {
+        KEY_TO_ELEMENT?.set(key, span)
+        NODE_TO_ELEMENT.set(text, span)
+        ELEMENT_TO_NODE.set(span, text)
+      } else {
+        KEY_TO_ELEMENT?.delete(key)
+        NODE_TO_ELEMENT.delete(text)
+        if (ref.current) {
+          ELEMENT_TO_NODE.delete(ref.current)
+        }
+      }
+      ref.current = span
+    },
+    [ref, editor, key, text]
+  )
   return (
-    <span data-slate-node="text" ref={ref}>
+    <span data-slate-node="text" ref={callbackRef}>
       {children}
     </span>
   )