From e09ee1c1fa8ac10dbec10a0fc47253bb51c7bbd7 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Fri, 15 Dec 2017 09:44:45 -0800
Subject: [PATCH] Fix error in TreeSitterHighlightIterator.seek

---
 spec/tree-sitter-language-mode-spec.js | 63 ++++++++++++++++++++------
 src/tree-sitter-language-mode.js       |  2 +-
 2 files changed, 51 insertions(+), 14 deletions(-)

diff --git a/spec/tree-sitter-language-mode-spec.js b/spec/tree-sitter-language-mode-spec.js
index 5ecc73308..91070710f 100644
--- a/spec/tree-sitter-language-mode-spec.js
+++ b/spec/tree-sitter-language-mode-spec.js
@@ -32,7 +32,7 @@ describe('TreeSitterLanguageMode', () => {
 
       buffer.setLanguageMode(new TreeSitterLanguageMode({buffer, grammar}))
       buffer.setText('aa.bbb = cc(d.eee());')
-      expectTokensToEqual(editor, [
+      expectTokensToEqual(editor, [[
         {text: 'aa.', scopes: ['source']},
         {text: 'bbb', scopes: ['source', 'property']},
         {text: ' = ', scopes: ['source']},
@@ -40,7 +40,7 @@ describe('TreeSitterLanguageMode', () => {
         {text: '(d.', scopes: ['source']},
         {text: 'eee', scopes: ['source', 'method']},
         {text: '());', scopes: ['source']}
-      ])
+      ]])
     })
 
     it('can start or end multiple scopes at the same position', () => {
@@ -58,7 +58,7 @@ describe('TreeSitterLanguageMode', () => {
 
       buffer.setLanguageMode(new TreeSitterLanguageMode({buffer, grammar}))
       buffer.setText('a = bb.ccc();')
-      expectTokensToEqual(editor, [
+      expectTokensToEqual(editor, [[
         {text: 'a', scopes: ['source', 'variable']},
         {text: ' = ', scopes: ['source']},
         {text: 'bb', scopes: ['source', 'call', 'member', 'variable']},
@@ -66,6 +66,31 @@ describe('TreeSitterLanguageMode', () => {
         {text: '(', scopes: ['source', 'call', 'open-paren']},
         {text: ')', scopes: ['source', 'call', 'close-paren']},
         {text: ';', scopes: ['source']}
+      ]])
+    })
+
+    it('can resume highlighting on a line that starts with whitespace', () => {
+      const grammar = new TreeSitterGrammar(atom.grammars, jsGrammarPath, {
+        parser: 'tree-sitter-javascript',
+        scopes: {
+          'call_expression > member_expression > property_identifier': 'function',
+          'property_identifier': 'member',
+          'identifier': 'variable'
+        }
+      })
+
+      buffer.setLanguageMode(new TreeSitterLanguageMode({buffer, grammar}))
+      buffer.setText('a\n  .b();')
+      expectTokensToEqual(editor, [
+        [
+          {text: 'a', scopes: ['variable']},
+        ],
+        [
+          {text: '  ', scopes: ['whitespace']},
+          {text: '.', scopes: []},
+          {text: 'b', scopes: ['function']},
+          {text: '();', scopes: []}
+        ]
       ])
     })
   })
@@ -432,22 +457,34 @@ function getDisplayText (editor) {
   return editor.displayLayer.getText()
 }
 
-function expectTokensToEqual (editor, expectedTokens) {
-  const tokens = []
-  for (let row = 0, lastRow = editor.getLastScreenRow(); row <= lastRow; row++) {
-    tokens.push(
-      ...editor.tokensForScreenRow(row).map(({text, scopes}) => ({
+function expectTokensToEqual (editor, expectedTokenLines) {
+  const lastRow = editor.getLastScreenRow()
+
+  // Assert that the correct tokens are returned regardless of which row
+  // the highlighting iterator starts on.
+  for (let startRow = 0; startRow <= lastRow; startRow++) {
+    editor.displayLayer.clearSpatialIndex()
+    editor.displayLayer.getScreenLines(startRow, Infinity)
+
+    const tokenLines = []
+    for (let row = startRow; row <= lastRow; row++) {
+      tokenLines[row] = editor.tokensForScreenRow(row).map(({text, scopes}) => ({
         text,
         scopes: scopes.map(scope => scope
           .split(' ')
           .map(className => className.slice('syntax--'.length))
           .join(' '))
       }))
-    )
-  }
+    }
 
-  expect(tokens.length).toEqual(expectedTokens.length)
-  for (let i = 0; i < tokens.length; i++) {
-    expect(tokens[i]).toEqual(expectedTokens[i], `Token ${i}`)
+    for (let row = startRow; row <= lastRow; row++) {
+      const tokenLine = tokenLines[row]
+      const expectedTokenLine = expectedTokenLines[row]
+
+      expect(tokenLine.length).toEqual(expectedTokenLine.length)
+      for (let i = 0; i < tokenLine.length; i++) {
+        expect(tokenLine[i]).toEqual(expectedTokenLine[i], `Token ${i}, startRow: ${startRow}`)
+      }
+    }
   }
 }
diff --git a/src/tree-sitter-language-mode.js b/src/tree-sitter-language-mode.js
index 33656cf35..5cd725108 100644
--- a/src/tree-sitter-language-mode.js
+++ b/src/tree-sitter-language-mode.js
@@ -349,9 +349,9 @@ class TreeSitterHighlightIterator {
     do {
       this.currentNode = node
       this.currentChildIndex = childIndex
+      if (!nodeContainsTarget) break
       this.containingNodeTypes.push(node.type)
       this.containingNodeChildIndices.push(childIndex)
-      if (!nodeContainsTarget) break
 
       const scopeName = this.currentScopeName()
       if (scopeName) {