diff --git a/examples/grammars/arithmetic.cc b/examples/grammars/arithmetic.cc
index e414711c..915c7ed9 100644
--- a/examples/grammars/arithmetic.cc
+++ b/examples/grammars/arithmetic.cc
@@ -6,7 +6,7 @@ namespace tree_sitter_examples {
 using tree_sitter::Grammar;
 using namespace tree_sitter::rules;
 
-extern const Grammar arithmetic({
+extern const Grammar arithmetic = Grammar({
     { "expression", choice({
       sym("sum"),
       sym("difference"),
@@ -26,6 +26,8 @@ extern const Grammar arithmetic({
 
     { "number", pattern("\\d+") },
     { "variable", pattern("\\a[\\w_]*") },
-});
+
+    { "comment", pattern("#.*") },
+}).ubiquitous_tokens({ "comment" });
 
 }  // namespace tree_sitter_examples
diff --git a/examples/parsers/arithmetic.c b/examples/parsers/arithmetic.c
index 7b835dbd..e629bfe3 100644
--- a/examples/parsers/arithmetic.c
+++ b/examples/parsers/arithmetic.c
@@ -1,7 +1,7 @@
 #include "tree_sitter/parser.h"
 
 #define STATE_COUNT 32
-#define SYMBOL_COUNT 19
+#define SYMBOL_COUNT 20
 
 enum {
     ts_sym_expression = ts_builtin_sym_start,
@@ -13,6 +13,7 @@ enum {
     ts_sym_group,
     ts_sym_number,
     ts_sym_variable,
+    ts_sym_comment,
     ts_aux_sym_1,
     ts_aux_sym_2,
     ts_aux_sym_3,
@@ -35,6 +36,7 @@ SYMBOL_NAMES = {
     [ts_builtin_sym_end] = "end",
     [ts_sym_number] = "number",
     [ts_sym_variable] = "variable",
+    [ts_sym_comment] = "comment",
     [ts_aux_sym_1] = "'+'",
     [ts_aux_sym_2] = "'-'",
     [ts_aux_sym_3] = "'*'",
@@ -64,148 +66,164 @@ LEX_FN() {
                 (lookahead == '\r') ||
                 (lookahead == ' '))
                 ADVANCE(1);
-            if (lookahead == '(')
+            if (lookahead == '#')
                 ADVANCE(2);
-            if ('0' <= lookahead && lookahead <= '9')
+            if (lookahead == '(')
                 ADVANCE(3);
+            if ('0' <= lookahead && lookahead <= '9')
+                ADVANCE(4);
             if (('A' <= lookahead && lookahead <= 'Z') ||
                 ('a' <= lookahead && lookahead <= 'z'))
-                ADVANCE(4);
+                ADVANCE(5);
             LEX_ERROR();
         case 2:
-            ACCEPT_TOKEN(ts_aux_sym_6);
+            if (!(lookahead == '\n'))
+                ADVANCE(2);
+            ACCEPT_TOKEN(ts_sym_comment);
         case 3:
-            if ('0' <= lookahead && lookahead <= '9')
-                ADVANCE(3);
-            ACCEPT_TOKEN(ts_sym_number);
+            ACCEPT_TOKEN(ts_aux_sym_6);
         case 4:
+            if ('0' <= lookahead && lookahead <= '9')
+                ADVANCE(4);
+            ACCEPT_TOKEN(ts_sym_number);
+        case 5:
             if (('0' <= lookahead && lookahead <= '9') ||
                 ('A' <= lookahead && lookahead <= 'Z') ||
                 (lookahead == '_') ||
                 ('a' <= lookahead && lookahead <= 'z'))
-                ADVANCE(4);
+                ADVANCE(5);
             ACCEPT_TOKEN(ts_sym_variable);
-        case 5:
+        case 6:
             START_TOKEN();
             if (lookahead == 0)
+                ADVANCE(7);
+            if ((lookahead == '\t') ||
+                (lookahead == '\n') ||
+                (lookahead == '\r') ||
+                (lookahead == ' '))
                 ADVANCE(6);
-            if ((lookahead == '\t') ||
-                (lookahead == '\n') ||
-                (lookahead == '\r') ||
-                (lookahead == ' '))
-                ADVANCE(5);
+            if (lookahead == '#')
+                ADVANCE(2);
             if (lookahead == '*')
-                ADVANCE(7);
-            if (lookahead == '+')
                 ADVANCE(8);
-            if (lookahead == '-')
+            if (lookahead == '+')
                 ADVANCE(9);
-            if (lookahead == '/')
+            if (lookahead == '-')
                 ADVANCE(10);
-            if (lookahead == '^')
+            if (lookahead == '/')
                 ADVANCE(11);
-            LEX_ERROR();
-        case 6:
-            ACCEPT_TOKEN(ts_builtin_sym_end);
-        case 7:
-            ACCEPT_TOKEN(ts_aux_sym_3);
-        case 8:
-            ACCEPT_TOKEN(ts_aux_sym_1);
-        case 9:
-            ACCEPT_TOKEN(ts_aux_sym_2);
-        case 10:
-            ACCEPT_TOKEN(ts_aux_sym_4);
-        case 11:
-            ACCEPT_TOKEN(ts_aux_sym_5);
-        case 12:
-            START_TOKEN();
-            if ((lookahead == '\t') ||
-                (lookahead == '\n') ||
-                (lookahead == '\r') ||
-                (lookahead == ' '))
+            if (lookahead == '^')
                 ADVANCE(12);
-            if (lookahead == ')')
-                ADVANCE(13);
-            if (lookahead == '*')
-                ADVANCE(7);
-            if (lookahead == '+')
-                ADVANCE(8);
-            if (lookahead == '-')
-                ADVANCE(9);
-            if (lookahead == '/')
-                ADVANCE(10);
-            if (lookahead == '^')
-                ADVANCE(11);
             LEX_ERROR();
+        case 7:
+            ACCEPT_TOKEN(ts_builtin_sym_end);
+        case 8:
+            ACCEPT_TOKEN(ts_aux_sym_3);
+        case 9:
+            ACCEPT_TOKEN(ts_aux_sym_1);
+        case 10:
+            ACCEPT_TOKEN(ts_aux_sym_2);
+        case 11:
+            ACCEPT_TOKEN(ts_aux_sym_4);
+        case 12:
+            ACCEPT_TOKEN(ts_aux_sym_5);
         case 13:
-            ACCEPT_TOKEN(ts_aux_sym_7);
-        case 14:
             START_TOKEN();
             if ((lookahead == '\t') ||
                 (lookahead == '\n') ||
                 (lookahead == '\r') ||
                 (lookahead == ' '))
-                ADVANCE(14);
-            if (lookahead == ')')
                 ADVANCE(13);
+            if (lookahead == '#')
+                ADVANCE(2);
+            if (lookahead == ')')
+                ADVANCE(14);
+            if (lookahead == '*')
+                ADVANCE(8);
+            if (lookahead == '+')
+                ADVANCE(9);
+            if (lookahead == '-')
+                ADVANCE(10);
+            if (lookahead == '/')
+                ADVANCE(11);
+            if (lookahead == '^')
+                ADVANCE(12);
             LEX_ERROR();
+        case 14:
+            ACCEPT_TOKEN(ts_aux_sym_7);
         case 15:
             START_TOKEN();
-            if (lookahead == 0)
-                ADVANCE(6);
             if ((lookahead == '\t') ||
                 (lookahead == '\n') ||
                 (lookahead == '\r') ||
                 (lookahead == ' '))
                 ADVANCE(15);
-            if (lookahead == '(')
+            if (lookahead == '#')
                 ADVANCE(2);
             if (lookahead == ')')
-                ADVANCE(13);
-            if (lookahead == '*')
+                ADVANCE(14);
+            LEX_ERROR();
+        case 16:
+            START_TOKEN();
+            if (lookahead == 0)
                 ADVANCE(7);
-            if (lookahead == '+')
-                ADVANCE(8);
-            if (lookahead == '-')
-                ADVANCE(9);
-            if (lookahead == '/')
-                ADVANCE(10);
-            if ('0' <= lookahead && lookahead <= '9')
+            if ((lookahead == '\t') ||
+                (lookahead == '\n') ||
+                (lookahead == '\r') ||
+                (lookahead == ' '))
+                ADVANCE(16);
+            if (lookahead == '#')
+                ADVANCE(2);
+            if (lookahead == '(')
                 ADVANCE(3);
+            if (lookahead == ')')
+                ADVANCE(14);
+            if (lookahead == '*')
+                ADVANCE(8);
+            if (lookahead == '+')
+                ADVANCE(9);
+            if (lookahead == '-')
+                ADVANCE(10);
+            if (lookahead == '/')
+                ADVANCE(11);
+            if ('0' <= lookahead && lookahead <= '9')
+                ADVANCE(4);
             if (('A' <= lookahead && lookahead <= 'Z') ||
                 ('a' <= lookahead && lookahead <= 'z'))
-                ADVANCE(4);
+                ADVANCE(5);
             if (lookahead == '^')
-                ADVANCE(11);
+                ADVANCE(12);
             LEX_ERROR();
         case ts_lex_state_error:
             START_TOKEN();
             if (lookahead == 0)
-                ADVANCE(6);
+                ADVANCE(7);
             if ((lookahead == '\t') ||
                 (lookahead == '\n') ||
                 (lookahead == '\r') ||
                 (lookahead == ' '))
-                ADVANCE(15);
-            if (lookahead == '(')
+                ADVANCE(16);
+            if (lookahead == '#')
                 ADVANCE(2);
-            if (lookahead == ')')
-                ADVANCE(13);
-            if (lookahead == '*')
-                ADVANCE(7);
-            if (lookahead == '+')
-                ADVANCE(8);
-            if (lookahead == '-')
-                ADVANCE(9);
-            if (lookahead == '/')
-                ADVANCE(10);
-            if ('0' <= lookahead && lookahead <= '9')
+            if (lookahead == '(')
                 ADVANCE(3);
+            if (lookahead == ')')
+                ADVANCE(14);
+            if (lookahead == '*')
+                ADVANCE(8);
+            if (lookahead == '+')
+                ADVANCE(9);
+            if (lookahead == '-')
+                ADVANCE(10);
+            if (lookahead == '/')
+                ADVANCE(11);
+            if ('0' <= lookahead && lookahead <= '9')
+                ADVANCE(4);
             if (('A' <= lookahead && lookahead <= 'Z') ||
                 ('a' <= lookahead && lookahead <= 'z'))
-                ADVANCE(4);
+                ADVANCE(5);
             if (lookahead == '^')
-                ADVANCE(11);
+                ADVANCE(12);
             LEX_ERROR();
         default:
             LEX_ERROR();
@@ -214,37 +232,37 @@ LEX_FN() {
 
 LEX_STATES = {
     [0] = 1,
-    [1] = 5,
-    [2] = 5,
+    [1] = 6,
+    [2] = 6,
     [3] = 1,
-    [4] = 12,
-    [5] = 12,
-    [6] = 14,
+    [4] = 13,
+    [5] = 13,
+    [6] = 15,
     [7] = 1,
-    [8] = 12,
-    [9] = 14,
-    [10] = 12,
+    [8] = 13,
+    [9] = 15,
+    [10] = 13,
     [11] = 1,
     [12] = 1,
     [13] = 1,
     [14] = 1,
     [15] = 1,
-    [16] = 12,
-    [17] = 12,
-    [18] = 12,
-    [19] = 12,
-    [20] = 12,
-    [21] = 5,
+    [16] = 13,
+    [17] = 13,
+    [18] = 13,
+    [19] = 13,
+    [20] = 13,
+    [21] = 6,
     [22] = 1,
     [23] = 1,
     [24] = 1,
     [25] = 1,
     [26] = 1,
-    [27] = 5,
-    [28] = 5,
-    [29] = 5,
-    [30] = 5,
-    [31] = 5,
+    [27] = 6,
+    [28] = 6,
+    [29] = 6,
+    [30] = 6,
+    [31] = 6,
 };
 
 #pragma GCC diagnostic push
@@ -261,10 +279,12 @@ PARSE_TABLE = {
         [ts_sym_group] = SHIFT(2),
         [ts_sym_number] = SHIFT(2),
         [ts_sym_variable] = SHIFT(2),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_6] = SHIFT(3),
     },
     [1] = {
         [ts_builtin_sym_end] = ACCEPT_INPUT(),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = SHIFT(22),
         [ts_aux_sym_2] = SHIFT(23),
         [ts_aux_sym_3] = SHIFT(24),
@@ -273,6 +293,7 @@ PARSE_TABLE = {
     },
     [2] = {
         [ts_builtin_sym_end] = REDUCE(ts_sym_expression, 1),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = REDUCE(ts_sym_expression, 1),
         [ts_aux_sym_2] = REDUCE(ts_sym_expression, 1),
         [ts_aux_sym_3] = REDUCE(ts_sym_expression, 1),
@@ -290,9 +311,11 @@ PARSE_TABLE = {
         [ts_builtin_sym_error] = SHIFT(6),
         [ts_sym_number] = SHIFT(5),
         [ts_sym_variable] = SHIFT(5),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_6] = SHIFT(7),
     },
     [4] = {
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = SHIFT(11),
         [ts_aux_sym_2] = SHIFT(12),
         [ts_aux_sym_3] = SHIFT(13),
@@ -301,6 +324,7 @@ PARSE_TABLE = {
         [ts_aux_sym_7] = SHIFT(21),
     },
     [5] = {
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = REDUCE(ts_sym_expression, 1),
         [ts_aux_sym_2] = REDUCE(ts_sym_expression, 1),
         [ts_aux_sym_3] = REDUCE(ts_sym_expression, 1),
@@ -309,6 +333,7 @@ PARSE_TABLE = {
         [ts_aux_sym_7] = REDUCE(ts_sym_expression, 1),
     },
     [6] = {
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_7] = SHIFT(21),
     },
     [7] = {
@@ -322,9 +347,11 @@ PARSE_TABLE = {
         [ts_builtin_sym_error] = SHIFT(9),
         [ts_sym_number] = SHIFT(5),
         [ts_sym_variable] = SHIFT(5),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_6] = SHIFT(7),
     },
     [8] = {
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = SHIFT(11),
         [ts_aux_sym_2] = SHIFT(12),
         [ts_aux_sym_3] = SHIFT(13),
@@ -333,9 +360,11 @@ PARSE_TABLE = {
         [ts_aux_sym_7] = SHIFT(10),
     },
     [9] = {
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_7] = SHIFT(10),
     },
     [10] = {
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = REDUCE(ts_sym_group, 3),
         [ts_aux_sym_2] = REDUCE(ts_sym_group, 3),
         [ts_aux_sym_3] = REDUCE(ts_sym_group, 3),
@@ -353,6 +382,7 @@ PARSE_TABLE = {
         [ts_sym_group] = SHIFT(5),
         [ts_sym_number] = SHIFT(5),
         [ts_sym_variable] = SHIFT(5),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_6] = SHIFT(7),
     },
     [12] = {
@@ -365,6 +395,7 @@ PARSE_TABLE = {
         [ts_sym_group] = SHIFT(5),
         [ts_sym_number] = SHIFT(5),
         [ts_sym_variable] = SHIFT(5),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_6] = SHIFT(7),
     },
     [13] = {
@@ -377,6 +408,7 @@ PARSE_TABLE = {
         [ts_sym_group] = SHIFT(5),
         [ts_sym_number] = SHIFT(5),
         [ts_sym_variable] = SHIFT(5),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_6] = SHIFT(7),
     },
     [14] = {
@@ -389,6 +421,7 @@ PARSE_TABLE = {
         [ts_sym_group] = SHIFT(5),
         [ts_sym_number] = SHIFT(5),
         [ts_sym_variable] = SHIFT(5),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_6] = SHIFT(7),
     },
     [15] = {
@@ -401,9 +434,11 @@ PARSE_TABLE = {
         [ts_sym_group] = SHIFT(5),
         [ts_sym_number] = SHIFT(5),
         [ts_sym_variable] = SHIFT(5),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_6] = SHIFT(7),
     },
     [16] = {
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = REDUCE(ts_sym_exponent, 3),
         [ts_aux_sym_2] = REDUCE(ts_sym_exponent, 3),
         [ts_aux_sym_3] = REDUCE(ts_sym_exponent, 3),
@@ -412,6 +447,7 @@ PARSE_TABLE = {
         [ts_aux_sym_7] = REDUCE(ts_sym_exponent, 3),
     },
     [17] = {
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = REDUCE(ts_sym_quotient, 3),
         [ts_aux_sym_2] = REDUCE(ts_sym_quotient, 3),
         [ts_aux_sym_3] = SHIFT(13),
@@ -420,6 +456,7 @@ PARSE_TABLE = {
         [ts_aux_sym_7] = REDUCE(ts_sym_quotient, 3),
     },
     [18] = {
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = REDUCE(ts_sym_product, 3),
         [ts_aux_sym_2] = REDUCE(ts_sym_product, 3),
         [ts_aux_sym_3] = SHIFT(13),
@@ -428,6 +465,7 @@ PARSE_TABLE = {
         [ts_aux_sym_7] = REDUCE(ts_sym_product, 3),
     },
     [19] = {
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = SHIFT(11),
         [ts_aux_sym_2] = SHIFT(12),
         [ts_aux_sym_3] = SHIFT(13),
@@ -436,6 +474,7 @@ PARSE_TABLE = {
         [ts_aux_sym_7] = REDUCE(ts_sym_difference, 3),
     },
     [20] = {
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = SHIFT(11),
         [ts_aux_sym_2] = SHIFT(12),
         [ts_aux_sym_3] = SHIFT(13),
@@ -445,6 +484,7 @@ PARSE_TABLE = {
     },
     [21] = {
         [ts_builtin_sym_end] = REDUCE(ts_sym_group, 3),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = REDUCE(ts_sym_group, 3),
         [ts_aux_sym_2] = REDUCE(ts_sym_group, 3),
         [ts_aux_sym_3] = REDUCE(ts_sym_group, 3),
@@ -461,6 +501,7 @@ PARSE_TABLE = {
         [ts_sym_group] = SHIFT(2),
         [ts_sym_number] = SHIFT(2),
         [ts_sym_variable] = SHIFT(2),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_6] = SHIFT(3),
     },
     [23] = {
@@ -473,6 +514,7 @@ PARSE_TABLE = {
         [ts_sym_group] = SHIFT(2),
         [ts_sym_number] = SHIFT(2),
         [ts_sym_variable] = SHIFT(2),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_6] = SHIFT(3),
     },
     [24] = {
@@ -485,6 +527,7 @@ PARSE_TABLE = {
         [ts_sym_group] = SHIFT(2),
         [ts_sym_number] = SHIFT(2),
         [ts_sym_variable] = SHIFT(2),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_6] = SHIFT(3),
     },
     [25] = {
@@ -497,6 +540,7 @@ PARSE_TABLE = {
         [ts_sym_group] = SHIFT(2),
         [ts_sym_number] = SHIFT(2),
         [ts_sym_variable] = SHIFT(2),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_6] = SHIFT(3),
     },
     [26] = {
@@ -509,10 +553,12 @@ PARSE_TABLE = {
         [ts_sym_group] = SHIFT(2),
         [ts_sym_number] = SHIFT(2),
         [ts_sym_variable] = SHIFT(2),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_6] = SHIFT(3),
     },
     [27] = {
         [ts_builtin_sym_end] = REDUCE(ts_sym_exponent, 3),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = REDUCE(ts_sym_exponent, 3),
         [ts_aux_sym_2] = REDUCE(ts_sym_exponent, 3),
         [ts_aux_sym_3] = REDUCE(ts_sym_exponent, 3),
@@ -521,6 +567,7 @@ PARSE_TABLE = {
     },
     [28] = {
         [ts_builtin_sym_end] = REDUCE(ts_sym_quotient, 3),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = REDUCE(ts_sym_quotient, 3),
         [ts_aux_sym_2] = REDUCE(ts_sym_quotient, 3),
         [ts_aux_sym_3] = SHIFT(24),
@@ -529,6 +576,7 @@ PARSE_TABLE = {
     },
     [29] = {
         [ts_builtin_sym_end] = REDUCE(ts_sym_product, 3),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = REDUCE(ts_sym_product, 3),
         [ts_aux_sym_2] = REDUCE(ts_sym_product, 3),
         [ts_aux_sym_3] = SHIFT(24),
@@ -537,6 +585,7 @@ PARSE_TABLE = {
     },
     [30] = {
         [ts_builtin_sym_end] = REDUCE(ts_sym_difference, 3),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = SHIFT(22),
         [ts_aux_sym_2] = SHIFT(23),
         [ts_aux_sym_3] = SHIFT(24),
@@ -545,6 +594,7 @@ PARSE_TABLE = {
     },
     [31] = {
         [ts_builtin_sym_end] = REDUCE(ts_sym_sum, 3),
+        [ts_sym_comment] = SHIFT_EXTRA(),
         [ts_aux_sym_1] = SHIFT(22),
         [ts_aux_sym_2] = SHIFT(23),
         [ts_aux_sym_3] = SHIFT(24),
diff --git a/spec/runtime/document_spec.cc b/spec/runtime/document_spec.cc
index 033f3b54..ec970c49 100644
--- a/spec/runtime/document_spec.cc
+++ b/spec/runtime/document_spec.cc
@@ -3,6 +3,7 @@
 
 extern "C" const TSLanguage * ts_language_json();
 extern "C" const TSLanguage * ts_language_javascript();
+extern "C" const TSLanguage * ts_language_arithmetic();
 
 START_TEST
 
@@ -64,19 +65,19 @@ describe("Document", [&]() {
     SpyReader *reader;
 
     before_each([&]() {
-      reader = new SpyReader("{ \"key\": [1, 2] }", 3);
-      ts_document_set_language(doc, ts_language_json());
-      ts_document_set_input(doc, reader->input);
-      reader->clear();
-
-      AssertThat(ts_node_string(ts_document_root_node(doc)), Equals(
-          "(DOCUMENT (object (string) (array (number) (number))))"));
+      ts_document_set_language(doc, ts_language_arithmetic());
     });
 
     after_each([&]() {
       delete reader;
     });
 
+    auto set_text = [&](const char *text) {
+      reader = new SpyReader(text, 3);
+      ts_document_set_input(doc, reader->input);
+      reader->clear();
+    };
+
     auto insert_text = [&](size_t position, string text) {
       reader->content.insert(position, text);
       ts_document_edit(doc, { position, 0, text.length() });
@@ -89,87 +90,116 @@ describe("Document", [&]() {
 
     describe("inserting new tokens near the end of the input", [&]() {
       before_each([&]() {
-        insert_text(
-            strlen("{ \"key\": [1, 2]"),
-            ", \"key2\": 4");
+        set_text("x ^ (100 + abc)");
+        AssertThat(string(ts_node_string(ts_document_root_node(doc))), Equals(
+            "(DOCUMENT (exponent "
+              "(variable) "
+              "(group (sum (number) (variable)))))"));
+
+        insert_text(strlen("x ^ (100 + abc"), " * 5");
       });
 
       it("updates the parse tree", [&]() {
         AssertThat(string(ts_node_string(ts_document_root_node(doc))), Equals(
-            "(DOCUMENT (object (string) (array (number) (number)) (string) (number)))"));
+            "(DOCUMENT (exponent "
+              "(variable) " 
+              "(group (sum (number) (product (variable) (number))))))"));
       });
 
       it("re-reads only the changed portion of the input", [&]() {
-        AssertThat(reader->strings_read, Equals(vector<string>({ "], \"key2\": 4 }" })));
+        AssertThat(reader->strings_read, Equals(vector<string>({ " abc * 5)" })));
       });
     });
 
     describe("inserting text into the middle of an existing token", [&]() {
-      it("updates the parse three", [&]() {
-        insert_text(strlen("{ \"key"), "123");
-
+      before_each([&]() {
+        set_text("abc * 123");
         AssertThat(string(ts_node_string(ts_document_root_node(doc))), Equals(
-            "(DOCUMENT (object (string) (array (number) (number))))"));
-        TSNode *node = ts_node_find_for_pos(ts_document_root_node(doc), 3);
-        AssertThat(ts_node_name(node), Equals("string"));
-        AssertThat(ts_node_size(node), Equals(strlen("\"key123\"")));
+            "(DOCUMENT (product (variable) (number)))"));
+
+        insert_text(strlen("ab"), "XYZ");
+      });
+
+      it("updates the parse three", [&]() {
+        AssertThat(string(ts_node_string(ts_document_root_node(doc))), Equals(
+            "(DOCUMENT (product (variable) (number)))"));
+
+        TSNode *node = ts_node_find_for_pos(ts_document_root_node(doc), 1);
+        AssertThat(ts_node_name(node), Equals("variable"));
+        AssertThat(ts_node_size(node), Equals(strlen("abXYZc")));
         ts_node_release(node);
       });
     });
 
     describe("appending text to the end of an existing token", [&]() {
-      it("updates the parse three", [&]() {
-        insert_text(strlen("{ \"key\": [1"), "001");
-
+      before_each([&]() {
+        set_text("abc * 123");
         AssertThat(string(ts_node_string(ts_document_root_node(doc))), Equals(
-            "(DOCUMENT (object (string) (array (number) (number))))"));
+            "(DOCUMENT (product (variable) (number)))"));
 
-        TSNode *node = ts_node_find_for_pos(ts_document_root_node(doc), strlen("{ \"key\": ["));
-        AssertThat(ts_node_name(node), Equals("number"));
-        AssertThat(ts_node_size(node), Equals(strlen("1001")));
+        insert_text(strlen("abc"), "XYZ");
+      });
+
+      it("updates the parse three", [&]() {
+        AssertThat(string(ts_node_string(ts_document_root_node(doc))), Equals(
+            "(DOCUMENT (product (variable) (number)))"));
+
+        TSNode *node = ts_node_find_for_pos(ts_document_root_node(doc), 1);
+        AssertThat(ts_node_name(node), Equals("variable"));
+        AssertThat(ts_node_size(node), Equals(strlen("abcXYZ")));
         ts_node_release(node);
       });
     });
 
     describe("editing text inside a node containing a ubiquitous token", [&]() {
       before_each([&]() {
-        ts_document_set_language(doc, ts_language_javascript());
+        set_text("123 *\n"
+            "# a-comment\n"
+            "abc");
 
-        delete reader;
-        reader = new SpyReader("{\nx;\n}", 3);
-
-        ts_document_set_input(doc, reader->input);
         AssertThat(string(ts_node_string(ts_document_root_node(doc))), Equals(
-            "(DOCUMENT (program (statement_block (expression_statement (identifier)))))"));
+            "(DOCUMENT (product (number) (comment) (variable)))"));
 
+        insert_text(
+            strlen("123 *\n"
+                "# a-comment\n"
+                "abc"),
+            "XYZ");
       });
 
       it("updates the parse tree", [&]() {
-        insert_text(strlen("{\nx"), "y");
         AssertThat(string(ts_node_string(ts_document_root_node(doc))), Equals(
-            "(DOCUMENT (program (statement_block (expression_statement (identifier)))))"));
+            "(DOCUMENT (product (number) (comment) (variable)))"));
       });
     });
 
-    describe("deleting an important part of a token", [&]() {
-      it("updates the parse tree, creating an error", [&]() {
-        delete_text(strlen("{ \"key"), 1);
-
+    describe("deleting an important token", [&]() {
+      before_each([&]() {
+        set_text("123 * 456");
         AssertThat(string(ts_node_string(ts_document_root_node(doc))), Equals(
-            "(DOCUMENT (end))"));
+            "(DOCUMENT (product (number) (number)))"));
+
+        delete_text(strlen("123 "), 2);
+      });
+
+      it("updates the parse tree, creating an error", [&]() {
+        AssertThat(string(ts_node_string(ts_document_root_node(doc))), Equals(
+            "(DOCUMENT (number) (ERROR '4'))"));
       });
     });
 
     describe("inserting tokens near the beginning of the input", [&]() {
       before_each([&]() {
-        insert_text(
-            strlen("{ "),
-            "\"key2\": 4, ");
+        set_text("123 * 456");
+        AssertThat(string(ts_node_string(ts_document_root_node(doc))), Equals(
+            "(DOCUMENT (product (number) (number)))"));
+
+        insert_text(strlen("123"), " + 5 ");
       });
 
       it("updates the parse tree", [&]() {
         AssertThat(string(ts_node_string(ts_document_root_node(doc))), Equals(
-            "(DOCUMENT (object (string) (number) (string) (array (number) (number))))"));
+            "(DOCUMENT (sum (number) (product (number) (number))))"));
       });
 
       it_skip("re-reads only the changed portion of the input", [&]() {
diff --git a/src/runtime/parser.c b/src/runtime/parser.c
index 8a55cf63..f1d746fd 100644
--- a/src/runtime/parser.c
+++ b/src/runtime/parser.c
@@ -18,18 +18,19 @@ static TSParseAction action_for(const TSLanguage *lang, TSStateId state,
 }
 
 static size_t breakdown_stack(TSParser *parser, TSInputEdit *edit) {
-  if (!edit)
+  if (!edit) {
+    ts_stack_shrink(&parser->stack, 0);
     return 0;
+  }
 
   TSStack *stack = &parser->stack;
-  size_t position = 0;
+  size_t position = ts_stack_right_position(stack);
 
   for (;;) {
     TSTree *node = ts_stack_top_node(stack);
     if (!node)
       break;
 
-    position = ts_stack_right_position(stack);
     size_t child_count;
     TSTree **children = ts_tree_children(node, &child_count);
     if (position < edit->position && !children)
@@ -37,23 +38,20 @@ static size_t breakdown_stack(TSParser *parser, TSInputEdit *edit) {
 
     stack->size--;
     position -= ts_tree_total_size(node);
-    DEBUG_PARSE("BREAKDOWN %s %u",
-        parser->language->symbol_names[node->symbol],
-        ts_stack_top_state(stack));
+    DEBUG_PARSE("BREAKDOWN %s %u", parser->language->symbol_names[node->symbol],
+                ts_stack_top_state(stack));
 
     for (size_t i = 0; i < child_count && position < edit->position; i++) {
       TSTree *child = children[i];
       TSStateId state = ts_stack_top_state(stack);
       TSParseAction action = action_for(parser->language, state, child->symbol);
-      TSStateId next_state = action.type == TSParseActionTypeShift ?
-          action.data.to_state :
-          state;
+      TSStateId next_state =
+          action.type == TSParseActionTypeShift ? action.data.to_state : state;
       ts_stack_push(stack, next_state, child);
       ts_tree_retain(child);
       position += ts_tree_total_size(child);
       DEBUG_PARSE("PUT_BACK %s %u",
-          parser->language->symbol_names[child->symbol],
-          next_state);
+                  parser->language->symbol_names[child->symbol], next_state);
     }
 
     ts_tree_release(node);
@@ -214,9 +212,6 @@ void ts_parser_destroy(TSParser *parser) {
 
 const TSTree *ts_parser_parse(TSParser *parser, TSInput input,
                               TSInputEdit *edit) {
-  if (!edit)
-    ts_stack_shrink(&parser->stack, 0);
-
   parser->lookahead = NULL;
   parser->next_lookahead = NULL;
   parser->lexer = ts_lexer_make();