From 45f7cee0c890195876950f37488f5e23cc86b1df Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Sat, 18 Jun 2016 20:35:33 -0700
Subject: [PATCH] Handle extra tokens properly during error recovery

---
 include/tree_sitter/parser.h                   |  3 +++
 src/compiler/build_tables/build_parse_table.cc |  5 +++++
 src/compiler/generate_code/c_code.cc           | 10 +++++++---
 3 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/include/tree_sitter/parser.h b/include/tree_sitter/parser.h
index 5b4dcb09..6178ca21 100644
--- a/include/tree_sitter/parser.h
+++ b/include/tree_sitter/parser.h
@@ -175,6 +175,9 @@ enum {
     .type = TSParseActionTypeRecover, .data = {.to_state = to_state_value } \
   }
 
+#define RECOVER_EXTRA() \
+  { .type = TSParseActionTypeShift, .extra = true, }
+
 #define SHIFT_EXTRA()                                 \
   {                                                   \
     { .type = TSParseActionTypeShift, .extra = true } \
diff --git a/src/compiler/build_tables/build_parse_table.cc b/src/compiler/build_tables/build_parse_table.cc
index f59095df..98072941 100644
--- a/src/compiler/build_tables/build_parse_table.cc
+++ b/src/compiler/build_tables/build_parse_table.cc
@@ -109,6 +109,11 @@ class ParseTableBuilder {
       add_out_of_context_parse_state(symbol);
     }
 
+    for (const Symbol &symbol : grammar.extra_tokens) {
+      parse_table.error_state.actions[symbol].push_back(
+        ParseAction::ShiftExtra());
+    }
+
     for (size_t i = 0; i < grammar.variables.size(); i++) {
       Symbol symbol(i, false);
       add_out_of_context_parse_state(symbol);
diff --git a/src/compiler/generate_code/c_code.cc b/src/compiler/generate_code/c_code.cc
index 2a0a528f..63c8c1ad 100644
--- a/src/compiler/generate_code/c_code.cc
+++ b/src/compiler/generate_code/c_code.cc
@@ -224,9 +224,13 @@ class CCodeGenerator {
       for (const auto &entry : parse_table.error_state.actions) {
         const rules::Symbol &symbol = entry.first;
         if (!entry.second.empty()) {
-          ParseStateId state = entry.second[0].state_index;
-          line("[" + symbol_id(symbol) + "] = RECOVER(" + to_string(state) +
-               "),");
+          line("[" + symbol_id(symbol) + "] = ");
+          ParseAction action = entry.second[0];
+          if (action.extra) {
+            add("RECOVER_EXTRA(),");
+          } else {
+            add("RECOVER(" + to_string(action.state_index) + "),");
+          }
         }
       }
     });