From ac3bc51c2a07945fb181808d9d031c379946757b Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Fri, 20 Jul 2018 10:43:05 -0700
Subject: [PATCH] Improve criteria for when to fold partial vs entire buffer
 rows

---
 spec/tree-sitter-language-mode-spec.js | 52 +++++++++++++++++++++++++-
 src/tree-sitter-language-mode.js       | 17 ++++-----
 2 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/spec/tree-sitter-language-mode-spec.js b/spec/tree-sitter-language-mode-spec.js
index 22d80b823..6d1381a93 100644
--- a/spec/tree-sitter-language-mode-spec.js
+++ b/spec/tree-sitter-language-mode-spec.js
@@ -722,6 +722,55 @@ describe('TreeSitterLanguageMode', () => {
       `)
     })
 
+    it('folds entire buffer rows when necessary to keep words on separate lines', async () => {
+      const grammar = new TreeSitterGrammar(atom.grammars, jsGrammarPath, {
+        parser: 'tree-sitter-javascript',
+        folds: [
+          {
+            start: {type: '{', index: 0},
+            end: {type: '}', index: -1}
+          },
+          {
+            start: {type: '(', index: 0},
+            end: {type: ')', index: -1}
+          }
+        ]
+      })
+
+      buffer.setText(dedent `
+        if (a) {
+          b
+        } else if (c) {
+          d
+        } else {
+          e
+        }
+      `)
+
+      const languageMode = new TreeSitterLanguageMode({buffer, grammar})
+      buffer.setLanguageMode(languageMode)
+      await nextHighlightingUpdate(languageMode)
+
+      // Avoid bringing the `else if...` up onto the same screen line as the preceding `if`.
+      editor.foldBufferRow(1)
+      editor.foldBufferRow(3)
+      expect(getDisplayText(editor)).toBe(dedent `
+        if (a) {…
+        } else if (c) {…
+        } else {
+          e
+        }
+      `)
+
+      // It's ok to bring the final `}` onto the same screen line as the preceding `else`.
+      editor.foldBufferRow(5)
+      expect(getDisplayText(editor)).toBe(dedent `
+        if (a) {…
+        } else if (c) {…
+        } else {…}
+      `)
+    })
+
     it('can fold nodes of specified types', async () => {
       const grammar = new TreeSitterGrammar(atom.grammars, jsGrammarPath, {
         parser: 'tree-sitter-javascript',
@@ -1152,7 +1201,8 @@ describe('TreeSitterLanguageMode', () => {
       expect(getDisplayText(editor)).toBe(
         `a = html \`
             <div>
-              c\${def(…)}e\${f}g
+              c\${def(…
+              )}e\${f}g
             </div>
           \`
         `
diff --git a/src/tree-sitter-language-mode.js b/src/tree-sitter-language-mode.js
index b8a400ce7..70de21021 100644
--- a/src/tree-sitter-language-mode.js
+++ b/src/tree-sitter-language-mode.js
@@ -9,6 +9,7 @@ const TextMateLanguageMode = require('./text-mate-language-mode')
 let nextId = 0
 const MAX_RANGE = new Range(Point.ZERO, Point.INFINITY).freeze()
 const PARSER_POOL = []
+const WORD_REGEX = /\w/
 
 class TreeSitterLanguageMode {
   static _patchSyntaxNode () {
@@ -306,11 +307,13 @@ class TreeSitterLanguageMode {
           if (!foldEndNode) continue
         }
 
-        if (foldEndNode.endIndex - foldEndNode.startIndex > 1 && foldEndNode.startPosition.row > foldStart.row) {
-          foldEnd = new Point(foldEndNode.startPosition.row - 1, Infinity)
-        } else {
-          foldEnd = foldEndNode.startPosition
-          if (!pointIsGreater(foldEnd, foldStart)) continue
+        if (foldEndNode.startPosition.row <= foldStart.row) continue
+
+        foldEnd = foldEndNode.startPosition
+        if (this.buffer.findInRangeSync(
+          WORD_REGEX, new Range(foldEnd, new Point(foldEnd.row, Infinity))
+        )) {
+          foldEnd = new Point(foldEnd.row - 1, Infinity)
         }
       } else {
         const {endPosition} = node
@@ -1036,10 +1039,6 @@ function compareScopeDescriptorIterators (a, b) {
   )
 }
 
-function pointIsGreater (left, right) {
-  return left.row > right.row || left.row === right.row && left.column > right.column
-}
-
 function last (array) {
   return array[array.length - 1]
 }