From 136dc86584a04f42b69b01685af8ec8c8403ca88 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Thu, 7 Dec 2017 15:29:11 -0800
Subject: [PATCH] Leave muli-character fold end tokens on their own line

Signed-off-by: Nathan Sobo <nathan@github.com>
---
 spec/tree-sitter-language-mode-spec.js | 33 ++++++++++++++++++++++----
 src/tree-sitter-language-mode.js       | 16 +++++++++----
 2 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/spec/tree-sitter-language-mode-spec.js b/spec/tree-sitter-language-mode-spec.js
index 426291e5f..0eeeb8b93 100644
--- a/spec/tree-sitter-language-mode-spec.js
+++ b/spec/tree-sitter-language-mode-spec.js
@@ -186,7 +186,8 @@ describe('TreeSitterLanguageMode', () => {
       expect(getDisplayText(editor)).toBe(dedent `
         const element1 = <Element…/>
 
-        const element2 = <Element>…</Element>
+        const element2 = <Element>…
+        </Element>
       `)
     })
 
@@ -239,10 +240,15 @@ describe('TreeSitterLanguageMode', () => {
         folds: [
           // If the #ifdef has an `#else` clause, then end the fold there.
           {
-            type: 'preproc_ifdef',
+            type: ['preproc_ifdef', 'preproc_elif'],
             start: {index: 1},
             end: {type: 'preproc_else'}
           },
+          {
+            type: ['preproc_ifdef', 'preproc_elif'],
+            start: {index: 1},
+            end: {type: 'preproc_elif'}
+          },
 
           // Otherwise, end the fold at the last child - the `#endif`.
           {
@@ -270,6 +276,11 @@ describe('TreeSitterLanguageMode', () => {
         #include <windows.h>
         const char *path_separator = "\\";
 
+        #elif defined MACOS
+
+        #include <carbon.h>
+        const char *path_separator = "/";
+
         #else
 
         #include <dirent.h>
@@ -287,7 +298,13 @@ describe('TreeSitterLanguageMode', () => {
         #ifndef FOO_H_
         #define FOO_H_
 
-        #ifdef _WIN32…#else
+        #ifdef _WIN32…
+        #elif defined MACOS
+
+        #include <carbon.h>
+        const char *path_separator = "/";
+
+        #else
 
         #include <dirent.h>
         const char *path_separator = "/";
@@ -302,7 +319,12 @@ describe('TreeSitterLanguageMode', () => {
         #ifndef FOO_H_
         #define FOO_H_
 
-        #ifdef _WIN32…#else…
+        #ifdef _WIN32…
+        #elif defined MACOS…
+        #else
+
+        #include <dirent.h>
+        const char *path_separator = "/";
 
         #endif
 
@@ -311,7 +333,8 @@ describe('TreeSitterLanguageMode', () => {
 
       editor.foldBufferRow(0)
       expect(getDisplayText(editor)).toBe(dedent `
-        #ifndef FOO_H_…#endif
+        #ifndef FOO_H_…
+        #endif
       `)
     })
 
diff --git a/src/tree-sitter-language-mode.js b/src/tree-sitter-language-mode.js
index f47d89db7..a76043638 100644
--- a/src/tree-sitter-language-mode.js
+++ b/src/tree-sitter-language-mode.js
@@ -221,16 +221,22 @@ class TreeSitterLanguageMode {
       let foldEnd
       const endEntry = foldEntry.end
       if (endEntry) {
+        let foldEndNode
         if (endEntry.index != null) {
           const index = endEntry.index < 0 ? childCount + endEntry.index : endEntry.index
-          const child = children[index]
-          if (!child || (endEntry.type && endEntry.type !== child.type)) continue
-          foldEnd = child.startPosition
+          foldEndNode = children[index]
+          if (!foldEndNode || (endEntry.type && endEntry.type !== foldEndNode.type)) continue
         } else {
-          if (!childTypes) childTypes = children.map(child => child.type)
+          if (!childTypes) childTypes = children.map(foldEndNode => foldEndNode.type)
           const index = childTypes.lastIndexOf(endEntry.type)
           if (index === -1) continue
-          foldEnd = children[index].startPosition
+          foldEndNode = children[index]
+        }
+
+        if (foldEndNode.endIndex - foldEndNode.startIndex > 1 && foldEndNode.startPosition.row > foldStart.row) {
+          foldEnd = new Point(foldEndNode.startPosition.row - 1, Infinity)
+        } else {
+          foldEnd = foldEndNode.startPosition
         }
       }