fix: Use native vim.treesitter module when available.

Author: kylechui

lua/nvim-surround/queries.lua

@@ -1,33 +1,22 @@
-local utils = require("nvim-surround.utils")
-local ts_query = require("nvim-treesitter.query")
-local ts_utils = require("nvim-treesitter.ts_utils")
-local ts_parsers = require("nvim-treesitter.parsers")
-
 local M = {}
 
 -- Retrieves the node that corresponds exactly to a given selection.
 ---@param selection selection The given selection.
----@return _ @The corresponding node.
+---@return TSNode|nil @The corresponding node.
 ---@nodiscard
 M.get_node = function(selection)
-    -- Convert the selection into a list
-    local range = {
-        selection.first_pos[1],
-        selection.first_pos[2],
-        selection.last_pos[1],
-        selection.last_pos[2],
-    }
+    local treesitter = require("nvim-surround.treesitter")
 
-    -- Get the root node of the current tree
-    local lang_tree = ts_parsers.get_parser(0)
-    local tree = lang_tree:trees()[1]
-    local root = tree:root()
+    local root = treesitter.get_root()
+    if root == nil then
+        return nil
+    end
     -- DFS through the tree and find all nodes that have the given type
     local stack = { root }
     while #stack > 0 do
         local cur = stack[#stack]
         -- If the current node's range is equal to the desired selection, return the node
-        if vim.deep_equal(range, { ts_utils.get_vim_range({ cur:range() }) }) then
+        if vim.deep_equal(selection, treesitter.get_node_selection(cur)) then
             return cur
         end
         -- Pop off of the stack
@@ -40,48 +29,24 @@ M.get_node = function(selection)
     return nil
 end
 
--- Filters an existing parent selection down to a capture.
----@param sexpr string The given S-expression containing the capture.
----@param capture string The name of the capture to be returned.
----@param parent_selection selection The parent selection to be filtered down.
-M.filter_selection = function(sexpr, capture, parent_selection)
-    local parent_node = M.get_node(parent_selection)
-
-    local range = { ts_utils.get_vim_range({ parent_node:range() }) }
-    local lang_tree = ts_parsers.get_parser(0)
-    local ok, parsed_query = pcall(function()
-        return vim.treesitter.query.parse and vim.treesitter.query.parse(lang_tree:lang(), sexpr)
-            or vim.treesitter.parse_query(lang_tree:lang(), sexpr)
-    end)
-    if not ok or not parent_node then
-        return {}
-    end
-
-    for id, node in parsed_query:iter_captures(parent_node, 0, 0, -1) do
-        local name = parsed_query.captures[id]
-        if name == capture then
-            range = { ts_utils.get_vim_range({ node:range() }) }
-            return {
-                first_pos = { range[1], range[2] },
-                last_pos = { range[3], range[4] },
-            }
-        end
-    end
-    return nil
-end
-
 -- Finds the nearest selection of a given query capture and its source.
 ---@param capture string The capture to be retrieved.
 ---@param type string The type of query to get the capture from.
 ---@return selection|nil @The selection of the capture.
 ---@nodiscard
 M.get_selection = function(capture, type)
+    local utils = require("nvim-surround.utils")
+    local treesitter = require("nvim-surround.treesitter")
+    local ts_query = require("nvim-treesitter.query")
+
     -- Get a table of all nodes that match the query
     local table_list = ts_query.get_capture_matches_recursively(0, capture, type)
     -- Convert the list of nodes into a list of selections
     local selections_list = {}
     for _, tab in ipairs(table_list) do
-        local range = { ts_utils.get_vim_range({ tab.node:range() }) }
+        local selection = treesitter.get_node_selection(tab.node)
+
+        local range = { selection.first_pos[1], selection.first_pos[2], selection.last_pos[1], selection.last_pos[2] }
         selections_list[#selections_list + 1] = {
             left = {
                 first_pos = { range[1], range[2] },

lua/nvim-surround/treesitter.lua

@@ -14,28 +14,63 @@ local function is_any_of(target, types)
     return false
 end
 
+-- Gets the root node for the buffer.
+---@return TSNode|nil @The root node for the buffer.
+---@nodiscard
+M.get_root = function()
+    if vim.treesitter then
+        return vim.treesitter.get_node():tree():root()
+    end
+
+    local ok, ts_utils = pcall(require, "nvim-treesitter.ts_utils")
+    if not ok then
+        return nil
+    end
+
+    local node = ts_utils.get_node_at_cursor()
+    if node == nil then
+        return nil
+    end
+    return ts_utils.get_root_for_node(node)
+end
+
+-- Gets the selection that a TreeSitter node spans.
+---@param node TSNode The given TreeSitter node.
+---@return selection @The span of the input node.
+---@nodiscard
+M.get_node_selection = function(node)
+    ---@type integer, integer, integer, integer
+    local srow, scol, erow, ecol = node:range()
+    srow = srow + 1
+    scol = scol + 1
+    erow = erow + 1
+
+    if ecol == 0 then
+        -- Use the value of the last col of the previous row instead.
+        erow = erow - 1
+        ecol = math.max(1, vim.fn.col({ erow, "$" }) - 1)
+    end
+    return {
+        first_pos = { srow, scol },
+        last_pos = { erow, ecol },
+    }
+end
+
 -- Finds the nearest selection of a given Tree-sitter node type or types.
 ---@param node_types string|string[] The Tree-sitter node type(s) to be retrieved.
 ---@return selection|nil @The selection of the node.
 ---@nodiscard
 M.get_selection = function(node_types)
+    local utils = require("nvim-surround.utils")
+
     if type(node_types) == "string" then
         node_types = { node_types }
     end
 
-    local utils = require("nvim-surround.utils")
-    local ok, ts_utils = pcall(require, "nvim-treesitter.ts_utils")
-    if not ok then
+    local root = M.get_root()
+    if root == nil then
         return nil
     end
-    -- Find the root node of the given buffer
-    local root = ts_utils.get_node_at_cursor()
-    if not root then
-        return {}
-    end
-    while root:parent() do
-        root = root:parent()
-    end
     -- DFS through the tree and find all nodes that have the given type
     local stack = { root }
     local nodes, selections_list = {}, {}
@@ -46,15 +81,8 @@ M.get_selection = function(node_types)
             -- Add the current node to the stack
             nodes[#nodes + 1] = cur
             -- Compute the node's selection and add it to the list
-            local range = { ts_utils.get_vim_range({ cur:range() }) }
-            selections_list[#selections_list + 1] = {
-                left = {
-                    first_pos = { range[1], range[2] },
-                },
-                right = {
-                    last_pos = { range[3], range[4] },
-                },
-            }
+            local selection = M.get_node_selection(cur)
+            selections_list[#selections_list + 1] = selection
         end
         -- Pop off of the stack
         stack[#stack] = nil