From 4b5dab2d18f26ac0393a0984045e4d4ca903ff91 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Fri, 15 Jun 2018 12:25:17 -0700
Subject: [PATCH] Convert keywords back to word token at runtime if needed

---
 src/runtime/parser.c  | 44 ++++++++++++++++++++++++++++++++++---------
 src/runtime/subtree.c |  3 ++-
 src/runtime/subtree.h |  1 +
 3 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/src/runtime/parser.c b/src/runtime/parser.c
index dcb55a62..6e2299d4 100644
--- a/src/runtime/parser.c
+++ b/src/runtime/parser.c
@@ -263,7 +263,7 @@ static bool ts_parser__can_reuse_first_leaf(TSParser *self, TSStateId state, con
   if (tree->first_leaf.lex_mode.lex_state == current_lex_mode.lex_state &&
       tree->first_leaf.lex_mode.external_lex_state == current_lex_mode.external_lex_state &&
       (tree->first_leaf.symbol != self->language->keyword_capture_token ||
-       tree->parse_state == state)) return true;
+       (!tree->is_keyword && tree->parse_state == state))) return true;
 
   // Empty tokens are not reusable in states with different lookaheads.
   if (tree->size.bytes == 0 && tree->symbol != ts_builtin_sym_end) return false;
@@ -382,6 +382,7 @@ static const Subtree *ts_parser__lex(TSParser *self, StackVersion version, TSSta
       self->lexer.token_start_position = self->lexer.token_end_position;
     }
 
+    bool is_keyword = false;
     TSSymbol symbol = self->lexer.data.result_symbol;
     Length padding = length_sub(self->lexer.token_start_position, start_position);
     Length size = length_sub(self->lexer.token_end_position, self->lexer.token_start_position);
@@ -397,11 +398,13 @@ static const Subtree *ts_parser__lex(TSParser *self, StackVersion version, TSSta
         self->lexer.token_end_position.bytes == end_byte &&
         ts_language_has_actions(self->language, parse_state, self->lexer.data.result_symbol)
       ) {
+        is_keyword = true;
         symbol = self->lexer.data.result_symbol;
       }
     }
 
     result = ts_subtree_new_leaf(&self->tree_pool, symbol, padding, size, self->language);
+    result->is_keyword = is_keyword;
 
     if (found_external_token) {
       result->has_external_tokens = true;
@@ -1200,18 +1203,41 @@ static void ts_parser__advance(TSParser *self, StackVersion version, bool allow_
     if (last_reduction_version != STACK_VERSION_NONE) {
       ts_stack_renumber_version(self->stack, last_reduction_version, version);
       LOG_STACK();
-    } else if (state == ERROR_STATE) {
+      state = ts_stack_state(self->stack, version);
+      ts_language_table_entry(self->language, state, lookahead->first_leaf.symbol, &table_entry);
+      continue;
+    }
+
+    if (lookahead->is_keyword && lookahead->symbol != self->language->keyword_capture_token) {
+      ts_language_table_entry(self->language, state, self->language->keyword_capture_token, &table_entry);
+      if (table_entry.action_count > 0) {
+        LOG(
+          "switch from_keyword:%s, to_word_token:%s",
+          SYM_NAME(lookahead->symbol),
+          SYM_NAME(self->language->keyword_capture_token)
+        );
+
+        Subtree *mutable_lookahead = ts_subtree_make_mut(&self->tree_pool, lookahead);
+        mutable_lookahead->symbol = self->language->keyword_capture_token;
+        mutable_lookahead->first_leaf.symbol = self->language->keyword_capture_token;
+        lookahead = mutable_lookahead;
+        continue;
+      }
+    }
+
+    if (state == ERROR_STATE) {
       ts_parser__recover(self, version, lookahead);
       return;
-    } else if (!ts_parser__breakdown_top_of_stack(self, version)) {
-      LOG("detect_error");
-      ts_stack_pause(self->stack, version, lookahead->first_leaf.symbol);
-      ts_subtree_release(&self->tree_pool, lookahead);
-      return;
     }
 
-    state = ts_stack_state(self->stack, version);
-    ts_language_table_entry(self->language, state, lookahead->first_leaf.symbol, &table_entry);
+    if (ts_parser__breakdown_top_of_stack(self, version)) {
+      continue;
+    }
+
+    LOG("detect_error");
+    ts_stack_pause(self->stack, version, lookahead->first_leaf.symbol);
+    ts_subtree_release(&self->tree_pool, lookahead);
+    return;
   }
 }
 
diff --git a/src/runtime/subtree.c b/src/runtime/subtree.c
index 83a1ad4e..625f757c 100644
--- a/src/runtime/subtree.c
+++ b/src/runtime/subtree.c
@@ -19,7 +19,7 @@ typedef struct {
 
 TSStateId TS_TREE_STATE_NONE = USHRT_MAX;
 
-static const uint32_t MAX_TREE_POOL_SIZE = 1024;
+static const uint32_t MAX_TREE_POOL_SIZE = 0;
 
 static const ExternalScannerState empty_state = {.length = 0, .short_data = {0}};
 
@@ -157,6 +157,7 @@ Subtree *ts_subtree_new_leaf(SubtreePool *pool, TSSymbol symbol, Length padding,
     .named = metadata.named,
     .node_count = 1,
     .has_changes = false,
+    .is_keyword = false,
     .first_leaf = {
       .symbol = symbol,
       .lex_mode = {0, 0},
diff --git a/src/runtime/subtree.h b/src/runtime/subtree.h
index 0fd3311a..66fc1f24 100644
--- a/src/runtime/subtree.h
+++ b/src/runtime/subtree.h
@@ -45,6 +45,7 @@ struct Subtree {
   bool has_changes : 1;
   bool has_external_tokens : 1;
   bool is_missing : 1;
+  bool is_keyword : 1;
   TSSymbol symbol;
   TSStateId parse_state;
   struct {