feat(treesitter): add support for setting query depths

cmake/FindTreesitter.cmake

@@ -7,3 +7,23 @@ mark_as_advanced(TREESITTER_LIBRARY TREESITTER_INCLUDE_DIR)
 add_library(treesitter INTERFACE)
 target_include_directories(treesitter SYSTEM BEFORE INTERFACE ${TREESITTER_INCLUDE_DIR})
 target_link_libraries(treesitter INTERFACE ${TREESITTER_LIBRARY})
+
+# TODO(lewis6991): remove when min TS version is 0.20.9
+list(APPEND CMAKE_REQUIRED_INCLUDES "${TREESITTER_INCLUDE_DIR}")
+list(APPEND CMAKE_REQUIRED_LIBRARIES "${TREESITTER_LIBRARY}")
+check_c_source_compiles("
+#include <tree_sitter/api.h>
+int
+main(void)
+{
+  TSQueryCursor *cursor = ts_query_cursor_new();
+  ts_query_cursor_set_max_start_depth(cursor, 32);
+  return 0;
+}
+" TS_HAS_SET_MAX_START_DEPTH)
+list(REMOVE_ITEM CMAKE_REQUIRED_INCLUDES "${TREESITTER_INCLUDE_DIR}")
+list(REMOVE_ITEM CMAKE_REQUIRED_LIBRARIES "${TREESITTER_LIBRARY}")
+
+if(TS_HAS_SET_MAX_START_DEPTH)
+  target_compile_definitions(treesitter INTERFACE NVIM_TS_HAS_SET_MAX_START_DEPTH)
+endif()

runtime/doc/news.txt

@@ -29,6 +29,7 @@ The following changes may require adaptations in user config or plugins.
     set selectmode=mouse,key
     set mousemodel=popup
     set keymodel=startsel,stopsel
+<
 
 ==============================================================================
 ADDED FEATURES                                                     *news-added*
@@ -52,6 +53,9 @@ iterators |luaref-in|.
 • |'smoothscroll'| option to scroll by screen line rather than by text line
 when |'wrap'| is set.
 
+• |Query:iter_matches()| now has the ability to set the maximum start depth
+  for matches.
+
 ==============================================================================
 CHANGED FEATURES                                                 *news-changed*
 

runtime/doc/treesitter.txt

@@ -939,7 +939,7 @@ Query:iter_captures({self}, {node}, {source}, {start}, {stop})
         metadata
 
                                                         *Query:iter_matches()*
-Query:iter_matches({self}, {node}, {source}, {start}, {stop})
+Query:iter_matches({self}, {node}, {source}, {start}, {stop}, {opts})
     Iterates the matches of self on a given range.
 
     Iterate over all matches within a {node}. The arguments are the same as
@@ -966,6 +966,11 @@ Query:iter_matches({self}, {node}, {source}, {start}, {stop})
       • {source}  (integer|string) Source buffer or string to search
       • {start}   (integer) Starting line for the search
       • {stop}    (integer) Stopping line for the search (end-exclusive)
+      • {opts}    (table|nil) Options:
+                  • max_start_depth (integer) if non-zero, sets the maximum
+                    start depth for each match. This is used to prevent
+                    traversing too deep into a tree. Requires treesitter >=
+                    0.20.9.
       • {self}
 
     Return: ~

runtime/lua/vim/treesitter/_meta.lua

@@ -31,17 +31,19 @@ local TSNode = {}
 
 ---@param query userdata
 ---@param captures true
----@param start integer
+---@param start? integer
----@param end_ integer
+---@param end_? integer
+---@param opts? table
 ---@return fun(): integer, TSNode, any
-function TSNode:_rawquery(query, captures, start, end_) end
+function TSNode:_rawquery(query, captures, start, end_, opts) end
 
 ---@param query userdata
 ---@param captures false
----@param start integer
+---@param start? integer
----@param end_ integer
+---@param end_? integer
+---@param opts? table
 ---@return fun(): string, any
-function TSNode:_rawquery(query, captures, start, end_) end
+function TSNode:_rawquery(query, captures, start, end_, opts) end
 
 ---@class TSParser
 ---@field parse fun(self: TSParser, tree: TSTree?, source: integer|string, include_bytes: boolean?): TSTree, integer[]

runtime/lua/vim/treesitter/query.lua

@@ -686,16 +686,20 @@ end
 ---@param source (integer|string) Source buffer or string to search
 ---@param start integer Starting line for the search
 ---@param stop integer Stopping line for the search (end-exclusive)
+---@param opts table|nil Options:
+---   - max_start_depth (integer) if non-zero, sets the maximum start depth
+---     for each match. This is used to prevent traversing too deep into a tree.
+---     Requires treesitter >= 0.20.9.
 ---
 ---@return (fun(): integer, table<integer,TSNode>, table): pattern id, match, metadata
-function Query:iter_matches(node, source, start, stop)
+function Query:iter_matches(node, source, start, stop, opts)
   if type(source) == 'number' and source == 0 then
     source = api.nvim_get_current_buf()
   end
 
   start, stop = value_or_node_range(start, stop, node)
 
-  local raw_iter = node:_rawquery(self.query, false, start, stop)
+  local raw_iter = node:_rawquery(self.query, false, start, stop, opts)
   ---@cast raw_iter fun(): string, any
   local function iter()
     local pattern, match = raw_iter()

src/nvim/lua/treesitter.c

@@ -1360,6 +1360,29 @@ static int node_rawquery(lua_State *L)
     ts_query_cursor_set_point_range(cursor, (TSPoint){ start, 0 }, (TSPoint){ end, 0 });
   }
 
+  if (lua_gettop(L) >= 6 && !lua_isnil(L, 6)) {
+    if (!lua_istable(L, 6)) {
+      return luaL_error(L, "table expected");
+    }
+    lua_pushnil(L);
+    // stack: [dict, ..., nil]
+    while (lua_next(L, 6)) {
+      // stack: [dict, ..., key, value]
+      if (lua_type(L, -2) == LUA_TSTRING) {
+        char *k = (char *)lua_tostring(L, -2);
+        if (strequal("max_start_depth", k)) {
+          // TODO(lewis6991): remove ifdef when min TS version is 0.20.9
+#ifdef NVIM_TS_HAS_SET_MAX_START_DEPTH
+          uint32_t max_start_depth = (uint32_t)lua_tointeger(L, -1);
+          ts_query_cursor_set_max_start_depth(cursor, max_start_depth);
+#endif
+        }
+      }
+      lua_pop(L, 1);  // pop the value; lua_next will pop the key.
+      // stack: [dict, ..., key]
+    }
+  }
+
   TSLua_cursor *ud = lua_newuserdata(L, sizeof(*ud));  // [udata]
   ud->cursor = cursor;
   ud->predicated_match = -1;