fix(spellcheck): Unwrap nodes with simpler code with less side effects

internal_packages/composer-spellcheck/lib/spellcheck-composer-extension.coffee

@@ -13,7 +13,7 @@ class SpellcheckComposerExtension extends ComposerExtension
     SpellcheckCache[word]
 
   @onContentChanged: ({editor}) =>
-    @walkTree(editor)
+    @update(editor)
 
   @onShowContextMenu: ({editor, event, menu}) =>
     selection = editor.currentSelection()
@@ -39,31 +39,63 @@ class SpellcheckComposerExtension extends ComposerExtension
 
   @applyCorrection: (editor, range, selection, correction) =>
     DOMUtils.Mutating.applyTextInRange(range, selection, correction)
-    @walkTree(editor)
+    @update(editor)
 
   @learnSpelling: (editor, word) =>
     spellchecker.add(word)
     delete SpellcheckCache[word]
-    @walkTree(editor)
+    @update(editor)
 
-  @walkTree: (editor) =>
-    # Remove all existing spellcheck nodes
-    spellingNodes = editor.rootNode.querySelectorAll('spelling')
-    for node in spellingNodes
-      editor.whilePreservingSelection =>
-        DOMUtils.unwrapNode(node)
-
-    # Normalize to make sure words aren't split across text nodes
-    editor.rootNode.normalize()
+  @update: (editor) =>
+    @_unwrapWords(editor)
+    @_wrapMisspelledWords(editor)
 
+  # Creates a shallow copy of a selection object where anchorNode / focusNode
+  # can be changed, and provides it to the callback provided. After the callback
+  # runs, it applies the new selection if `snapshot.modified` has been set.
+  #
+  # Note: This is different from ExposedSelection because the nodes are not cloned.
+  # In the callback functions, we need to check whether the anchor/focus nodes
+  # are INSIDE the nodes we're adjusting.
+  #
+  @_whileApplyingSelectionChanges: (cb) =>
     selection = document.getSelection()
     selectionSnapshot =
       anchorNode: selection.anchorNode
       anchorOffset: selection.anchorOffset
       focusNode: selection.focusNode
       focusOffset: selection.focusOffset
-    selectionImpacted = false
+      modified: false
 
+    cb(selectionSnapshot)
+
+    if selectionSnapshot.modified
+      selection.setBaseAndExtent(selectionSnapshot.anchorNode, selectionSnapshot.anchorOffset, selectionSnapshot.focusNode, selectionSnapshot.focusOffset)
+
+  # Removes all of the <spelling> nodes found in the provided `editor`.
+  # It normalizes the DOM after removing spelling nodes to ensure that words
+  # are not split between text nodes. (ie: doesn, 't => doesn't)
+  @_unwrapWords: (editor) =>
+    @_whileApplyingSelectionChanges (selectionSnapshot) =>
+      spellingNodes = editor.rootNode.querySelectorAll('spelling')
+
+      for node in spellingNodes
+        if selectionSnapshot.anchorNode is node
+          selectionSnapshot.anchorNode = node.firstChild
+        if selectionSnapshot.focusNode is node
+          selectionSnapshot.focusNode = node.firstChild
+
+        selectionSnapshot.modified = true
+        node.parentNode.insertBefore(node.firstChild, node) while (node.firstChild)
+        node.parentNode.removeChild(node)
+
+    editor.rootNode.normalize()
+
+  # Traverses all of the text nodes within the provided `editor`. If it finds a
+  # text node with a misspelled word, it splits it, wraps the misspelled word
+  # with a <spelling> node and updates the selection to account for the change.
+  @_wrapMisspelledWords: (editor) =>
+    @_whileApplyingSelectionChanges (selectionSnapshot) =>
       treeWalker = document.createTreeWalker(editor.rootNode, NodeFilter.SHOW_TEXT)
       nodeList = []
       nodeMisspellingsFound = 0
@@ -72,17 +104,17 @@ class SpellcheckComposerExtension extends ComposerExtension
         nodeList.push(treeWalker.currentNode)
 
       # Note: As a performance optimization, we stop spellchecking after encountering
-    # 10 misspelled words. This keeps the runtime of this method bounded!
+      # 30 misspelled words. This keeps the runtime of this method bounded!
 
       while (node = nodeList.shift())
-      break if nodeMisspellingsFound > 10
+        break if nodeMisspellingsFound > 30
         str = node.textContent
 
         # https://regex101.com/r/bG5yC4/1
         wordRegexp = /(\w[\w'’-]*\w|\w)/g
 
         while ((match = wordRegexp.exec(str)) isnt null)
-        break if nodeMisspellingsFound > 10
+          break if nodeMisspellingsFound > 30
           misspelled = @isMisspelled(match[0])
 
           if misspelled
@@ -105,11 +137,11 @@ class SpellcheckComposerExtension extends ComposerExtension
             for prop in ['anchor', 'focus']
               if selectionSnapshot["#{prop}Node"] is node
                 if selectionSnapshot["#{prop}Offset"] > match.index + match[0].length
-                selectionImpacted = true
+                  selectionSnapshot.modified = true
                   selectionSnapshot["#{prop}Node"] = afterMatchNode
                   selectionSnapshot["#{prop}Offset"] -= match.index + match[0].length
                 else if selectionSnapshot["#{prop}Offset"] > match.index
-                selectionImpacted = true
+                  selectionSnapshot.modified = true
                   selectionSnapshot["#{prop}Node"] = spellingSpan.childNodes[0]
                   selectionSnapshot["#{prop}Offset"] -= match.index
 
@@ -117,9 +149,6 @@ class SpellcheckComposerExtension extends ComposerExtension
             nodeList.unshift(afterMatchNode)
             break
 
-    if selectionImpacted
-      selection.setBaseAndExtent(selectionSnapshot.anchorNode, selectionSnapshot.anchorOffset, selectionSnapshot.focusNode, selectionSnapshot.focusOffset)
-
   @finalizeSessionBeforeSending: ({session}) ->
     body = session.draft().body
     clean = body.replace(/<\/?spelling[^>]*>/g, '')