captures are not local to nested parser

fixes #3
microsoft · Jul 30, 2024 · e629e6c · e629e6c
1 parent 1ba49cb
commit e629e6c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/parser/src/earley/parser.rs b/parser/src/earley/parser.rs
@@ -1613,8 +1613,9 @@ impl Parser {
         &self.state.stats
     }
 
-    pub(crate) fn set_stats(&mut self, stats: ParserStats) {
-        self.state.stats = stats;
+    pub(crate) fn take_global_state_from(&mut self, other: &mut Parser) {
+        self.state.stats = other.state.stats.clone();
+        self.state.captures = std::mem::take(&mut other.state.captures);
     }
 
     pub fn hidden_start(&self) -> usize {

diff --git a/parser/src/tokenparser.rs b/parser/src/tokenparser.rs
@@ -609,8 +609,8 @@ impl TokenParser {
             let grm = Arc::clone(&self.compiled_grammars[gen_grammar.grammar.0]);
             let max_tokens = self.parser.grammar().sym_data(symidx).props.max_tokens;
             let parser = Parser::new(grm, gen_grammar)?;
-            let old_parser = std::mem::replace(&mut self.parser, parser);
-            self.parser.set_stats(old_parser.stats().clone());
+            let mut old_parser = std::mem::replace(&mut self.parser, parser);
+            self.parser.take_global_state_from(&mut old_parser);
             let mut entry = ParserStackEntry {
                 parser: old_parser,
                 parser_llm_tokens_offset: self.parser_llm_tokens_offset,
@@ -632,9 +632,8 @@ impl TokenParser {
     fn pop_parser(&mut self) {
         let inner_bytes = self.parser.get_bytes();
         let entry = self.parser_stack.pop().unwrap();
-        let stats = self.parser.stats().clone();
-        self.parser = entry.parser;
-        self.parser.set_stats(stats);
+        let mut prev_parser = std::mem::replace(&mut self.parser, entry.parser);
+        self.parser.take_global_state_from(&mut prev_parser);
         self.parser_llm_tokens_offset = entry.parser_llm_tokens_offset;
         self.previous_grm_bytes
             .truncate(entry.previous_grm_bytes_len);