|
| 1 | +NOTE: has been slightly edited in order to apply cleanly |
| 2 | + |
| 3 | +From a9c6564bcbee9fb0e00c056d8c88fd9f60e49e7e Mon Sep 17 00:00:00 2001 |
| 4 | +From: Riley Bruins <ribru17@hotmail.com> |
| 5 | +Date: Thu, 20 Mar 2025 14:50:15 -0700 |
| 6 | +Subject: [PATCH] refactor(treesitter): migrate to ts parser callback API |
| 7 | + |
| 8 | +--- |
| 9 | + runtime/lua/vim/treesitter/_meta/misc.lua | 4 +- |
| 10 | + runtime/lua/vim/treesitter/languagetree.lua | 43 ++++++++++------ |
| 11 | + src/nvim/lua/treesitter.c | 57 +++++++++++---------- |
| 12 | + 3 files changed, 58 insertions(+), 46 deletions(-) |
| 13 | + |
| 14 | +diff --git a/runtime/lua/vim/treesitter/_meta/misc.lua b/runtime/lua/vim/treesitter/_meta/misc.lua |
| 15 | +index 07a1c921c744c3..9b9cc4eb5469d5 100644 |
| 16 | +--- a/runtime/lua/vim/treesitter/_meta/misc.lua |
| 17 | ++++ b/runtime/lua/vim/treesitter/_meta/misc.lua |
| 18 | +@@ -5,12 +5,10 @@ error('Cannot require a meta file') |
| 19 | + ---@alias TSLoggerCallback fun(logtype: 'parse'|'lex', msg: string) |
| 20 | + |
| 21 | + ---@class TSParser: userdata |
| 22 | +----@field parse fun(self: TSParser, tree: TSTree?, source: integer|string, include_bytes: boolean): TSTree, (Range4|Range6)[] |
| 23 | ++---@field parse fun(self: TSParser, tree: TSTree?, source: integer|string, include_bytes: boolean, timeout_ns: integer?): TSTree?, (Range4|Range6)[] |
| 24 | + ---@field reset fun(self: TSParser) |
| 25 | + ---@field included_ranges fun(self: TSParser, include_bytes: boolean?): integer[] |
| 26 | + ---@field set_included_ranges fun(self: TSParser, ranges: (Range6|TSNode)[]) |
| 27 | +----@field set_timeout fun(self: TSParser, timeout: integer) |
| 28 | +----@field timeout fun(self: TSParser): integer |
| 29 | + ---@field _set_logger fun(self: TSParser, lex: boolean, parse: boolean, cb: TSLoggerCallback) |
| 30 | + ---@field _logger fun(self: TSParser): TSLoggerCallback |
| 31 | + |
| 32 | +diff --git a/runtime/lua/vim/treesitter/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua |
| 33 | +index e7cee33a038866..6f0e377d2f920f 100644 |
| 34 | +--- a/runtime/lua/vim/treesitter/languagetree.lua |
| 35 | ++++ b/runtime/lua/vim/treesitter/languagetree.lua |
| 36 | +@@ -43,8 +43,10 @@ |
| 37 | + local query = require('vim.treesitter.query') |
| 38 | + local language = require('vim.treesitter.language') |
| 39 | + local Range = require('vim.treesitter._range') |
| 40 | ++local hrtime = vim.uv.hrtime |
| 41 | + |
| 42 | +-local default_parse_timeout_ms = 3 |
| 43 | ++-- Parse in 3ms chunks. |
| 44 | ++local default_parse_timeout_ns = 3 * 1000000 |
| 45 | + |
| 46 | + ---@type Range2 |
| 47 | + local entire_document_range = { 0, math.huge } |
| 48 | +@@ -198,16 +200,16 @@ function LanguageTree:_set_logger() |
| 49 | + self._parser:_set_logger(log_lex, log_parse, self._logger) |
| 50 | + end |
| 51 | + |
| 52 | +----Measure execution time of a function |
| 53 | ++---Measure execution time of a function, in nanoseconds. |
| 54 | + ---@generic R1, R2, R3 |
| 55 | + ---@param f fun(): R1, R2, R3 |
| 56 | + ---@return number, R1, R2, R3 |
| 57 | + local function tcall(f, ...) |
| 58 | +- local start = vim.uv.hrtime() |
| 59 | ++ local start = hrtime() |
| 60 | + ---@diagnostic disable-next-line |
| 61 | + local r = { f(...) } |
| 62 | + --- @type number |
| 63 | +- local duration = (vim.uv.hrtime() - start) / 1000000 |
| 64 | ++ local duration = hrtime() - start |
| 65 | + --- @diagnostic disable-next-line: redundant-return-value |
| 66 | + return duration, unpack(r) |
| 67 | + end |
| 68 | +@@ -388,18 +390,29 @@ function LanguageTree:_parse_regions(range, thread_state) |
| 69 | + ) |
| 70 | + then |
| 71 | + self._parser:set_included_ranges(ranges) |
| 72 | +- self._parser:set_timeout(thread_state.timeout and thread_state.timeout * 1000 or 0) -- ms -> micros |
| 73 | + |
| 74 | +- local parse_time, tree, tree_changes = |
| 75 | +- tcall(self._parser.parse, self._parser, self._trees[i], self._source, true) |
| 76 | ++ local parse_time, tree, tree_changes = tcall( |
| 77 | ++ self._parser.parse, |
| 78 | ++ self._parser, |
| 79 | ++ self._trees[i], |
| 80 | ++ self._source, |
| 81 | ++ true, |
| 82 | ++ thread_state.timeout |
| 83 | ++ ) |
| 84 | + while true do |
| 85 | + if tree then |
| 86 | + break |
| 87 | + end |
| 88 | + coroutine.yield(self._trees, false) |
| 89 | + |
| 90 | +- parse_time, tree, tree_changes = |
| 91 | +- tcall(self._parser.parse, self._parser, self._trees[i], self._source, true) |
| 92 | ++ parse_time, tree, tree_changes = tcall( |
| 93 | ++ self._parser.parse, |
| 94 | ++ self._parser, |
| 95 | ++ self._trees[i], |
| 96 | ++ self._source, |
| 97 | ++ true, |
| 98 | ++ thread_state.timeout |
| 99 | ++ ) |
| 100 | + end |
| 101 | + |
| 102 | + self:_subtract_time(thread_state, parse_time) |
| 103 | +@@ -503,7 +516,7 @@ function LanguageTree:_async_parse(range, on_parse) |
| 104 | + local buf = is_buffer_parser and vim.b[source] or nil |
| 105 | + local ct = is_buffer_parser and buf.changedtick or nil |
| 106 | + local total_parse_time = 0 |
| 107 | +- local redrawtime = vim.o.redrawtime |
| 108 | ++ local redrawtime = vim.o.redrawtime * 1000000 |
| 109 | + |
| 110 | + local thread_state = {} ---@type ParserThreadState |
| 111 | + |
| 112 | +diff --git a/runtime/lua/vim/treesitter/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua |
| 113 | +index a346bf59637fa0..5d3599920ea308 100644 |
| 114 | +--- a/runtime/lua/vim/treesitter/languagetree.lua |
| 115 | ++++ b/runtime/lua/vim/treesitter/languagetree.lua |
| 116 | +@@ -997,7 +997,7 @@ |
| 117 | + |
| 118 | + ---@type table<integer,vim.treesitter.languagetree.Injection> |
| 119 | + local injections = {} |
| 120 | +- local start = vim.uv.hrtime() |
| 121 | ++ local start = hrtime() |
| 122 | + |
| 123 | + local full_scan = range == true or self._injection_query.has_combined_injections |
| 124 | + |
| 125 | +@@ -1021,9 +1021,9 @@ |
| 126 | + end |
| 127 | + |
| 128 | + -- Check the current function duration against the timeout, if it exists. |
| 129 | +- local current_time = vim.uv.hrtime() |
| 130 | +- self:_subtract_time(thread_state, (current_time - start) / 1000000) |
| 131 | +- start = current_time |
| 132 | ++ local current_time = hrtime() |
| 133 | ++ self:_subtract_time(thread_state, current_time - start) |
| 134 | ++ start = hrtime() |
| 135 | + end |
| 136 | + end |
| 137 | + |
| 138 | +diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c |
| 139 | +index a346bf59637fa0..5d3599920ea308 100644 |
| 140 | +--- a/src/nvim/lua/treesitter.c |
| 141 | ++++ b/src/nvim/lua/treesitter.c |
| 142 | +@@ -15,6 +15,8 @@ |
| 143 | + #include <tree_sitter/api.h> |
| 144 | + #include <uv.h> |
| 145 | + |
| 146 | ++#include "nvim/os/time.h" |
| 147 | ++ |
| 148 | + #ifdef HAVE_WASMTIME |
| 149 | + # include <wasm.h> |
| 150 | + |
| 151 | +@@ -54,5 +56,10 @@ typedef struct { |
| 152 | + } TSLuaTree; |
| 153 | + |
| 154 | ++typedef struct { |
| 155 | ++ uint64_t parse_start_time; |
| 156 | ++ uint64_t timeout_threshold_ns; |
| 157 | ++} TSLuaParserCallbackPayload; |
| 158 | ++ |
| 159 | + #ifdef INCLUDE_GENERATED_DECLARATIONS |
| 160 | + # include "lua/treesitter.c.generated.h" |
| 161 | + #endif |
| 162 | +@@ -362,8 +369,6 @@ static struct luaL_Reg parser_meta[] = { |
| 163 | + { "reset", parser_reset }, |
| 164 | + { "set_included_ranges", parser_set_ranges }, |
| 165 | + { "included_ranges", parser_get_ranges }, |
| 166 | +- { "set_timeout", parser_set_timeout }, |
| 167 | +- { "timeout", parser_get_timeout }, |
| 168 | + { "_set_logger", parser_set_logger }, |
| 169 | + { "_logger", parser_get_logger }, |
| 170 | + { NULL, NULL } |
| 171 | +@@ -487,6 +492,13 @@ static void push_ranges(lua_State *L, const TSRange *ranges, const size_t length |
| 172 | + } |
| 173 | + } |
| 174 | + |
| 175 | ++static bool on_parser_progress(TSParseState *state) |
| 176 | ++{ |
| 177 | ++ TSLuaParserCallbackPayload *payload = state->payload; |
| 178 | ++ uint64_t parse_time = os_hrtime() - payload->parse_start_time; |
| 179 | ++ return parse_time >= payload->timeout_threshold_ns; |
| 180 | ++} |
| 181 | ++ |
| 182 | + static int parser_parse(lua_State *L) |
| 183 | + { |
| 184 | + TSParser *p = parser_check(L, 1); |
| 185 | +@@ -524,7 +536,17 @@ static int parser_parse(lua_State *L) |
| 186 | + } |
| 187 | + |
| 188 | + input = (TSInput){ (void *)buf, input_cb, TSInputEncodingUTF8, NULL }; |
| 189 | +- new_tree = ts_parser_parse(p, old_tree, input); |
| 190 | ++ if (!lua_isnil(L, 5)) { |
| 191 | ++ uint64_t timeout_ns = (uint64_t)lua_tointeger(L, 5); |
| 192 | ++ TSLuaParserCallbackPayload payload = |
| 193 | ++ (TSLuaParserCallbackPayload){ .parse_start_time = os_hrtime(), |
| 194 | ++ .timeout_threshold_ns = timeout_ns }; |
| 195 | ++ TSParseOptions parse_options = { .payload = &payload, |
| 196 | ++ .progress_callback = on_parser_progress }; |
| 197 | ++ new_tree = ts_parser_parse_with_options(p, old_tree, input, parse_options); |
| 198 | ++ } else { |
| 199 | ++ new_tree = ts_parser_parse(p, old_tree, input); |
| 200 | ++ } |
| 201 | + |
| 202 | + break; |
| 203 | + |
| 204 | +@@ -534,12 +556,11 @@ static int parser_parse(lua_State *L) |
| 205 | + |
| 206 | + bool include_bytes = (lua_gettop(L) >= 4) && lua_toboolean(L, 4); |
| 207 | + |
| 208 | +- // Sometimes parsing fails (timeout, or wrong parser ABI) |
| 209 | +- // In those case, just return an error. |
| 210 | + if (!new_tree) { |
| 211 | +- if (ts_parser_timeout_micros(p) == 0) { |
| 212 | +- // No timeout set, must have had an error |
| 213 | +- return luaL_error(L, "An error occurred when parsing."); |
| 214 | ++ // Sometimes parsing fails (no language was set, or it was set to one with an incompatible ABI) |
| 215 | ++ // In those cases, just return an error. |
| 216 | ++ if (!ts_parser_language(p)) { |
| 217 | ++ return luaL_error(L, "Language was unset, or has an incompatible ABI."); |
| 218 | + } |
| 219 | + return 0; |
| 220 | + } |
| 221 | +@@ -670,26 +691,6 @@ static int parser_get_ranges(lua_State *L) |
| 222 | + return 1; |
| 223 | + } |
| 224 | + |
| 225 | +-static int parser_set_timeout(lua_State *L) |
| 226 | +-{ |
| 227 | +- TSParser *p = parser_check(L, 1); |
| 228 | +- |
| 229 | +- if (lua_gettop(L) < 2) { |
| 230 | +- luaL_error(L, "integer expected"); |
| 231 | +- } |
| 232 | +- |
| 233 | +- uint32_t timeout = (uint32_t)luaL_checkinteger(L, 2); |
| 234 | +- ts_parser_set_timeout_micros(p, timeout); |
| 235 | +- return 0; |
| 236 | +-} |
| 237 | +- |
| 238 | +-static int parser_get_timeout(lua_State *L) |
| 239 | +-{ |
| 240 | +- TSParser *p = parser_check(L, 1); |
| 241 | +- lua_pushinteger(L, (lua_Integer)ts_parser_timeout_micros(p)); |
| 242 | +- return 1; |
| 243 | +-} |
| 244 | +- |
| 245 | + static void logger_cb(void *payload, TSLogType logtype, const char *s) |
| 246 | + { |
| 247 | + TSLuaLoggerOpts *opts = (TSLuaLoggerOpts *)payload; |
0 commit comments