fix(spellcheck): Redo node creation each time, optimize findSimilarNodes

findSimilarNodes was taking 20% of total execution time because my test email was a jenkins error report with thousands of text nodes.

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

@@ -13,11 +13,10 @@ class SpellcheckComposerExtension extends ComposerExtension
     SpellcheckCache[word]
 
   @onContentChanged: ({editor}) =>
-    @walkTree(editor.rootNode)
+    @walkTree(editor)
 
   @onShowContextMenu: ({editor, event, menu}) =>
     selection = editor.currentSelection()
-    editableNode = editor.rootNode
     range = DOMUtils.Mutating.getRangeAtAndSelectWord(selection, 0)
     word = range.toString()
     if @isMisspelled(word)
@@ -26,7 +25,7 @@ class SpellcheckComposerExtension extends ComposerExtension
         corrections.forEach (correction) =>
           menu.append(new MenuItem({
             label: correction,
-            click: @applyCorrection.bind(@, editableNode, range, selection, correction)
+            click: @applyCorrection.bind(@, editor, range, selection, correction)
           }))
       else
         menu.append(new MenuItem({ label: 'No Guesses Found', enabled: false}))
@@ -34,23 +33,29 @@ class SpellcheckComposerExtension extends ComposerExtension
       menu.append(new MenuItem({ type: 'separator' }))
       menu.append(new MenuItem({
         label: 'Learn Spelling',
-        click: @learnSpelling.bind(@, editableNode, word)
+        click: @learnSpelling.bind(@, editor, word)
       }))
       menu.append(new MenuItem({ type: 'separator' }))
 
-  @applyCorrection: (editableNode, range, selection, correction) =>
+  @applyCorrection: (editor, range, selection, correction) =>
     DOMUtils.Mutating.applyTextInRange(range, selection, correction)
-    @walkTree(editableNode)
+    @walkTree(editor)
 
-  @learnSpelling: (editableNode, word) =>
+  @learnSpelling: (editor, word) =>
     spellchecker.add(word)
     delete SpellcheckCache[word]
-    @walkTree(editableNode)
+    @walkTree(editor)
 
-  @walkTree: (editableNode) =>
-    treeWalker = document.createTreeWalker(editableNode, NodeFilter.SHOW_TEXT)
+  @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()
 
-    nodeList = []
     selection = document.getSelection()
     selectionSnapshot =
       anchorNode: selection.anchorNode
@@ -59,62 +64,58 @@ class SpellcheckComposerExtension extends ComposerExtension
       focusOffset: selection.focusOffset
     selectionImpacted = false
 
+    treeWalker = document.createTreeWalker(editor.rootNode, NodeFilter.SHOW_TEXT)
+    nodeList = []
+    nodeMisspellingsFound = 0
+
     while (treeWalker.nextNode())
       nodeList.push(treeWalker.currentNode)
 
