From c14a776a3d84419ffd52cb0d1e65e7bfc6c22e64 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Sun, 19 Feb 2017 13:53:28 -0800
Subject: [PATCH] Avoid including trailing extra tokens within error nodes
 unnecessarily

---
 spec/runtime/parser_spec.cc | 19 +++++++++++++++++++
 src/runtime/array.h         |  3 +++
 src/runtime/parser.c        |  4 ++++
 src/runtime/tree.c          | 15 +++++++++++++++
 src/runtime/tree.h          |  1 +
 5 files changed, 42 insertions(+)

diff --git a/spec/runtime/parser_spec.cc b/spec/runtime/parser_spec.cc
index 88633f1f..0b4c0a3a 100644
--- a/spec/runtime/parser_spec.cc
+++ b/spec/runtime/parser_spec.cc
@@ -4,6 +4,7 @@
 #include "helpers/spy_input.h"
 #include "helpers/load_language.h"
 #include "helpers/record_alloc.h"
+#include "helpers/point_helpers.h"
 #include "helpers/stderr_logger.h"
 #include "helpers/dedent.h"
 
@@ -168,6 +169,24 @@ describe("Parser", [&]() {
           "(ERROR (program (expression_statement (identifier))) (UNEXPECTED EOF))");
       });
     });
+
+    describe("when there are extra tokens at the end of the viable prefix", [&]() {
+      it("does not include them in the error node", [&]() {
+        ts_document_set_language(document, get_test_language("javascript"));
+        set_text(
+          "var x;\n"
+          "\n"
+          "if\n"
+          "\n"
+          "var y;"
+        );
+
+        TSNode error = ts_node_named_child(root, 1);
+        AssertThat(ts_node_type(error, document), Equals("ERROR"));
+        AssertThat(ts_node_start_point(error), Equals<TSPoint>({2, 0}));
+        AssertThat(ts_node_end_point(error), Equals<TSPoint>({2, 2}));
+      });
+    });
   });
 
   describe("handling extra tokens", [&]() {
diff --git a/src/runtime/array.h b/src/runtime/array.h
index 88581b71..a8899c88 100644
--- a/src/runtime/array.h
+++ b/src/runtime/array.h
@@ -46,6 +46,9 @@ extern "C" {
   (array_grow((self), (self)->size + 1), \
    (self)->contents[(self)->size++] = (element))
 
+#define array_push_all(self, other)                                       \
+  array_splice((self), (self)->size, 0, (other)->size, (other)->contents)
+
 #define array_splice(self, index, old_count, new_count, new_elements)          \
   array__splice((VoidArray *)(self), array__elem_size(self), index, old_count, \
                 new_count, (new_elements))
diff --git a/src/runtime/parser.c b/src/runtime/parser.c
index d2ea8747..523c1f42 100644
--- a/src/runtime/parser.c
+++ b/src/runtime/parser.c
@@ -755,8 +755,12 @@ static bool parser__repair_error(Parser *self, StackSlice slice,
   }
 
   TreeArray skipped_children = ts_tree_array_remove_last_n(&children, skip_count);
+  TreeArray trailing_extras = ts_tree_array_remove_trailing_extras(&skipped_children);
   Tree *error = ts_tree_make_error_node(&skipped_children);
   array_push(&children, error);
+  array_push_all(&children, &trailing_extras);
+  trailing_extras.size = 0;
+  array_delete(&trailing_extras);
 
   for (uint32_t i = 0; i < slice.trees.size; i++)
     array_push(&children, slice.trees.contents[i]);
diff --git a/src/runtime/tree.c b/src/runtime/tree.c
index 0b10604e..e437a0dc 100644
--- a/src/runtime/tree.c
+++ b/src/runtime/tree.c
@@ -84,6 +84,21 @@ TreeArray ts_tree_array_remove_last_n(TreeArray *self, uint32_t remove_count) {
   return result;
 }
 
+TreeArray ts_tree_array_remove_trailing_extras(TreeArray *self) {
+  TreeArray result = array_new();
+
+  uint32_t i = self->size - 1;
+  for (; i + 1 > 0; i--) {
+    Tree *child = self->contents[i];
+    if (!child->extra) break;
+    array_push(&result, child);
+  }
+
+  self->size = i + 1;
+  array_reverse(&result);
+  return result;
+}
+
 Tree *ts_tree_make_error(Length size, Length padding, char lookahead_char) {
   Tree *result = ts_tree_make_leaf(ts_builtin_sym_error, padding, size,
                                      (TSSymbolMetadata){
diff --git a/src/runtime/tree.h b/src/runtime/tree.h
index 69ba8c38..58452fc2 100644
--- a/src/runtime/tree.h
+++ b/src/runtime/tree.h
@@ -69,6 +69,7 @@ bool ts_tree_array_copy(TreeArray, TreeArray *);
 void ts_tree_array_delete(TreeArray *);
 uint32_t ts_tree_array_essential_count(const TreeArray *);
 TreeArray ts_tree_array_remove_last_n(TreeArray *, uint32_t);
+TreeArray ts_tree_array_remove_trailing_extras(TreeArray *);
 
 Tree *ts_tree_make_leaf(TSSymbol, Length, Length, TSSymbolMetadata);
 Tree *ts_tree_make_node(TSSymbol, uint32_t, Tree **, TSSymbolMetadata);