From d57043b66555b631236e54c7ac47aa1fdd20d627 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Wed, 4 Jan 2017 21:22:23 -0800
Subject: [PATCH] Add ability to store external token state per stack version

---
 spec/runtime/stack_spec.cc | 25 +++++++++++++++++++++++++
 src/runtime/stack.c        | 28 ++++++++++++++++++++++++----
 src/runtime/stack.h        |  6 +++++-
 3 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/spec/runtime/stack_spec.cc b/spec/runtime/stack_spec.cc
index 4d4b01fd..20180843 100644
--- a/spec/runtime/stack_spec.cc
+++ b/spec/runtime/stack_spec.cc
@@ -521,6 +521,31 @@ describe("Stack", [&]() {
       free_slice_array(&pop.slices);
     });
   });
+
+  describe("setting external token state", [&]() {
+    TSExternalTokenState external_token_state1, external_token_state2;
+
+    it("allows the state to be retrieved", [&]() {
+      AssertThat(ts_stack_external_token_state(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_copy_version(stack, 0);
+      AssertThat(ts_stack_external_token_state(stack, 0), Equals(&external_token_state1));
+    });
+
+    it("does not merge stack versions with different external token states", [&]() {
+      ts_stack_copy_version(stack, 0);
+      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);
+
+      AssertThat(ts_stack_merge(stack, 0, 1), IsFalse());
+    });
+  });
 });
 
 END_TEST
diff --git a/src/runtime/stack.c b/src/runtime/stack.c
index bdc5945c..198cce4d 100644
--- a/src/runtime/stack.c
+++ b/src/runtime/stack.c
@@ -50,6 +50,7 @@ typedef struct {
   StackNode *node;
   bool is_halted;
   unsigned push_count;
+  const TSExternalTokenState *external_token_state;
 } StackHead;
 
 struct Stack {
@@ -288,7 +289,12 @@ Stack *ts_stack_new() {
   self->base_node =
     stack_node_new(NULL, NULL, false, 1, length_zero(), &self->node_pool);
   stack_node_retain(self->base_node);
-  array_push(&self->heads, ((StackHead){ self->base_node, false, 0 }));
+  array_push(&self->heads, ((StackHead){
+    self->base_node,
+    false,
+    0,
+    NULL
+  }));
 
   return self;
 }
@@ -327,11 +333,19 @@ unsigned ts_stack_push_count(const Stack *self, StackVersion version) {
   return array_get(&self->heads, version)->push_count;
 }
 
-void ts_stack_decrease_push_count(const Stack *self, StackVersion version,
+void ts_stack_decrease_push_count(Stack *self, StackVersion version,
                                   unsigned decrement) {
   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;
+}
+
+void ts_stack_set_external_token_state(Stack *self, StackVersion version, const TSExternalTokenState *state) {
+  array_get(&self->heads, version)->external_token_state = state;
+}
+
 ErrorStatus ts_stack_error_status(const Stack *self, StackVersion version) {
   StackHead *head = array_get(&self->heads, version);
   return (ErrorStatus){
@@ -480,7 +494,8 @@ bool ts_stack_merge(Stack *self, StackVersion version, StackVersion new_version)
   if (new_node->state == node->state &&
       new_node->position.chars == node->position.chars &&
       new_node->error_count == node->error_count &&
-      new_node->error_cost == node->error_cost) {
+      new_node->error_cost == node->error_cost &&
+      new_head->external_token_state == head->external_token_state) {
     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)
@@ -505,7 +520,12 @@ void ts_stack_clear(Stack *self) {
   for (uint32_t i = 0; i < self->heads.size; i++)
     stack_node_release(self->heads.contents[i].node, &self->node_pool);
   array_clear(&self->heads);
-  array_push(&self->heads, ((StackHead){ self->base_node, false, 0 }));
+  array_push(&self->heads, ((StackHead){
+    self->base_node,
+    false,
+    0,
+    NULL
+  }));
 }
 
 bool ts_stack_print_dot_graph(Stack *self, const char **symbol_names, FILE *f) {
diff --git a/src/runtime/stack.h b/src/runtime/stack.h
index 64d9842b..2e88d72a 100644
--- a/src/runtime/stack.h
+++ b/src/runtime/stack.h
@@ -65,7 +65,11 @@ TSStateId ts_stack_top_state(const Stack *, StackVersion);
 
 unsigned ts_stack_push_count(const Stack *, StackVersion);
 
-void ts_stack_decrease_push_count(const Stack *, StackVersion, unsigned);
+void ts_stack_decrease_push_count(Stack *, StackVersion, unsigned);
+
+const TSExternalTokenState *ts_stack_external_token_state(const Stack *, StackVersion);
+
+void ts_stack_set_external_token_state(Stack *, StackVersion, const TSExternalTokenState *);
 
 /*
  *  Get the position at the top of the given version of the stack. If the stack