-    while (node = nodeList.pop())
+    # Note: As a performance optimization, we stop spellchecking after encountering
+    # 10 misspelled words. This keeps the runtime of this method bounded!
+
+    while (node = nodeList.shift())
+      break if nodeMisspellingsFound > 10
       str = node.textContent
 
       # https://regex101.com/r/bG5yC4/1
       wordRegexp = /(\w[\w'’-]*\w|\w)/g
 
       while ((match = wordRegexp.exec(str)) isnt null)
-        spellingSpan = null
-        if node.parentNode and node.parentNode.nodeName is 'SPELLING'
-          if match[0] is str
-            spellingSpan = node.parentNode
-          else
-            node.parentNode.classList.remove('misspelled')
-
+        break if nodeMisspellingsFound > 10
         misspelled = @isMisspelled(match[0])
-        markedAsMisspelled = spellingSpan?.classList.contains('misspelled')
 
-        if misspelled and not markedAsMisspelled
+        if misspelled
           # The insertion point is currently at the end of this misspelled word.
           # Do not mark it until the user types a space or leaves.
           if selectionSnapshot.focusNode is node and selectionSnapshot.focusOffset is match.index + match[0].length
             continue
 
-          if spellingSpan
-            spellingSpan.classList.add('misspelled')
+          if match.index is 0
+            matchNode = node
           else
-            if match.index is 0
-              matchNode = node
-            else
-              matchNode = node.splitText(match.index)
-            afterMatchNode = matchNode.splitText(match[0].length)
+            matchNode = node.splitText(match.index)
+          afterMatchNode = matchNode.splitText(match[0].length)
 
-            spellingSpan = document.createElement('spelling')
-            spellingSpan.classList.add('misspelled')
-            spellingSpan.innerText = match[0]
-            matchNode.parentNode.replaceChild(spellingSpan, matchNode)
+          spellingSpan = document.createElement('spelling')
+          spellingSpan.classList.add('misspelled')
+          spellingSpan.innerText = match[0]
+          matchNode.parentNode.replaceChild(spellingSpan, matchNode)
 
-            for prop in ['anchor', 'focus']
-              if selectionSnapshot["#{prop}Node"] is node
-                if selectionSnapshot["#{prop}Offset"] > match.index + match[0].length
-                  selectionImpacted = true
-                  selectionSnapshot["#{prop}Node"] = afterMatchNode
-                  selectionSnapshot["#{prop}Offset"] -= match.index + match[0].length
-                else if selectionSnapshot["#{prop}Offset"] > match.index
-                  selectionImpacted = true
-                  selectionSnapshot["#{prop}Node"] = spellingSpan.childNodes[0]
-                  selectionSnapshot["#{prop}Offset"] -= match.index
+          for prop in ['anchor', 'focus']
+            if selectionSnapshot["#{prop}Node"] is node
+              if selectionSnapshot["#{prop}Offset"] > match.index + match[0].length
+                selectionImpacted = true
+                selectionSnapshot["#{prop}Node"] = afterMatchNode
+                selectionSnapshot["#{prop}Offset"] -= match.index + match[0].length
+              else if selectionSnapshot["#{prop}Offset"] > match.index
+                selectionImpacted = true
+                selectionSnapshot["#{prop}Node"] = spellingSpan.childNodes[0]
+                selectionSnapshot["#{prop}Offset"] -= match.index
 
-            nodeList.unshift(afterMatchNode)
-            break
-
-        else if not misspelled and markedAsMisspelled
-          spellingSpan.classList.remove('misspelled')
+          nodeMisspellingsFound += 1
+          nodeList.unshift(afterMatchNode)
+          break
 
     if selectionImpacted
       selection.setBaseAndExtent(selectionSnapshot.anchorNode, selectionSnapshot.anchorOffset, selectionSnapshot.focusNode, selectionSnapshot.focusOffset)

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

@@ -16,7 +16,12 @@ describe "SpellcheckComposerExtension", ->
     it "correctly walks a DOM tree and surrounds mispelled words", ->
       dom = document.createElement('div')
       dom.innerHTML = initialHTML
-      SpellcheckComposerExtension.walkTree(dom)
+
+      editor =
+        rootNode: dom
+        whilePreservingSelection: (cb) -> cb()
+
+      SpellcheckComposerExtension.walkTree(editor)
       expect(dom.innerHTML).toEqual(expectedHTML)
 
   describe "finalizeSessionBeforeSending", ->

src/components/contenteditable/contenteditable.cjsx

@@ -413,8 +413,8 @@ class Contenteditable extends React.Component
   # we currently have rendered. Every time React re-renders the component
   # we get new DOM objects. When the `exportedSelection` is re-imported
   # during `_restoreSelection`, the `ExtendedSelection` class will attempt
-  # to find the appropriate DOM Nodes via the `DOMUtils.findSimilarNodes`
-  # method.
+  # to find the appropriate DOM Nodes via the `similar nodes` conveience methods
+  # in DOMUtils.
   #
   # When React re-renders it doesn't restore the Selection. We need to do
   # this manually with `_restoreSelection`

src/components/contenteditable/extended-selection.coffee

@@ -72,9 +72,9 @@ class ExtendedSelection
   # focus node aka extent node).
   importSelection: (exportedSelection) ->
     return unless exportedSelection instanceof ExportedSelection
-    newAnchorNode = DOMUtils.findSimilarNodes(@scopeNode, exportedSelection.anchorNode)[exportedSelection.anchorNodeIndex]
+    newAnchorNode = DOMUtils.findSimilarNodeAtIndex(@scopeNode, exportedSelection.anchorNode, exportedSelection.anchorNodeIndex)
 
-    newFocusNode = DOMUtils.findSimilarNodes(@scopeNode, exportedSelection.focusNode)[exportedSelection.focusNodeIndex]
+    newFocusNode = DOMUtils.findSimilarNodeAtIndex(@scopeNode, exportedSelection.focusNode, exportedSelection.focusNodeIndex)
 
     @setBaseAndExtent(newAnchorNode,
                       exportedSelection.anchorOffset,

src/components/contenteditable/floating-toolbar-container.cjsx

@@ -96,7 +96,7 @@ class FloatingToolbarContainer extends React.Component
   _onSaveUrl: (url, linkToModify) =>
     @props.atomicEdit ({editor}) ->
       if linkToModify?
-        equivalentNode = DOMUtils.findSimilarNodes(editor.rootNode, linkToModify)?[0]
+        equivalentNode = DOMUtils.findSimilarNodeAtIndex(editor.rootNode, linkToModify, 0)
         return unless equivalentNode?
         equivalentLinkText = DOMUtils.findFirstTextNode(equivalentNode)
         return if linkToModify.getAttribute?('href')?.trim() is url.trim()

src/dom-utils.coffee

@@ -216,7 +216,7 @@ DOMUtils =
     return adjacent
 
   getNodeIndex: (context, nodeToFind) =>
-    DOMUtils.findSimilarNodes(context, nodeToFind).indexOf nodeToFind
+    DOMUtils.indexOfNodeInSimilarNodes(context, nodeToFind)
 
   getRangeInScope: (scope) =>
     selection = document.getSelection()
@@ -332,17 +332,35 @@ DOMUtils =
     # \u00a0 is  
     node.data.replace(/\u00a0/g, "x").trim().length is 0
 
-  findSimilarNodes: (context, nodeToFind) ->
-    nodeList = []
+  indexOfNodeInSimilarNodes: (context, nodeToFind) ->
     if nodeToFind.isEqualNode(context)
-      nodeList.push(context)
-      return nodeList
+      return 0
+
     treeWalker = document.createTreeWalker context
+    idx = 0
     while treeWalker.nextNode()
       if treeWalker.currentNode.isEqualNode nodeToFind
-        nodeList.push(treeWalker.currentNode)
+        if treeWalker.currentNode.isSameNode nodeToFind
+          return idx
+        idx += 1
 
-    return nodeList
+    return -1
+
+  # This is an optimization of findSimilarNodes which avoids tons of extra work
+  # scanning a large DOM if all we're going to do is get item at index [0]. It
+  # returns once it has found the similar node at the index desired.
+  findSimilarNodeAtIndex: (context, nodeToFind, desiredIdx) ->
+    if desiredIdx is 0 and nodeToFind.isEqualNode(context)
+      return context
+
+    treeWalker = document.createTreeWalker context
+    idx = 0
+    while treeWalker.nextNode()
+      if treeWalker.currentNode.isEqualNode nodeToFind
+        return treeWalker.currentNode if desiredIdx is idx
+        idx += 1
+
+    return null
 
   findCharacter: (context, character) ->
     node = null