Restore external scanner's state only after repositioning lexer

Also, properly identify the leaf node with the external token state

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) {

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 *);

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;
       }

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");

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;