diff --git a/src/runtime/document.c b/src/runtime/document.c
index 64677cb4..bae6ca6f 100644
--- a/src/runtime/document.c
+++ b/src/runtime/document.c
@@ -24,8 +24,7 @@ error:
 
 void ts_document_free(TSDocument *self) {
   parser_destroy(&self->parser);
-  if (self->tree)
-    ts_tree_release(self->tree);
+  if (self->tree) ts_tree_release(self->tree);
   ts_document_set_input(self, (TSInput){
     NULL,
     NULL,
diff --git a/src/runtime/lexer.c b/src/runtime/lexer.c
index 96bc2d13..3e9f9f70 100644
--- a/src/runtime/lexer.c
+++ b/src/runtime/lexer.c
@@ -110,7 +110,7 @@ void ts_lexer_init(Lexer *self) {
       .payload = NULL,
       .log = NULL
     },
-    .last_external_token_state = NULL,
+    .last_external_token = NULL,
   };
   ts_lexer_reset(self, length_zero());
 }
@@ -134,7 +134,7 @@ 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->last_external_token_state = NULL;
+  ts_lexer_set_last_external_token(self, NULL);
 }
 
 void ts_lexer_reset(Lexer *self, Length position) {
@@ -157,3 +157,9 @@ void ts_lexer_start(Lexer *self) {
 void ts_lexer_advance_to_end(Lexer *self) {
   while (self->data.lookahead != 0) ts_lexer__advance(self, false);
 }
+
+void ts_lexer_set_last_external_token(Lexer *self, Tree *token) {
+  if (token) ts_tree_retain(token);
+  if (self->last_external_token) ts_tree_release(self->last_external_token);
+  self->last_external_token = token;
+}
diff --git a/src/runtime/lexer.h b/src/runtime/lexer.h
index f6f127f5..5ae977ad 100644
--- a/src/runtime/lexer.h
+++ b/src/runtime/lexer.h
@@ -8,6 +8,7 @@ extern "C" {
 #include "tree_sitter/parser.h"
 #include "tree_sitter/runtime.h"
 #include "runtime/length.h"
+#include "runtime/tree.h"
 
 #define TS_DEBUG_BUFFER_SIZE 512
 
@@ -25,7 +26,7 @@ typedef struct {
   TSInput input;
   TSLogger logger;
   char debug_buffer[TS_DEBUG_BUFFER_SIZE];
-  const TSExternalTokenState *last_external_token_state;
+  Tree *last_external_token;
 } Lexer;
 
 void ts_lexer_init(Lexer *);
@@ -33,6 +34,7 @@ void ts_lexer_set_input(Lexer *, TSInput);
 void ts_lexer_reset(Lexer *, Length);
 void ts_lexer_start(Lexer *);
 void ts_lexer_advance_to_end(Lexer *);
+void ts_lexer_set_last_external_token(Lexer *, Tree *);
 
 #ifdef __cplusplus
 }
diff --git a/src/runtime/parser.c b/src/runtime/parser.c
index c678d881..492efc18 100644
--- a/src/runtime/parser.c
+++ b/src/runtime/parser.c
@@ -56,8 +56,7 @@ typedef struct {
   TSSymbol lookahead_symbol;
 } SkipPrecedingTreesSession;
 
-static void parser__push(Parser *self, StackVersion version, Tree *tree,
-                         TSStateId state) {
+static void parser__push(Parser *self, StackVersion version, Tree *tree, TSStateId state) {
   ts_stack_push(self->stack, version, tree, false, state);
   ts_tree_release(tree);
 }
@@ -201,19 +200,16 @@ static CondenseResult parser__condense_stack(Parser *self) {
   return result;
 }
 
-static void parser__restore_external_scanner(Parser *self, StackVersion version) {
-  const TSExternalTokenState *state = ts_stack_external_token_state(self->stack, version);
-  if (self->lexer.last_external_token_state != state) {
-    LOG("restore_external_scanner");
-    self->lexer.last_external_token_state = state;
-    if (state) {
-      self->language->external_scanner.deserialize(
-        self->external_scanner_payload,
-        *state
-      );
-    } else {
-      self->language->external_scanner.reset(self->external_scanner_payload);
-    }
+static void parser__restore_external_scanner(Parser *self, Tree *external_token) {
+  LOG("restore_external_scanner");
+  ts_lexer_set_last_external_token(&self->lexer, external_token);
+  if (external_token) {
+    self->language->external_scanner.deserialize(
+      self->external_scanner_payload,
+      external_token->external_token_state
+    );
+  } else {
+    self->language->external_scanner.reset(self->external_scanner_payload);
   }
 }
 
@@ -239,19 +235,33 @@ static Tree *parser__lex(Parser *self, StackVersion version) {
     if (valid_external_tokens) {
       LOG("lex_external state:%d, row:%u, column:%u", lex_mode.external_lex_state,
           current_position.extent.row, current_position.extent.column);
-      parser__restore_external_scanner(self, version);
+
+      Tree *external_token = ts_stack_last_external_token(self->stack, version);
+      if (!ts_tree_external_token_state_eq(self->lexer.last_external_token, external_token)) {
+        parser__restore_external_scanner(self, external_token);
+      }
+
       ts_lexer_start(&self->lexer);
-      if (self->language->external_scanner.scan(self->external_scanner_payload,
-                                                &self->lexer.data, valid_external_tokens)) {
+      found_external_token = self->language->external_scanner.scan(
+        self->external_scanner_payload,
+        &self->lexer.data,
+        valid_external_tokens
+      );
+
+      if (found_external_token) {
         if (length_has_unknown_chars(self->lexer.token_end_position)) {
           self->lexer.token_end_position = self->lexer.current_position;
         }
-        if (lex_mode.lex_state != 0 ||
-            self->lexer.token_end_position.bytes > current_position.bytes) {
-          found_external_token = true;
+
+        // Don't allow zero-length external tokens durring error recovery.
+        if (lex_mode.lex_state == 0 && self->lexer.token_end_position.bytes <= current_position.bytes) {
+          parser__restore_external_scanner(self, external_token);
+          found_external_token = false;
+        } else {
           break;
         }
       }
+
       ts_lexer_reset(&self->lexer, current_position);
     }
 
@@ -314,10 +324,9 @@ static Tree *parser__lex(Parser *self, StackVersion version) {
 
     if (found_external_token) {
       result->has_external_tokens = true;
-      result->has_external_token_state = true;
       memset(result->external_token_state, 0, sizeof(TSExternalTokenState));
       self->language->external_scanner.serialize(self->external_scanner_payload, result->external_token_state);
-      self->lexer.last_external_token_state = &result->external_token_state;
+      ts_lexer_set_last_external_token(&self->lexer, result);
     }
   }
 
@@ -330,7 +339,7 @@ static Tree *parser__lex(Parser *self, StackVersion version) {
 }
 
 static void parser__clear_cached_token(Parser *self) {
-  ts_tree_release(self->cached_token);
+  if (self->cached_token) ts_tree_release(self->cached_token);
   self->cached_token = NULL;
 }
 
@@ -373,9 +382,9 @@ static Tree *parser__get_lookahead(Parser *self, StackVersion version,
       continue;
     }
 
-    if (!ts_external_token_state_eq(
-          reusable_node->preceding_external_token_state,
-          ts_stack_external_token_state(self->stack, version))) {
+    if (!ts_tree_external_token_state_eq(
+          reusable_node->preceding_external_token,
+          ts_stack_last_external_token(self->stack, version))) {
       LOG("cant_reuse_external_tokens tree:%s, size:%u",
           SYM_NAME(reusable_node->tree->symbol),
           reusable_node->tree->size.bytes);
@@ -475,9 +484,10 @@ static void parser__shift(Parser *self, StackVersion version, TSStateId state,
 
   bool is_pending = lookahead->child_count > 0;
   ts_stack_push(self->stack, version, lookahead, is_pending, state);
-  if (lookahead->has_external_token_state) {
-    ts_stack_set_external_token_state(
-      self->stack, version, ts_tree_last_external_token_state(lookahead));
+  if (lookahead->has_external_tokens) {
+    ts_stack_set_last_external_token(
+      self->stack, version, ts_tree_last_external_token(lookahead)
+    );
   }
   ts_tree_release(lookahead);
 }
@@ -808,7 +818,7 @@ static void parser__start(Parser *self, TSInput input, Tree *previous_tree) {
   ts_lexer_set_input(&self->lexer, input);
   ts_stack_clear(self->stack);
   self->reusable_node = reusable_node_new(previous_tree);
-  self->cached_token = NULL;
+  parser__clear_cached_token(self);
   self->finished_tree = NULL;
 }
 
@@ -843,12 +853,17 @@ static void parser__accept(Parser *self, StackVersion version,
       }
     }
 
-    if (parser__select_tree(self, self->finished_tree, root)) {
-      ts_tree_release(self->finished_tree);
-      assert(root && root->ref_count > 0);
-      self->finished_tree = root;
+    assert(root && root->ref_count > 0);
+
+    if (self->finished_tree) {
+      if (parser__select_tree(self, self->finished_tree, root)) {
+        ts_tree_release(self->finished_tree);
+        self->finished_tree = root;
+      } else {
+        ts_tree_release(root);
+      }
     } else {
-      ts_tree_release(root);
+      self->finished_tree = root;
     }
   }
 
@@ -1113,8 +1128,9 @@ static void parser__advance(Parser *self, StackVersion version,
 
           parser__shift(self, version, next_state, lookahead, extra);
 
-          if (lookahead == reusable_node->tree)
+          if (lookahead == reusable_node->tree) {
             reusable_node_pop(reusable_node);
+          }
 
           ts_tree_release(lookahead);
           return;
@@ -1163,8 +1179,9 @@ static void parser__advance(Parser *self, StackVersion version,
           }
 
           parser__recover(self, version, action.params.to_state, lookahead);
-          if (lookahead == reusable_node->tree)
+          if (lookahead == reusable_node->tree) {
             reusable_node_pop(reusable_node);
+          }
           ts_tree_release(lookahead);
           return;
         }
@@ -1227,6 +1244,7 @@ void parser_destroy(Parser *self) {
     array_delete(&self->tree_path1);
   if (self->tree_path2.contents)
     array_delete(&self->tree_path2);
+  ts_lexer_set_last_external_token(&self->lexer, NULL);
   parser_set_language(self, NULL);
 }
 
diff --git a/src/runtime/reusable_node.h b/src/runtime/reusable_node.h
index b9777638..2c7bc58f 100644
--- a/src/runtime/reusable_node.h
+++ b/src/runtime/reusable_node.h
@@ -3,24 +3,21 @@
 typedef struct {
   Tree *tree;
   uint32_t byte_index;
-  bool has_preceding_external_token;
-  const TSExternalTokenState *preceding_external_token_state;
+  Tree *preceding_external_token;
 } ReusableNode;
 
 static inline ReusableNode reusable_node_new(Tree *tree) {
   return (ReusableNode){
     .tree = tree,
     .byte_index = 0,
-    .has_preceding_external_token = false,
-    .preceding_external_token_state = NULL,
+    .preceding_external_token = NULL,
   };
 }
 
 static inline void reusable_node_pop(ReusableNode *self) {
   self->byte_index += ts_tree_total_bytes(self->tree);
   if (self->tree->has_external_tokens) {
-    self->has_preceding_external_token = true;
-    self->preceding_external_token_state = ts_tree_last_external_token_state(self->tree);
+    self->preceding_external_token = ts_tree_last_external_token(self->tree);
   }
 
   while (self->tree) {
diff --git a/src/runtime/stack.c b/src/runtime/stack.c
index 7e52335a..a8fafa43 100644
--- a/src/runtime/stack.c
+++ b/src/runtime/stack.c
@@ -51,7 +51,7 @@ typedef struct {
   StackNode *node;
   bool is_halted;
   unsigned push_count;
-  const TSExternalTokenState *external_token_state;
+  Tree *last_external_token;
 } StackHead;
 
 struct Stack {
@@ -70,14 +70,14 @@ static void stack_node_retain(StackNode *self) {
 }
 
 static void stack_node_release(StackNode *self, StackNodeArray *pool) {
-  if (!self)
-    return;
+  if (!self) return;
   assert(self->ref_count != 0);
   self->ref_count--;
   if (self->ref_count == 0) {
     for (int i = 0; i < self->link_count; i++) {
-      if (self->links[i].tree)
+      if (self->links[i].tree) {
         ts_tree_release(self->links[i].tree);
+      }
       stack_node_release(self->links[i].node, pool);
     }
 
@@ -173,21 +173,22 @@ static void stack_node_add_link(StackNode *self, StackLink link) {
 
 static StackVersion ts_stack__add_version(Stack *self, StackNode *node,
                                           unsigned push_count,
-                                          const TSExternalTokenState *external_token_state) {
+                                          Tree *last_external_token) {
   StackHead head = {
     .node = node,
     .is_halted = false,
     .push_count = push_count,
-    .external_token_state = external_token_state,
+    .last_external_token = last_external_token,
   };
   array_push(&self->heads, head);
   stack_node_retain(node);
+  if (last_external_token) ts_tree_retain(last_external_token);
   return (StackVersion)(self->heads.size - 1);
 }
 
 static void ts_stack__add_slice(Stack *self, StackNode *node, TreeArray *trees,
                                 unsigned push_count,
-                                const TSExternalTokenState *external_token_state) {
+                                Tree *last_external_token) {
   for (uint32_t i = self->slices.size - 1; i + 1 > 0; i--) {
     StackVersion version = self->slices.contents[i].version;
     if (self->heads.contents[version].node == node) {
@@ -197,7 +198,7 @@ static void ts_stack__add_slice(Stack *self, StackNode *node, TreeArray *trees,
     }
   }
 
-  StackVersion version = ts_stack__add_version(self, node, push_count, external_token_state);
+  StackVersion version = ts_stack__add_version(self, node, push_count, last_external_token);
   StackSlice slice = { *trees, version };
   array_push(&self->slices, slice);
 }
@@ -209,7 +210,7 @@ INLINE StackPopResult stack__iter(Stack *self, StackVersion version,
 
   StackHead *head = array_get(&self->heads, version);
   unsigned push_count = head->push_count;
-  const TSExternalTokenState *external_token_state = head->external_token_state;
+  Tree *last_external_token = head->last_external_token;
   Iterator iterator = {
     .node = head->node,
     .trees = array_new(),
@@ -238,7 +239,7 @@ INLINE StackPopResult stack__iter(Stack *self, StackVersion version,
           ts_tree_array_copy(trees, &trees);
         array_reverse(&trees);
         ts_stack__add_slice(self, node, &trees, push_count + iterator->push_count,
-                            external_token_state);
+                            last_external_token);
       }
 
       if (should_stop) {
@@ -315,8 +316,11 @@ void ts_stack_delete(Stack *self) {
   if (self->iterators.contents)
     array_delete(&self->iterators);
   stack_node_release(self->base_node, &self->node_pool);
-  for (uint32_t i = 0; i < self->heads.size; i++)
-    stack_node_release(self->heads.contents[i].node, &self->node_pool);
+  for (uint32_t i = 0; i < self->heads.size; i++) {
+    StackHead head = self->heads.contents[i];
+    stack_node_release(head.node, &self->node_pool);
+    if (head.last_external_token) ts_tree_release(head.last_external_token);
+  }
   array_clear(&self->heads);
   if (self->node_pool.contents) {
     for (uint32_t i = 0; i < self->node_pool.size; i++)
@@ -348,12 +352,15 @@ void ts_stack_decrease_push_count(Stack *self, StackVersion version,
   array_get(&self->heads, version)->push_count -= decrement;
 }
 
-const TSExternalTokenState *ts_stack_external_token_state(const Stack *self, StackVersion version) {
-  return array_get(&self->heads, version)->external_token_state;
+Tree *ts_stack_last_external_token(const Stack *self, StackVersion version) {
+  return array_get(&self->heads, version)->last_external_token;
 }
 
-void ts_stack_set_external_token_state(Stack *self, StackVersion version, const TSExternalTokenState *state) {
-  array_get(&self->heads, version)->external_token_state = state;
+void ts_stack_set_last_external_token(Stack *self, StackVersion version, Tree *token) {
+  StackHead *head = array_get(&self->heads, version);
+  if (token) ts_tree_retain(token);
+  if (head->last_external_token) ts_tree_release(head->last_external_token);
+  head->last_external_token = token;
 }
 
 ErrorStatus ts_stack_error_status(const Stack *self, StackVersion version) {
@@ -475,15 +482,18 @@ StackPopResult ts_stack_pop_all(Stack *self, StackVersion version) {
 }
 
 void ts_stack_remove_version(Stack *self, StackVersion version) {
-  StackNode *node = array_get(&self->heads, version)->node;
-  stack_node_release(node, &self->node_pool);
+  StackHead *head = array_get(&self->heads, version);
+  if (head->last_external_token) ts_tree_release(head->last_external_token);
+  stack_node_release(head->node, &self->node_pool);
   array_erase(&self->heads, version);
 }
 
 void ts_stack_renumber_version(Stack *self, StackVersion v1, StackVersion v2) {
   assert(v2 < v1);
   assert((uint32_t)v1 < self->heads.size);
-  stack_node_release(self->heads.contents[v2].node, &self->node_pool);
+  StackHead *head_to_remove = &self->heads.contents[v2];
+  stack_node_release(head_to_remove->node, &self->node_pool);
+  if (head_to_remove->last_external_token) ts_tree_release(head_to_remove->last_external_token);
   self->heads.contents[v2] = self->heads.contents[v1];
   array_erase(&self->heads, v1);
 }
@@ -491,7 +501,9 @@ void ts_stack_renumber_version(Stack *self, StackVersion v1, StackVersion v2) {
 StackVersion ts_stack_copy_version(Stack *self, StackVersion version) {
   assert(version < self->heads.size);
   array_push(&self->heads, self->heads.contents[version]);
-  stack_node_retain(array_back(&self->heads)->node);
+  StackHead *head = array_back(&self->heads);
+  stack_node_retain(head->node);
+  if (head->last_external_token) ts_tree_retain(head->last_external_token);
   return self->heads.size - 1;
 }
 
@@ -505,7 +517,7 @@ bool ts_stack_merge(Stack *self, StackVersion version, StackVersion new_version)
       new_node->position.chars == node->position.chars &&
       new_node->error_count == node->error_count &&
       new_node->error_cost == node->error_cost &&
-      ts_external_token_state_eq(new_head->external_token_state, head->external_token_state)) {
+      ts_tree_external_token_state_eq(new_head->last_external_token, head->last_external_token)) {
     for (uint32_t j = 0; j < new_node->link_count; j++)
       stack_node_add_link(node, new_node->links[j]);
     if (new_head->push_count > head->push_count)
@@ -527,8 +539,11 @@ bool ts_stack_is_halted(Stack *self, StackVersion version) {
 
 void ts_stack_clear(Stack *self) {
   stack_node_retain(self->base_node);
-  for (uint32_t i = 0; i < self->heads.size; i++)
-    stack_node_release(self->heads.contents[i].node, &self->node_pool);
+  for (uint32_t i = 0; i < self->heads.size; i++) {
+    StackHead head = self->heads.contents[i];
+    stack_node_release(head.node, &self->node_pool);
+    if (head.last_external_token) ts_tree_release(head.last_external_token);
+  }
   array_clear(&self->heads);
   array_push(&self->heads, ((StackHead){
     self->base_node,
@@ -561,8 +576,8 @@ bool ts_stack_print_dot_graph(Stack *self, const char **symbol_names, FILE *f) {
       "labeltooltip=\"push_count: %u",
       i, head->node, i, head->push_count);
 
-    if (head->external_token_state) {
-      const TSExternalTokenState *s = head->external_token_state;
+    if (head->last_external_token) {
+      const TSExternalTokenState *s = &head->last_external_token->external_token_state;
       fprintf(f,
         "\nexternal_token_state: "
         "%2X %2X %2X %2X %2X %2X %2X %2X %2X %2X %2X %2X %2X %2X %2X %2X",
diff --git a/src/runtime/stack.h b/src/runtime/stack.h
index 2e88d72a..e64332dc 100644
--- a/src/runtime/stack.h
+++ b/src/runtime/stack.h
@@ -67,9 +67,9 @@ unsigned ts_stack_push_count(const Stack *, StackVersion);
 
 void ts_stack_decrease_push_count(Stack *, StackVersion, unsigned);
 
-const TSExternalTokenState *ts_stack_external_token_state(const Stack *, StackVersion);
+Tree *ts_stack_last_external_token(const Stack *, StackVersion);
 
-void ts_stack_set_external_token_state(Stack *, StackVersion, const TSExternalTokenState *);
+void ts_stack_set_last_external_token(Stack *, StackVersion, Tree *);
 
 /*
  *  Get the position at the top of the given version of the stack. If the stack
diff --git a/src/runtime/tree.c b/src/runtime/tree.c
index 5d8e4019..c11843fb 100644
--- a/src/runtime/tree.c
+++ b/src/runtime/tree.c
@@ -27,6 +27,7 @@ Tree *ts_tree_make_leaf(TSSymbol sym, Length padding, Length size,
     .named = metadata.named,
     .has_changes = false,
     .first_leaf.symbol = sym,
+    .has_external_tokens = false,
   };
   return result;
 }
@@ -47,8 +48,9 @@ bool ts_tree_array_copy(TreeArray self, TreeArray *dest) {
 }
 
 void ts_tree_array_delete(TreeArray *self) {
-  for (uint32_t i = 0; i < self->size; i++)
+  for (uint32_t i = 0; i < self->size; i++) {
     ts_tree_release(self->contents[i]);
+  }
   array_delete(self);
 }
 
@@ -148,7 +150,6 @@ void ts_tree_set_children(Tree *self, uint32_t child_count, Tree **children) {
   self->visible_child_count = 0;
   self->error_cost = 0;
   self->has_external_tokens = false;
-  self->has_external_token_state = false;
 
   for (uint32_t i = 0; i < child_count; i++) {
     Tree *child = children[i];
@@ -175,7 +176,6 @@ void ts_tree_set_children(Tree *self, uint32_t child_count, Tree **children) {
     }
 
     if (child->has_external_tokens) self->has_external_tokens = true;
-    if (child->has_external_token_state) self->has_external_token_state = true;
 
     if (child->symbol == ts_builtin_sym_error) {
       self->fragile_left = self->fragile_right = true;
@@ -235,17 +235,15 @@ void ts_tree_retain(Tree *self) {
 }
 
 void ts_tree_release(Tree *self) {
-  if (!self)
-    return;
-
 recur:
   assert(self->ref_count > 0);
   self->ref_count--;
 
   if (self->ref_count == 0) {
     if (self->child_count > 0) {
-      for (uint32_t i = 0; i < self->child_count - 1; i++)
+      for (uint32_t i = 0; i < self->child_count - 1; i++) {
         ts_tree_release(self->children[i]);
+      }
       Tree *last_child = self->children[self->child_count - 1];
       ts_free(self->children);
       ts_free(self);
@@ -311,12 +309,7 @@ bool ts_tree_tokens_eq(const Tree *self, const Tree *other) {
   if (self->padding.bytes != other->padding.bytes) return false;
   if (self->size.bytes != other->size.bytes) return false;
   if (self->extra != other->extra) return false;
-  if (self->has_external_token_state) {
-    if (!other->has_external_token_state) return false;
-    if (!ts_external_token_state_eq(&self->external_token_state, &other->external_token_state)) return false;
-  } else {
-    if (other->has_external_token_state) return false;
-  }
+  if (!ts_tree_external_token_state_eq(self, other)) return false;
   return true;
 }
 
@@ -451,19 +444,18 @@ void ts_tree_edit(Tree *self, const TSInputEdit *edit) {
   }
 }
 
-const TSExternalTokenState *ts_tree_last_external_token_state(const Tree *tree) {
+Tree *ts_tree_last_external_token(Tree *tree) {
+  if (!tree->has_external_tokens) return NULL;
   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) {
+      if (child->has_external_tokens) {
         tree = child;
         break;
-      } else if (child->has_external_tokens) {
-        return NULL;
       }
     }
   }
-  return &tree->external_token_state;
+  return tree;
 }
 
 static size_t ts_tree__write_char_to_string(char *s, size_t n, int32_t c) {
@@ -560,13 +552,13 @@ void ts_tree_print_dot_graph(const Tree *self, const TSLanguage *language,
   fprintf(f, "}\n");
 }
 
-bool ts_external_token_state_eq(const TSExternalTokenState *self,
-                                const TSExternalTokenState *other) {
-  if (self == other) {
-    return true;
-  } else if (!self || !other) {
-    return false;
-  } else {
-    return memcmp(self, other, sizeof(TSExternalTokenState)) == 0;
-  }
+bool ts_tree_external_token_state_eq(const Tree *self, const Tree *other) {
+  return self == other || (
+    self &&
+    other &&
+    self->has_external_tokens == other->has_external_tokens && (
+      !self->has_external_tokens ||
+      memcmp(&self->external_token_state, &other->external_token_state, sizeof(TSExternalTokenState)) == 0
+    )
+  );
 }
diff --git a/src/runtime/tree.h b/src/runtime/tree.h
index f205af97..8d1ce186 100644
--- a/src/runtime/tree.h
+++ b/src/runtime/tree.h
@@ -53,7 +53,6 @@ typedef struct Tree {
   bool fragile_right : 1;
   bool has_changes : 1;
   bool has_external_tokens : 1;
-  bool has_external_token_state : 1;
 } Tree;
 
 typedef struct {
@@ -90,7 +89,8 @@ 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 *);
+Tree *ts_tree_last_external_token(Tree *);
+bool ts_tree_external_token_state_eq(const Tree *, const Tree *);
 
 static inline uint32_t ts_tree_total_bytes(const Tree *self) {
   return self->padding.bytes + self->size.bytes;
@@ -109,8 +109,6 @@ static inline bool ts_tree_is_fragile(const Tree *tree) {
          ts_tree_total_bytes(tree) == 0;
 }
 
-bool ts_external_token_state_eq(const TSExternalTokenState *, const TSExternalTokenState *);
-
 #ifdef __cplusplus
 }
 #endif
diff --git a/test/runtime/stack_test.cc b/test/runtime/stack_test.cc
index bf0514b3..14bcbc7c 100644
--- a/test/runtime/stack_test.cc
+++ b/test/runtime/stack_test.cc
@@ -523,16 +523,19 @@ describe("Stack", [&]() {
   });
 
   describe("setting external token state", [&]() {
-    TSExternalTokenState external_token_state1, external_token_state2;
+    before_each([&]() {
+      trees[1]->external_token_state[0] = 'a';
+      trees[2]->external_token_state[0] = 'b';
+    });
 
     it("allows the state to be retrieved", [&]() {
-      AssertThat(ts_stack_external_token_state(stack, 0), Equals(nullptr));
+      AssertThat(ts_stack_last_external_token(stack, 0), Equals(nullptr));
 
-      ts_stack_set_external_token_state(stack, 0, &external_token_state1);
-      AssertThat(ts_stack_external_token_state(stack, 0), Equals(&external_token_state1));
+      ts_stack_set_last_external_token(stack, 0, trees[1]);
+      AssertThat(ts_stack_last_external_token(stack, 0), Equals(trees[1]));
 
       ts_stack_copy_version(stack, 0);
-      AssertThat(ts_stack_external_token_state(stack, 0), Equals(&external_token_state1));
+      AssertThat(ts_stack_last_external_token(stack, 0), Equals(trees[1]));
     });
 
     it("does not merge stack versions with different external token states", [&]() {
@@ -540,8 +543,8 @@ describe("Stack", [&]() {
       ts_stack_push(stack, 0, trees[0], false, 5);
       ts_stack_push(stack, 1, trees[0], false, 5);
 
-      ts_stack_set_external_token_state(stack, 0, &external_token_state1);
-      ts_stack_set_external_token_state(stack, 0, &external_token_state2);
+      ts_stack_set_last_external_token(stack, 0, trees[1]);
+      ts_stack_set_last_external_token(stack, 0, trees[1]);
 
       AssertThat(ts_stack_merge(stack, 0, 1), IsFalse());
     });
diff --git a/test/runtime/tree_test.cc b/test/runtime/tree_test.cc
index a3bdd563..d751c717 100644
--- a/test/runtime/tree_test.cc
+++ b/test/runtime/tree_test.cc
@@ -421,13 +421,12 @@ describe("Tree", []() {
     });
   });
 
-  describe("last_external_token_state", [&]() {
+  describe("last_external_token", [&]() {
     Length padding = {1, 1, {0, 1}};
     Length size = {2, 2, {0, 2}};
 
     auto make_external = [](Tree *tree) {
       tree->has_external_tokens = true;
-      tree->has_external_token_state = true;
       return tree;
     };
 
@@ -448,8 +447,8 @@ describe("Tree", []() {
         }), visible)),
       }), visible);
 
-      auto state = ts_tree_last_external_token_state(tree1);
-      AssertThat(state, Equals(&tree3->external_token_state));
+      auto token = ts_tree_last_external_token(tree1);
+      AssertThat(token, Equals(tree3));
     });
   });
 });