From f85230cdfd324a5943e3e905c0dd5b464d6f3409 Mon Sep 17 00:00:00 2001 From: Amaan Qureshi Date: Sun, 1 Sep 2024 04:46:05 -0400 Subject: [PATCH] feat: bump core library and add `QueryCursor` timeout functions --- include/tree_sitter/api.h | 16 ++++++++++++++++ query.go | 15 +++++++++++++++ src/parser.c | 4 ++-- src/query.c | 38 ++++++++++++++++++++++++++++++++++++-- src/tree_cursor.c | 5 +++-- tree-sitter | 2 +- 6 files changed, 73 insertions(+), 7 deletions(-) diff --git a/include/tree_sitter/api.h b/include/tree_sitter/api.h index c1fbad2..5ea845f 100644 --- a/include/tree_sitter/api.h +++ b/include/tree_sitter/api.h @@ -983,6 +983,22 @@ bool ts_query_cursor_did_exceed_match_limit(const TSQueryCursor *self); uint32_t ts_query_cursor_match_limit(const TSQueryCursor *self); void ts_query_cursor_set_match_limit(TSQueryCursor *self, uint32_t limit); +/** + * Set the maximum duration in microseconds that query execution should be allowed to + * take before halting. + * + * If query execution takes longer than this, it will halt early, returning NULL. + * See [`ts_query_cursor_next_match`] or [`ts_query_cursor_next_capture`] for more information. + */ +void ts_query_cursor_set_timeout_micros(TSQueryCursor *self, uint64_t timeout_micros); + +/** + * Get the duration in microseconds that query execution is allowed to take. + * + * This is set via [`ts_query_cursor_set_timeout_micros`]. + */ +uint64_t ts_query_cursor_timeout_micros(const TSQueryCursor *self); + /** * Set the range of bytes or (row, column) positions in which the query * will be executed. diff --git a/query.go b/query.go index 0fec7aa..d99a657 100644 --- a/query.go +++ b/query.go @@ -686,6 +686,21 @@ func (qc *QueryCursor) SetMatchLimit(limit uint) { C.ts_query_cursor_set_match_limit(qc._inner, C.uint32_t(limit)) } +// Set the maximum duration in microseconds that query execution should be allowed to +// take before halting. +// +// If query execution takes longer than this, it will halt early, returning None. +func (qc *QueryCursor) SetTimeoutMicros(timeoutMicros uint64) { + C.ts_query_cursor_set_timeout_micros(qc._inner, C.uint64_t(timeoutMicros)) +} + +// Get the duration in microseconds that query execution is allowed to take. +// +// This is set via [QueryCursor.SetTimeoutMicros] +func (qc *QueryCursor) TimeoutMicros() uint64 { + return uint64(C.ts_query_cursor_timeout_micros(qc._inner)) +} + // Check if, on its last execution, this cursor exceeded its maximum number // of in-progress matches. func (qc *QueryCursor) DidExceedMatchLimit() bool { diff --git a/src/parser.c b/src/parser.c index 2927d82..5db2cf5 100644 --- a/src/parser.c +++ b/src/parser.c @@ -83,7 +83,7 @@ static const unsigned MAX_VERSION_COUNT = 6; static const unsigned MAX_VERSION_COUNT_OVERFLOW = 4; static const unsigned MAX_SUMMARY_DEPTH = 16; static const unsigned MAX_COST_DIFFERENCE = 16 * ERROR_COST_PER_SKIPPED_TREE; -static const unsigned OP_COUNT_PER_TIMEOUT_CHECK = 100; +static const unsigned OP_COUNT_PER_PARSER_TIMEOUT_CHECK = 100; typedef struct { Subtree token; @@ -1565,7 +1565,7 @@ static bool ts_parser__advance( // If a cancellation flag or a timeout was provided, then check every // time a fixed number of parse actions has been processed. - if (++self->operation_count == OP_COUNT_PER_TIMEOUT_CHECK) { + if (++self->operation_count == OP_COUNT_PER_PARSER_TIMEOUT_CHECK) { self->operation_count = 0; } if ( diff --git a/src/query.c b/src/query.c index c9e8fbd..4941f50 100644 --- a/src/query.c +++ b/src/query.c @@ -1,6 +1,7 @@ #include "tree_sitter/api.h" #include "./alloc.h" #include "./array.h" +#include "./clock.h" #include "./language.h" #include "./point.h" #include "./tree_cursor.h" @@ -312,6 +313,9 @@ struct TSQueryCursor { TSPoint start_point; TSPoint end_point; uint32_t next_state_id; + TSClock end_clock; + TSDuration timeout_duration; + unsigned operation_count; bool on_visible_node; bool ascending; bool halted; @@ -322,6 +326,7 @@ static const TSQueryError PARENT_DONE = -1; static const uint16_t PATTERN_DONE_MARKER = UINT16_MAX; static const uint16_t NONE = UINT16_MAX; static const TSSymbol WILDCARD_SYMBOL = 0; +static const unsigned OP_COUNT_PER_QUERY_TIMEOUT_CHECK = 100; /********** * Stream @@ -2986,6 +2991,9 @@ TSQueryCursor *ts_query_cursor_new(void) { .start_point = {0, 0}, .end_point = POINT_MAX, .max_start_depth = UINT32_MAX, + .timeout_duration = 0, + .end_clock = clock_null(), + .operation_count = 0, }; array_reserve(&self->states, 8); array_reserve(&self->finished_states, 8); @@ -3012,6 +3020,14 @@ void ts_query_cursor_set_match_limit(TSQueryCursor *self, uint32_t limit) { self->capture_list_pool.max_capture_list_count = limit; } +uint64_t ts_query_cursor_timeout_micros(const TSQueryCursor *self) { + return duration_to_micros(self->timeout_duration); +} + +void ts_query_cursor_set_timeout_micros(TSQueryCursor *self, uint64_t timeout_micros) { + self->timeout_duration = duration_from_micros(timeout_micros); +} + #ifdef DEBUG_EXECUTE_QUERY #define LOG(...) fprintf(stderr, __VA_ARGS__) #else @@ -3023,7 +3039,7 @@ void ts_query_cursor_exec( const TSQuery *query, TSNode node ) { - if (query) { + if (query) { LOG("query steps:\n"); for (unsigned i = 0; i < query->steps.size; i++) { QueryStep *step = &query->steps.contents[i]; @@ -3060,6 +3076,12 @@ void ts_query_cursor_exec( self->halted = false; self->query = query; self->did_exceed_match_limit = false; + self->operation_count = 0; + if (self->timeout_duration) { + self->end_clock = clock_after(clock_now(), self->timeout_duration); + } else { + self->end_clock = clock_null(); + } } void ts_query_cursor_set_byte_range( @@ -3456,7 +3478,19 @@ static inline bool ts_query_cursor__advance( } } - if (did_match || self->halted) return did_match; + if (++self->operation_count == OP_COUNT_PER_QUERY_TIMEOUT_CHECK) { + self->operation_count = 0; + } + if ( + did_match || + self->halted || + ( + self->operation_count == 0 && + !clock_is_null(self->end_clock) && clock_is_gt(clock_now(), self->end_clock) + ) + ) { + return did_match; + } // Exit the current node. if (self->ascending) { diff --git a/src/tree_cursor.c b/src/tree_cursor.c index ddd7d66..2441666 100644 --- a/src/tree_cursor.c +++ b/src/tree_cursor.c @@ -475,8 +475,9 @@ uint32_t ts_tree_cursor_current_descendant_index(const TSTreeCursor *_self) { TSNode ts_tree_cursor_current_node(const TSTreeCursor *_self) { const TreeCursor *self = (const TreeCursor *)_self; TreeCursorEntry *last_entry = array_back(&self->stack); - TSSymbol alias_symbol = self->root_alias_symbol; - if (self->stack.size > 1 && !ts_subtree_extra(*last_entry->subtree)) { + bool is_extra = ts_subtree_extra(*last_entry->subtree); + TSSymbol alias_symbol = is_extra ? 0 : self->root_alias_symbol; + if (self->stack.size > 1 && !is_extra) { TreeCursorEntry *parent_entry = &self->stack.contents[self->stack.size - 2]; alias_symbol = ts_language_alias_at( self->tree->language, diff --git a/tree-sitter b/tree-sitter index 12fb318..ee06325 160000 --- a/tree-sitter +++ b/tree-sitter @@ -1 +1 @@ -Subproject commit 12fb31826b8469cc7b9788e72bceee5af1cf0977 +Subproject commit ee06325f67beba7413a4179018a87a417e2e4a49