Render correct classes on leading/trailing whitespace spans

spec/editor-component-spec.coffee

@@ -221,11 +221,15 @@ describe "EditorComponent", ->
         atom.config.set("editor.showInvisibles", true)
         expect(component.lineNodeForScreenRow(0).textContent).toBe "#{invisibles.space}a line with tabs#{invisibles.tab}and spaces#{invisibles.space}#{invisibles.eol}"
 
-      it "displays spaces, tabs, and newlines as visible characters", ->
+      it "displays leading/trailing spaces, tabs, and newlines as visible characters", ->
         editor.setText " a line with tabs\tand spaces "
         nextAnimationFrame()
         expect(component.lineNodeForScreenRow(0).textContent).toBe "#{invisibles.space}a line with tabs#{invisibles.tab}and spaces#{invisibles.space}#{invisibles.eol}"
 
+        leafNodes = getLeafNodes(component.lineNodeForScreenRow(0))
+        expect(leafNodes[0].classList.contains('invisible-character')).toBe true
+        expect(leafNodes[leafNodes.length - 1].classList.contains('invisible-character')).toBe true
+
       it "displays newlines as their own token outside of the other tokens' scopes", ->
         editor.setText "var"
         nextAnimationFrame()

src/token.coffee

@@ -2,11 +2,7 @@ _ = require 'underscore-plus'
 textUtils = require './text-utils'
 
 WhitespaceRegexesByTabLength = {}
-LeadingSpaceRegex = /^[ ]+/
-TrailingSpaceRegex = /[ ]+$/
 EscapeRegex = /[&"'<>]/g
-CharacterRegex = /./g
-StartCharacterRegex = /^./
 StartDotRegex = /^\.?/
 WhitespaceRegex = /\S/
 
@@ -22,6 +18,8 @@ class Token
   isHardTab: null
   hasLeadingWhitespace: false
   hasTrailingWhitespace: false
+  firstNonWhitespaceIndex: null
+  firstTrailingWhitespaceIndex: null
 
   constructor: ({@value, @scopes, @isAtomic, @bufferDelta, @isHardTab}) ->
     @screenDelta = @value.length
@@ -150,24 +148,27 @@ class Token
       leadingHtml = ''
       trailingHtml = ''
 
-      if @hasLeadingWhitespace and match = LeadingSpaceRegex.exec(@value)
+      if @hasLeadingWhitespace
+        leadingWhitespace = @value.substring(0, @firstNonWhitespaceIndex)
+
         classes = 'leading-whitespace'
         classes += ' indent-guide' if hasIndentGuide
         classes += ' invisible-character' if invisibles.space
 
-        leadingHtml = "<span class='#{classes}'>#{match[0]}</span>"
+        leadingHtml = "<span class='#{classes}'>#{leadingWhitespace}</span>"
+        startIndex = @firstNonWhitespaceIndex
 
-        startIndex = match[0].length
+      if @hasTrailingWhitespace
+        tokenIsOnlyWhitespace = @firstTrailingWhitespaceIndex is 0
+        trailingWhitespace = @value.substring(@firstTrailingWhitespaceIndex)
 
-      if @hasTrailingWhitespace and match = TrailingSpaceRegex.exec(@value)
-        tokenIsOnlyWhitespace = match[0].length is @value.length
         classes = 'trailing-whitespace'
         classes += ' indent-guide' if hasIndentGuide and not @hasLeadingWhitespace and tokenIsOnlyWhitespace
         classes += ' invisible-character' if invisibles.space
 
-        trailingHtml = "<span class='#{classes}'>#{match[0]}</span>"
+        trailingHtml = "<span class='#{classes}'>#{trailingWhitespace}</span>"
 
-        endIndex = match.index
+        endIndex = @firstTrailingWhitespaceIndex
 
       html = leadingHtml
       if @value.length > MaxTokenLength

src/tokenized-line.coffee

@@ -138,15 +138,17 @@ class TokenizedLine
     outputTokens
 
   markLeadingAndTrailingWhitespaceTokens: ->
-    firstNonWhitespacePosition = @text.search(NonWhitespaceRegex)
-    firstTrailingWhitespacePosition = @text.search(TrailingWhitespaceRegex)
-    lineIsWhitespaceOnly = firstTrailingWhitespacePosition is 0
-    position = 0
-    for token, i in @tokens
-      token.hasLeadingWhitespace =  position < firstNonWhitespacePosition
+    firstNonWhitespaceIndex = @text.search(NonWhitespaceRegex)
+    firstTrailingWhitespaceIndex = @text.search(TrailingWhitespaceRegex)
+    lineIsWhitespaceOnly = firstTrailingWhitespaceIndex is 0
+    index = 0
+    for token in @tokens
+      if token.hasLeadingWhitespace =  index < firstNonWhitespaceIndex
+        token.firstNonWhitespaceIndex = firstNonWhitespaceIndex - index
       # Only the *last* segment of a soft-wrapped line can have trailing whitespace
-      token.hasTrailingWhitespace = @lineEnding? and (position + token.value.length > firstTrailingWhitespacePosition)
-      position += token.value.length
+      if token.hasTrailingWhitespace = @lineEnding? and (index + token.value.length > firstTrailingWhitespaceIndex)
+        token.firstTrailingWhitespaceIndex = Math.max(0, firstTrailingWhitespaceIndex - index)
+      index += token.value.length
 
   substituteInvisibleCharacters: ->
     invisibles = @invisibles