From 5fd818babe8d41c0fbde45902023b49a5e557ebd Mon Sep 17 00:00:00 2001
From: Will Lillis <will.lillis24@gmail.com>
Date: Thu, 4 Sep 2025 03:14:28 -0400
Subject: [PATCH] fix(generate): use correct state id when adding terminal
 states to non terminal extras

---
 crates/generate/src/build_tables/build_parse_table.rs    | 7 ++++---
 .../test_grammars/extra_non_terminals/corpus.txt         | 9 +++++----
 .../test_grammars/extra_non_terminals/grammar.js         | 9 +++++++--
 3 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/crates/generate/src/build_tables/build_parse_table.rs b/crates/generate/src/build_tables/build_parse_table.rs
index 9c340db7..77f3f4e2 100644
--- a/crates/generate/src/build_tables/build_parse_table.rs
+++ b/crates/generate/src/build_tables/build_parse_table.rs
@@ -326,9 +326,10 @@ impl<'a> ParseTableBuilder<'a> {
                 ))?;
             }
 
-            self.non_terminal_extra_states
-                .push((terminal, self.parse_table.states.len()));
-            self.add_parse_state(&Vec::new(), &Vec::new(), item_set);
+            // Add the parse state, and *then* push the terminal and the state id into the
+            // list of nonterminal extra states
+            let state_id = self.add_parse_state(&Vec::new(), &Vec::new(), item_set);
+            self.non_terminal_extra_states.push((terminal, state_id));
         }
 
         while let Some(entry) = self.parse_state_queue.pop_front() {
diff --git a/test/fixtures/test_grammars/extra_non_terminals/corpus.txt b/test/fixtures/test_grammars/extra_non_terminals/corpus.txt
index 52b7d864..b58fa68b 100644
--- a/test/fixtures/test_grammars/extra_non_terminals/corpus.txt
+++ b/test/fixtures/test_grammars/extra_non_terminals/corpus.txt
@@ -12,11 +12,12 @@ a b c d
 Extras
 ==============
 
-a (one) b (two) (three) c d
+a (one) b (two) (three) c d // e
 
 ---
 
 (module
-  (comment)
-  (comment)
-  (comment))
+  (comment (paren_comment))
+  (comment (paren_comment))
+  (comment (paren_comment))
+  (comment (line_comment)))
diff --git a/test/fixtures/test_grammars/extra_non_terminals/grammar.js b/test/fixtures/test_grammars/extra_non_terminals/grammar.js
index d13cd68a..e66bc9ac 100644
--- a/test/fixtures/test_grammars/extra_non_terminals/grammar.js
+++ b/test/fixtures/test_grammars/extra_non_terminals/grammar.js
@@ -9,7 +9,12 @@ module.exports = grammar({
   ],
 
   rules: {
-    module: $ => seq('a', 'b', 'c', 'd'),
-    comment: $ => seq('(', repeat(/[a-z]+/), ')'),
+    module: _ => seq('a', 'b', 'c', 'd'),
+
+    comment: $ => choice($.paren_comment, $.line_comment),
+
+    paren_comment: _ => token(seq('(', repeat(/[a-z]+/), ')')),
+
+    line_comment: _ => token(seq('//', /.*/)),
   }
 })