From 2eaf18826988f921ddb39e4a2e7d23d95bb0e52a Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Thu, 21 Apr 2022 10:48:47 +0100
Subject: fix(fold): don't include whitespace end regions

Some languages that are difficult to parse via Treesitter may
incorrectly include whitespace lines at the end of regions. This can
makes the calculated folds sub-optimal.

To recitfy, use a custom directive (trim), to calculate the range with
the trailing whitespace lines removed. Note this only works if the
region end column is 0.

Also added folds for Make.
---
 lua/nvim-treesitter/fold.lua             |  9 +++++++--
 lua/nvim-treesitter/query.lua            |  4 +++-
 lua/nvim-treesitter/query_predicates.lua | 33 ++++++++++++++++++++++++++++++++
 3 files changed, 43 insertions(+), 3 deletions(-)

(limited to 'lua')

diff --git a/lua/nvim-treesitter/fold.lua b/lua/nvim-treesitter/fold.lua
index 01a7c08e9..1f417414f 100644
--- a/lua/nvim-treesitter/fold.lua
+++ b/lua/nvim-treesitter/fold.lua
@@ -39,8 +39,13 @@ local folds_levels = tsutils.memoize_by_buf_tick(function(bufnr)
 
   local min_fold_lines = api.nvim_win_get_option(0, "foldminlines")
 
-  for _, node in ipairs(matches) do
-    local start, _, stop, stop_col = node.node:range()
+  for _, match in ipairs(matches) do
+    local start, stop, stop_col
+    if match.metadata and match.metadata.range then
+      start, _, stop, stop_col = unpack(match.metadata.range)
+    else
+      start, _, stop, stop_col = match.node:range()
+    end
 
     if stop_col == 0 then
       stop = stop - 1
diff --git a/lua/nvim-treesitter/query.lua b/lua/nvim-treesitter/query.lua
index bad1c27e6..aeaa683be 100644
--- a/lua/nvim-treesitter/query.lua
+++ b/lua/nvim-treesitter/query.lua
@@ -209,7 +209,7 @@ function M.iter_prepared_matches(query, qnode, bufnr, start_row, end_row)
   local matches = query:iter_matches(qnode, bufnr, start_row, end_row)
 
   local function iterator()
-    local pattern, match = matches()
+    local pattern, match, metadata = matches()
     if pattern ~= nil then
       local prepared_match = {}
 
@@ -219,6 +219,8 @@ function M.iter_prepared_matches(query, qnode, bufnr, start_row, end_row)
         if name ~= nil then
           local path = split(name .. ".node")
           insert_to_path(prepared_match, path, node)
+          local metadata_path = split(name .. ".metadata")
+          insert_to_path(prepared_match, metadata_path, metadata[id])
         end
       end
 
diff --git a/lua/nvim-treesitter/query_predicates.lua b/lua/nvim-treesitter/query_predicates.lua
index 75f0efc04..0638b6e79 100644
--- a/lua/nvim-treesitter/query_predicates.lua
+++ b/lua/nvim-treesitter/query_predicates.lua
@@ -153,3 +153,36 @@ query.add_directive("exclude_children!", function(match, _pattern, _bufnr, pred,
   end
   metadata.content = ranges
 end)
+
+-- Trim blank lines from end of the region
+-- Arguments are the captures to trim.
+query.add_directive("trim!", function(match, _, bufnr, pred, metadata)
+  for _, id in ipairs { select(2, unpack(pred)) } do
+    local node = match[id]
+    local start_row, start_col, end_row, end_col = node:range()
+
+    -- Don't trim if region ends in middle of a line
+    if end_col ~= 0 then
+      return
+    end
+
+    while true do
+      -- As we only care when end_col == 0, always inspect one line above end_row.
+      local end_line = vim.api.nvim_buf_get_lines(bufnr, end_row - 1, end_row, true)[1]
+
+      if end_line ~= "" then
+        break
+      end
+
+      end_row = end_row - 1
+    end
+
+    -- If this produces an invalid range, we just skip it.
+    if start_row < end_row or (start_row == end_row and start_col <= end_col) then
+      if not metadata[id] then
+        metadata[id] = {}
+      end
+      metadata[id].range = { start_row, start_col, end_row, end_col }
+    end
+  end
+end)
-- 
cgit v1.2.3-70-g09d2