Merge #30085 #trim! all whitespace

5 files changed, 168 insertions, 39 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 858a44e62b..8b4b396a26 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -280,6 +280,8 @@ TREESITTER
 • |LanguageTree:node_for_range()| gets anonymous and named nodes for a range
 • |vim.treesitter.get_node()| now takes an option `include_anonymous`, default
   false, which allows it to return anonymous nodes as well as named nodes.
+• |treesitter-directive-trim!| can trim all whitespace (not just empty lines)
+  from both sides of a node.
 
 TUI
 
diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt
index 5fc6429f7a..877d90a3b7 100644
--- a/runtime/doc/treesitter.txt
+++ b/runtime/doc/treesitter.txt
@@ -245,15 +245,32 @@ The following directives are built in:
             (#gsub! @_node ".*%.(.*)" "%1")
 <
     `trim!`                                          *treesitter-directive-trim!*
-        Trim blank lines from the end of the node. This will set a new
-        `metadata[capture_id].range`.
+        Trims whitespace from the node. Sets a new
+        `metadata[capture_id].range`. Takes a capture ID and, optionally, four
+        integers to customize trimming behavior (`1` meaning trim, `0` meaning
+        don't trim). When only given a capture ID, trims blank lines (lines
+        that contain only whitespace, or are empty) from the end of the node
+        (for backwards compatibility). Can trim all whitespace from both sides
+        of the node if parameters are given.
 
+        Examples: >query
+            ; only trim blank lines from the end of the node
+            ; (equivalent to (#trim! @fold 0 0 1 0))
+            (#trim! @fold)
+
+            ; trim blank lines from both sides of the node
+            (#trim! @fold 1 0 1 0)
+
+            ; trim all whitespace around the node
+            (#trim! @fold 1 1 1 1)
+<
         Parameters: ~
             {capture_id}
+            {trim_start_linewise}
+            {trim_start_charwise}
+            {trim_end_linewise} (default `1` if only given {capture_id})
+            {trim_end_charwise}
 
-        Example: >query
-            (#trim! @fold)
-<
 Further directives can be added via |vim.treesitter.query.add_directive()|.
 Use |vim.treesitter.query.list_directives()| to list all available directives.
 
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index 1677e8d364..3c7bc2eb89 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -572,13 +572,17 @@ local directive_handlers = {
 
     metadata[id].text = text:gsub(pattern, replacement)
   end,
-  -- Trim blank lines from end of the node
-  -- Example: (#trim! @fold)
-  -- TODO(clason): generalize to arbitrary whitespace removal
+  -- Trim whitespace from both sides of the node
+  -- Example: (#trim! @fold 1 1 1 1)
   ['trim!'] = function(match, _, bufnr, pred, metadata)
     local capture_id = pred[2]
     assert(type(capture_id) == 'number')
 
+    local trim_start_lines = pred[3] == '1'
+    local trim_start_cols = pred[4] == '1'
+    local trim_end_lines = pred[5] == '1' or not pred[3] -- default true for backwards compatibility
+    local trim_end_cols = pred[6] == '1'
+
     local nodes = match[capture_id]
     if not nodes or #nodes == 0 then
       return
@@ -588,20 +592,36 @@ local directive_handlers = {
 
     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 end_row >= start_row do
-      -- As we only care when end_col == 0, always inspect one line above end_row.
-      local end_line = api.nvim_buf_get_lines(bufnr, end_row - 1, end_row, true)[1]
+    local node_text = vim.split(vim.treesitter.get_node_text(node, bufnr), '\n')
+    local end_idx = #node_text
+    local start_idx = 1
 
-      if end_line ~= '' then
-        break
+    if trim_end_lines then
+      while end_idx > 0 and node_text[end_idx]:find('^%s*$') do
+        end_idx = end_idx - 1
+        end_row = end_row - 1
       end
+    end
+    if trim_end_cols then
+      if end_idx == 0 then
+        end_row = start_row
+        end_col = start_col
+      else
+        local whitespace_start = node_text[end_idx]:find('(%s*)$')
+        end_col = (whitespace_start - 1) + (end_idx == 1 and start_col or 0)
+      end
+    end
 
-      end_row = end_row - 1
+    if trim_start_lines then
+      while start_idx <= end_idx and node_text[start_idx]:find('^%s*$') do
+        start_idx = start_idx + 1
+        start_row = start_row + 1
+      end
+    end
+    if trim_start_cols and node_text[start_idx] then
+      local _, whitespace_end = node_text[start_idx]:find('^(%s*)')
+      whitespace_end = whitespace_end or 0
+      start_col = (start_idx == 1 and start_col or 0) + whitespace_end
     end
 
     -- If this produces an invalid range, we just skip it.
diff --git a/test/functional/treesitter/parser_spec.lua b/test/functional/treesitter/parser_spec.lua
index 2f8d204d36..8a552c7d34 100644
--- a/test/functional/treesitter/parser_spec.lua
+++ b/test/functional/treesitter/parser_spec.lua
@@ -1,5 +1,6 @@
 local t = require('test.testutil')
 local n = require('test.functional.testnvim')()
+local ts_t = require('test.functional.treesitter.testutil')
 
 local clear = n.clear
 local dedent = t.dedent
@@ -8,6 +9,7 @@ local insert = n.insert
 local exec_lua = n.exec_lua
 local pcall_err = t.pcall_err
 local feed = n.feed
+local run_query = ts_t.run_query
 
 describe('treesitter parser API', function()
   before_each(function()
@@ -644,6 +646,82 @@ print()
     end)
   end)
 
+  describe('trim! directive', function()
+    it('can trim all whitespace', function()
+      -- luacheck: push ignore 611 613
+      insert([=[
+        print([[
+
+                  f
+           helllo
+        there
+        asdf
+        asdfassd   
+
+
+
+        ]])
+        print([[
+              
+              
+              
+        ]])
+
+        print([[]])
+
+        print([[
+        ]])
+
+        print([[     hello 😃    ]])
+      ]=])
+      -- luacheck: pop
+
+      local query_text = [[
+        ; query
+        ((string_content) @str
+          (#trim! @str 1 1 1 1))
+      ]]
+
+      exec_lua(function()
+        vim.treesitter.start(0, 'lua')
+      end)
+
+      eq({
+        { 'str', { 2, 12, 6, 10 } },
+        { 'str', { 11, 10, 11, 10 } },
+        { 'str', { 17, 10, 17, 10 } },
+        { 'str', { 19, 10, 19, 10 } },
+        { 'str', { 22, 15, 22, 25 } },
+      }, run_query('lua', query_text))
+    end)
+
+    it('trims only empty lines by default (backwards compatible)', function()
+      insert [[
+      ## Heading
+
+      With some text
+
+      ## And another
+
+      With some more]]
+
+      local query_text = [[
+        ; query
+        ((section) @fold
+          (#trim! @fold))
+      ]]
+
+      exec_lua(function()
+        vim.treesitter.start(0, 'markdown')
+      end)
+
+      eq({
+        { 'fold', { 0, 0, 3, 0 } },
+        { 'fold', { 4, 0, 7, 0 } },
+      }, run_query('markdown', query_text))
+    end)
+  end)
+
   it('tracks the root range properly (#22911)', function()
     insert([[
       int main() {
@@ -659,32 +737,19 @@ print()
       vim.treesitter.start(0, 'c')
     end)
 
-    local function run_query()
-      return exec_lua(function()
-        local query = vim.treesitter.query.parse('c', query0)
-        local parser = vim.treesitter.get_parser()
-        local tree = parser:parse()[1]
-        local res = {}
-        for id, node in query:iter_captures(tree:root()) do
-          table.insert(res, { query.captures[id], node:range() })
-        end
-        return res
-      end)
-    end
-
     eq({
-      { 'function', 0, 0, 2, 1 },
-      { 'declaration', 1, 2, 1, 12 },
-    }, run_query())
+      { 'function', { 0, 0, 2, 1 } },
+      { 'declaration', { 1, 2, 1, 12 } },
+    }, run_query('c', query0))
 
     n.command 'normal ggO'
     insert('int a;')
 
     eq({
-      { 'declaration', 0, 0, 0, 6 },
-      { 'function', 1, 0, 3, 1 },
-      { 'declaration', 2, 2, 2, 12 },
-    }, run_query())
+      { 'declaration', { 0, 0, 0, 6 } },
+      { 'function', { 1, 0, 3, 1 } },
+      { 'declaration', { 2, 2, 2, 12 } },
+    }, run_query('c', query0))
   end)
 
   it('handles ranges when source is a multiline string (#20419)', function()
diff --git a/test/functional/treesitter/testutil.lua b/test/functional/treesitter/testutil.lua
new file mode 100644
index 0000000000..f8934f06c3
--- /dev/null
+++ b/test/functional/treesitter/testutil.lua
@@ -0,0 +1,25 @@
+local n = require('test.functional.testnvim')()
+
+local exec_lua = n.exec_lua
+
+local M = {}
+
+---@param language string
+---@param query_string string
+function M.run_query(language, query_string)
+  return exec_lua(function(lang, query_str)
+    local query = vim.treesitter.query.parse(lang, query_str)
+    local parser = vim.treesitter.get_parser()
+    local tree = parser:parse()[1]
+    local res = {}
+    for id, node, metadata in query:iter_captures(tree:root(), 0) do
+      table.insert(
+        res,
+        { query.captures[id], metadata[id] and metadata[id].range or { node:range() } }
+      )
+    end
+    return res
+  end, language, query_string)
+end
+
+return M