From af040c3a079f6e25db0ad6b908aa1327f67deb82 Mon Sep 17 00:00:00 2001 From: Lewis Russell Date: Thu, 11 May 2023 11:13:32 +0100 Subject: [PATCH] feat(treesitter): add support for setting query depths --- cmake/FindTreesitter.cmake | 20 ++++++++++++++++++++ runtime/doc/news.txt | 4 ++++ runtime/doc/treesitter.txt | 7 ++++++- runtime/lua/vim/treesitter/_meta.lua | 14 ++++++++------ runtime/lua/vim/treesitter/query.lua | 8 ++++++-- src/nvim/lua/treesitter.c | 23 +++++++++++++++++++++++ 6 files changed, 67 insertions(+), 9 deletions(-) diff --git a/cmake/FindTreesitter.cmake b/cmake/FindTreesitter.cmake index ddea35fe66..23214283c0 100644 --- a/cmake/FindTreesitter.cmake +++ b/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 +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() diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt index bc357ac534..6b1e9112c1 100644 --- a/runtime/doc/news.txt +++ b/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* diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt index 0168b11499..d425c8dace 100644 --- a/runtime/doc/treesitter.txt +++ b/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: ~ diff --git a/runtime/lua/vim/treesitter/_meta.lua b/runtime/lua/vim/treesitter/_meta.lua index 4d0f43d030..c1009f5f5d 100644 --- a/runtime/lua/vim/treesitter/_meta.lua +++ b/runtime/lua/vim/treesitter/_meta.lua @@ -31,17 +31,19 @@ local TSNode = {} ---@param query userdata ---@param captures true ----@param start integer ----@param end_ integer +---@param start? 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 end_ integer +---@param start? 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[] diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua index 93841bb31e..e6a117557a 100644 --- a/runtime/lua/vim/treesitter/query.lua +++ b/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, 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() diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c index da64685a40..0c4700ccab 100644 --- a/src/nvim/lua/treesitter.c +++ b/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;