From 5ef9c4d6aa8e50ed36826baab0450c1db76bb241 Mon Sep 17 00:00:00 2001
From: Phil Turnbull <philipturnbull@github.com>
Date: Thu, 29 Jun 2017 06:35:26 -0700
Subject: [PATCH] Increase size of ref_counts and add assertions

This bumps the size of the reference counts from 16- to 32-bit counters to make
it less likely to overflow. Also assert in the retain function that the
reference count didn't overflow.

32-bits seems big enough for non-pathological examples but a more fool-proof
fix may be to bump it to 64-bits.
---
 src/runtime/stack.c | 5 +++--
 src/runtime/tree.c  | 1 +
 src/runtime/tree.h  | 2 +-
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/runtime/stack.c b/src/runtime/stack.c
index 55a124d6..5926ed78 100644
--- a/src/runtime/stack.c
+++ b/src/runtime/stack.c
@@ -26,7 +26,7 @@ struct StackNode {
   Length position;
   StackLink links[MAX_LINK_COUNT];
   short unsigned int link_count;
-  short unsigned int ref_count;
+  uint32_t ref_count;
   unsigned error_cost;
   unsigned error_count;
 };
@@ -65,8 +65,9 @@ struct Stack {
 static void stack_node_retain(StackNode *self) {
   if (!self)
     return;
-  assert(self->ref_count != 0);
+  assert(self->ref_count > 0);
   self->ref_count++;
+  assert(self->ref_count != 0);
 }
 
 static void stack_node_release(StackNode *self, StackNodeArray *pool) {
diff --git a/src/runtime/tree.c b/src/runtime/tree.c
index c11843fb..92003ff1 100644
--- a/src/runtime/tree.c
+++ b/src/runtime/tree.c
@@ -232,6 +232,7 @@ Tree *ts_tree_make_error_node(TreeArray *children) {
 void ts_tree_retain(Tree *self) {
   assert(self->ref_count > 0);
   self->ref_count++;
+  assert(self->ref_count != 0);
 }
 
 void ts_tree_release(Tree *self) {
diff --git a/src/runtime/tree.h b/src/runtime/tree.h
index 8d1ce186..f5f23a9f 100644
--- a/src/runtime/tree.h
+++ b/src/runtime/tree.h
@@ -45,7 +45,7 @@ typedef struct Tree {
     TSLexMode lex_mode;
   } first_leaf;
 
-  unsigned short ref_count;
+  uint32_t ref_count;
   bool visible : 1;
   bool named : 1;
   bool extra : 1;