fix: Enhance parser recursion depth handling in macro expansion and mhchem parsing

Author: erweixin

crates/ratex-parser/src/macro_expander.rs

@@ -55,6 +55,7 @@ struct MacroExpansion {
 pub struct MacroExpander<'a> {
     pub lexer: Lexer<'a>,
     pub mode: Mode,
+    parser_recursion_depth: usize,
     stack: Vec<Token>,
     macros: MacroNamespace,
     expansion_count: usize,
@@ -164,6 +165,7 @@ impl<'a> MacroExpander<'a> {
         let mut me = Self {
             lexer: Lexer::new(input),
             mode,
+            parser_recursion_depth: 0,
             stack: Vec::new(),
             macros: MacroNamespace::new(),
             expansion_count: 0,
@@ -173,6 +175,10 @@ impl<'a> MacroExpander<'a> {
         me
     }
 
+    pub fn set_parser_recursion_depth(&mut self, depth: usize) {
+        self.parser_recursion_depth = depth;
+    }
+
     fn load_builtins(&mut self) {
         let builtins: &[(&str, &str)] = &[
             // ── Grouping ──
@@ -898,8 +904,12 @@ impl<'a> MacroExpander<'a> {
             MacroDefinition::Function(|me: &mut MacroExpander| -> ParseResult<Vec<Token>> {
                 let args = me.consume_args(1)?;
                 let s = crate::mhchem::mhchem_arg_tokens_to_string(&args[0]);
-                let tex = crate::mhchem::chem_parse_str(&s, "ce")
-                    .map_err(|e| ParseError::msg(format!("\\ce: {e}")))?;
+                let tex = crate::mhchem::chem_parse_str_with_parser_depth(
+                    &s,
+                    "ce",
+                    me.parser_recursion_depth,
+                )
+                .map_err(|e| ParseError::msg(format!("\\ce: {e}")))?;
                 Ok(lex_string_to_stack_tokens(&tex))
             }),
         );
@@ -908,8 +918,12 @@ impl<'a> MacroExpander<'a> {
             MacroDefinition::Function(|me: &mut MacroExpander| -> ParseResult<Vec<Token>> {
                 let args = me.consume_args(1)?;
                 let s = crate::mhchem::mhchem_arg_tokens_to_string(&args[0]);
-                let tex = crate::mhchem::chem_parse_str(&s, "pu")
-                    .map_err(|e| ParseError::msg(format!("\\pu: {e}")))?;
+                let tex = crate::mhchem::chem_parse_str_with_parser_depth(
+                    &s,
+                    "pu",
+                    me.parser_recursion_depth,
+                )
+                .map_err(|e| ParseError::msg(format!("\\pu: {e}")))?;
                 Ok(lex_string_to_stack_tokens(&tex))
             }),
         );

crates/ratex-parser/src/mhchem/engine.rs

@@ -38,7 +38,7 @@ pub(crate) fn go_machine(
     // mhchem inputs fan out through helper machines before texify. Guard the
     // active engine-call depth directly so engine callers fail before building
     // an over-budget AST, while preserving the public parser boundary.
-    if previous_depth >= MAX_RECURSION_DEPTH {
+    if ctx.parser_recursion_depth + previous_depth >= MAX_RECURSION_DEPTH {
         return Err(MhchemError::msg("Recursion limit exceeded"));
     }
     MHCHEM_RECURSION_DEPTH.with(|depth| depth.set(previous_depth + 1));

crates/ratex-parser/src/mhchem/mod.rs

@@ -21,6 +21,7 @@ use serde_json::Value;
 /// Context for recursive `go` (used by actions).
 pub struct ParserCtx<'a> {
     pub data: &'a MhchemData,
+    pub parser_recursion_depth: usize,
 }
 
 impl ParserCtx<'_> {
@@ -31,8 +32,19 @@ impl ParserCtx<'_> {
 
 /// Parse `\ce` / `\pu` argument to TeX fragment (wrap `\mathrm` etc. is done here).
 pub fn chem_parse_str(input: &str, mode: &str) -> MhchemResult<String> {
+    chem_parse_str_with_parser_depth(input, mode, 0)
+}
+
+pub(crate) fn chem_parse_str_with_parser_depth(
+    input: &str,
+    mode: &str,
+    parser_recursion_depth: usize,
+) -> MhchemResult<String> {
     let d = data();
-    let ctx = ParserCtx { data: d };
+    let ctx = ParserCtx {
+        data: d,
+        parser_recursion_depth,
+    };
     let sm = match mode {
         "ce" => "ce",
         "pu" => "pu",
@@ -108,7 +120,10 @@ mod tests {
         let nested_empty_ce =
             |depth: usize| format!("{}{}", r"\ce{".repeat(depth), "}".repeat(depth));
         let d = data();
-        let ctx = ParserCtx { data: d };
+        let ctx = ParserCtx {
+            data: d,
+            parser_recursion_depth: 0,
+        };
 
         assert!(ctx.go(&nested_ce(30), "ce").is_ok());
         let error = ctx.go(&nested_ce(31), "ce").unwrap_err();

crates/ratex-parser/src/parser.rs

@@ -52,6 +52,7 @@ impl<'a> Parser<'a> {
     /// Return the current lookahead token (fetching from gullet if needed).
     pub fn fetch(&mut self) -> ParseResult<Token> {
         if self.next_token.is_none() {
+            self.gullet.set_parser_recursion_depth(self.recursion_depth);
             self.next_token = Some(self.gullet.expand_next_token()?);
         }
         Ok(self.next_token.clone().unwrap())

crates/ratex-parser/src/tests.rs

@@ -1235,6 +1235,16 @@ mod recursion_limit {
         format!("{}H{}", r"\ce{".repeat(depth), "}".repeat(depth))
     }
 
+    fn macro_expanded_nested_ce(enclosing_depth: usize, ce_depth: usize) -> String {
+        format!(
+            r"\def\foo{{{}}}{}{}{}",
+            nested_ce(ce_depth),
+            "{".repeat(enclosing_depth),
+            r"\foo",
+            "}".repeat(enclosing_depth)
+        )
+    }
+
     #[test]
     fn prooftree_depth_is_bounded() {
         assert!(parse(&unary_prooftree(31)).is_ok());
@@ -1264,6 +1274,12 @@ mod recursion_limit {
         assert_recursion_limit_err(&nested_ce(32));
     }
 
+    #[test]
+    fn mhchem_depth_accounts_for_macro_expanded_enclosing_groups() {
+        assert!(parse(&macro_expanded_nested_ce(10, 10)).is_ok());
+        assert_recursion_limit_err(&macro_expanded_nested_ce(20, 20));
+    }
+
     #[test]
     fn source_preflight_ignores_latex_macro_definition_bodies() {
         let deep_body = braced_body(40);