Merge #30085 #trim! all whitespace

3 files changed, 58 insertions, 19 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.