diff --git a/src/runtime/lexer.c b/src/runtime/lexer.c index acf394bb..f7ebf042 100644 --- a/src/runtime/lexer.c +++ b/src/runtime/lexer.c @@ -88,6 +88,8 @@ void ts_lexer_init(Lexer *self) { .payload = NULL, .log = NULL }, + .needs_to_restore_external_scanner = false, + .last_external_token_end_byte = 0, }; ts_lexer_reset(self, length_zero()); } @@ -110,12 +112,22 @@ static inline void ts_lexer__reset(Lexer *self, Length position) { void ts_lexer_set_input(Lexer *self, TSInput input) { self->input = input; ts_lexer__reset(self, length_zero()); + self->needs_to_restore_external_scanner = false; + self->last_external_token_end_byte = 0; } void ts_lexer_reset(Lexer *self, Length position) { - if (!length_eq(position, self->current_position)) + if (position.bytes > self->current_position.bytes) { + self->needs_to_restore_external_scanner = true; + self->last_external_token_end_byte = 0; ts_lexer__reset(self, position); - return; + } else if (position.bytes < self->current_position.bytes) { + if (position.bytes < self->last_external_token_end_byte) { + self->needs_to_restore_external_scanner = true; + self->last_external_token_end_byte = 0; + } + ts_lexer__reset(self, position); + } } void ts_lexer_start(Lexer *self) { diff --git a/src/runtime/lexer.h b/src/runtime/lexer.h index 682c3f93..76d863c4 100644 --- a/src/runtime/lexer.h +++ b/src/runtime/lexer.h @@ -25,6 +25,8 @@ typedef struct { TSInput input; TSLogger logger; char debug_buffer[TS_DEBUG_BUFFER_SIZE]; + bool needs_to_restore_external_scanner; + uint32_t last_external_token_end_byte; } Lexer; void ts_lexer_init(Lexer *); diff --git a/src/runtime/parser.c b/src/runtime/parser.c index f5b08f82..b503bea0 100644 --- a/src/runtime/parser.c +++ b/src/runtime/parser.c @@ -217,8 +217,11 @@ static StackIterateAction parser__restore_external_scanner_callback( Parser *self = payload; if (tree_count > 0) { Tree *tree = *array_back(trees); - if (tree->has_external_token_state && tree->child_count == 0) { - self->language->external_scanner.deserialize(self->external_scanner_payload, tree->external_token_state); + if (tree->has_external_token_state) { + self->language->external_scanner.deserialize( + self->external_scanner_payload, + *ts_tree_last_external_token_state(tree) + ); return StackIterateStop; } } else if (is_done) { @@ -230,6 +233,7 @@ static StackIterateAction parser__restore_external_scanner_callback( } static void parser__restore_external_scanner(Parser *self, StackVersion version) { + if (!self->lexer.needs_to_restore_external_scanner) return; StackPopResult pop = ts_stack_iterate(self->stack, version, parser__restore_external_scanner_callback, self); if (pop.slices.size > 0) { StackSlice slice = pop.slices.contents[0]; @@ -269,6 +273,8 @@ static Tree *parser__lex(Parser *self, StackVersion version) { ts_lexer_start(&self->lexer); if (self->language->external_scanner.scan(self->external_scanner_payload, &self->lexer.data, external_tokens)) { + self->lexer.last_external_token_end_byte = self->lexer.current_position.bytes; + self->lexer.needs_to_restore_external_scanner = false; found_external_token = true; break; } diff --git a/src/runtime/tree.c b/src/runtime/tree.c index dfc45c0c..ebcca441 100644 --- a/src/runtime/tree.c +++ b/src/runtime/tree.c @@ -379,6 +379,19 @@ void ts_tree_edit(Tree *self, const TSInputEdit *edit) { } } +const TSExternalTokenState *ts_tree_last_external_token_state(const Tree *tree) { + while (tree->child_count > 0) { + for (uint32_t i = tree->child_count - 1; i + 1 > 0; i--) { + Tree *child = tree->children[i]; + if (child->has_external_token_state) { + tree = child; + break; + } + } + } + return &tree->external_token_state; +} + static size_t ts_tree__write_char_to_string(char *s, size_t n, int32_t c) { if (c == 0) return snprintf(s, n, "EOF"); diff --git a/src/runtime/tree.h b/src/runtime/tree.h index 425fac51..d5916e31 100644 --- a/src/runtime/tree.h +++ b/src/runtime/tree.h @@ -86,6 +86,7 @@ void ts_tree_assign_parents(Tree *, TreePath *); void ts_tree_edit(Tree *, const TSInputEdit *edit); char *ts_tree_string(const Tree *, const TSLanguage *, bool include_all); void ts_tree_print_dot_graph(const Tree *, const TSLanguage *, FILE *); +const TSExternalTokenState *ts_tree_last_external_token_state(const Tree *); static inline uint32_t ts_tree_total_bytes(const Tree *self) { return self->padding.bytes + self->size.bytes;