Generate C code for out-of-context states

src/compiler/build_tables/build_parse_table.cc

@@ -107,25 +107,23 @@ class ParseTableBuilder {
   }
 
   void add_out_of_context_parse_states() {
-    map<Symbol, set<Symbol>> symbols_by_first_symbol =
-      build_tables::symbols_by_first_symbol(grammar);
-    for (size_t token_index = 0; token_index < lexical_grammar.variables.size();
-         token_index++) {
+    map<Symbol, set<Symbol>> symbols_by_token = symbols_by_first_symbol(grammar);
+    for (size_t i = 0; i < lexical_grammar.variables.size(); i++) {
+      Symbol token(i, true);
       ParseItemSet item_set;
-      const set<Symbol> &symbols =
-        symbols_by_first_symbol[Symbol(token_index, true)];
+      const set<Symbol> &symbols = symbols_by_token[token];
 
       for (const auto &parse_state_entry : parse_state_ids) {
         for (const auto &pair : parse_state_entry.first.entries) {
           const ParseItem &item = pair.first;
           const LookaheadSet &lookahead_set = pair.second;
-
           if (symbols.count(item.next_symbol()))
             item_set.entries[item].insert_all(lookahead_set);
         }
       }
 
-      add_parse_state(item_set);
+      ParseStateId state = add_parse_state(item_set);
+      parse_table.out_of_context_state_indices[token] = state;
     }
   }
 
@@ -235,7 +233,14 @@ class ParseTableBuilder {
   }
 
   void remove_duplicate_parse_states() {
-    remove_duplicate_states<ParseState, ParseAction>(&parse_table.states);
+    auto replacements = remove_duplicate_states<ParseState, ParseAction>(
+      &parse_table.states);
+
+    for (auto &pair : parse_table.out_of_context_state_indices) {
+      auto replacement = replacements.find(pair.second);
+      if (replacement != replacements.end())
+        pair.second = replacement->second;
+    }
   }
 
   ParseAction *add_action(ParseStateId state_id, Symbol lookahead,

src/compiler/generate_code/c_code.cc

@@ -98,6 +98,7 @@ class CCodeGenerator {
     add_symbol_node_types_list();
     add_lex_function();
     add_lex_states_list();
+    add_out_of_context_parse_states_list();
     add_parse_table();
     add_parser_export();
 
@@ -216,6 +217,21 @@ class CCodeGenerator {
     line();
   }
 
+  void add_out_of_context_parse_states_list() {
+    line("static TSStateId ts_out_of_context_states[SYMBOL_COUNT] = {");
+    indent([&]() {
+      for (const auto &entry : parse_table.symbols) {
+        const rules::Symbol &symbol = entry.first;
+        ParseStateId state = parse_table.out_of_context_state_indices.find(symbol)->second;
+        if (symbol.is_token && !symbol.is_built_in()) {
+          line("[" + symbol_id(symbol) + "] = " + to_string(state) + ",");
+        }
+      }
+    });
+    line("};");
+    line();
+  }
+
   void add_parse_table() {
     add_parse_actions({ ParseAction::Error() });
 

src/compiler/parse_table.h

@@ -102,6 +102,7 @@ class ParseTable {
 
   std::vector<ParseState> states;
   std::map<rules::Symbol, ParseTableSymbolMetadata> symbols;
+  std::map<rules::Symbol, size_t> out_of_context_state_indices;
 };
 
 }  // namespace tree_sitter