Refactor depth checks and apply them in math (#4845)

6 files changed, 64 insertions, 35 deletions

diff --git a/crates/typst/src/engine.rs b/crates/typst/src/engine.rs
index 2e2525b2..d0ab7d2c 100644
--- a/crates/typst/src/engine.rs
+++ b/crates/typst/src/engine.rs
@@ -7,7 +7,7 @@ use comemo::{Track, Tracked, TrackedMut, Validate};
 use ecow::EcoVec;
 use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
 
-use crate::diag::{SourceDiagnostic, SourceResult};
+use crate::diag::{bail, HintedStrResult, SourceDiagnostic, SourceResult, StrResult};
 use crate::foundations::{Styles, Value};
 use crate::introspection::Introspector;
 use crate::syntax::{FileId, Span};
@@ -229,21 +229,6 @@ pub struct Route<'a> {
     upper: AtomicUsize,
 }
 
-/// The maximum nesting depths. They are different so that even if show rule and
-/// call checks are interleaved, show rule problems we always get the show rule.
-/// The lower the max depth for a kind of error, the higher its precedence
-/// compared to the others.
-impl Route<'_> {
-    /// The maximum stack nesting depth.
-    pub const MAX_SHOW_RULE_DEPTH: usize = 64;
-
-    /// The maximum layout nesting depth.
-    pub const MAX_LAYOUT_DEPTH: usize = 72;
-
-    /// The maximum function call nesting depth.
-    pub const MAX_CALL_DEPTH: usize = 80;
-}
-
 impl<'a> Route<'a> {
     /// Create a new, empty route.
     pub fn root() -> Self {
@@ -297,6 +282,51 @@ impl<'a> Route<'a> {
     }
 }
 
+/// The maximum nesting depths. They are different so that even if show rule and
+/// call checks are interleaved, for show rule problems we always get the show
+/// rule error. The lower the max depth for a kind of error, the higher its
+/// precedence compared to the others.
+impl Route<'_> {
+    /// The maximum stack nesting depth.
+    const MAX_SHOW_RULE_DEPTH: usize = 64;
+
+    /// The maximum layout nesting depth.
+    const MAX_LAYOUT_DEPTH: usize = 72;
+
+    /// The maximum function call nesting depth.
+    const MAX_CALL_DEPTH: usize = 80;
+
+    /// Ensures that we are within the maximum show rule depth.
+    pub fn check_show_depth(&self) -> HintedStrResult<()> {
+        if !self.within(Route::MAX_SHOW_RULE_DEPTH) {
+            bail!(
+                "maximum show rule depth exceeded";
+                hint: "check whether the show rule matches its own output"
+            );
+        }
+        Ok(())
+    }
+
+    /// Ensures that we are within the maximum layout depth.
+    pub fn check_layout_depth(&self) -> HintedStrResult<()> {
+        if !self.within(Route::MAX_LAYOUT_DEPTH) {
+            bail!(
+                "maximum layout depth exceeded";
+                hint: "try to reduce the amount of nesting in your layout",
+            );
+        }
+        Ok(())
+    }
+
+    /// Ensures that we are within the maximum function call depth.
+    pub fn check_call_depth(&self) -> StrResult<()> {
+        if !self.within(Route::MAX_CALL_DEPTH) {
+            bail!("maximum function call depth exceeded");
+        }
+        Ok(())
+    }
+}
+
 #[comemo::track]
 impl<'a> Route<'a> {
     /// Whether the given id is part of the route.
diff --git a/crates/typst/src/eval/call.rs b/crates/typst/src/eval/call.rs
index ee4c4787..4331e187 100644
--- a/crates/typst/src/eval/call.rs
+++ b/crates/typst/src/eval/call.rs
@@ -30,9 +30,7 @@ impl Eval for ast::FuncCall<'_> {
         let args = self.args();
         let trailing_comma = args.trailing_comma();
 
-        if !vm.engine.route.within(Route::MAX_CALL_DEPTH) {
-            bail!(span, "maximum function call depth exceeded");
-        }
+        vm.engine.route.check_call_depth().at(span)?;
 
         // Try to evaluate as a call to an associated function or field.
         let (callee, args) = if let ast::Expr::FieldAccess(access) = callee {
diff --git a/crates/typst/src/layout/flow.rs b/crates/typst/src/layout/flow.rs
index f1148604..f3b105de 100644
--- a/crates/typst/src/layout/flow.rs
+++ b/crates/typst/src/layout/flow.rs
@@ -7,7 +7,7 @@ use std::num::NonZeroUsize;
 
 use comemo::{Track, Tracked, TrackedMut};
 
-use crate::diag::{bail, SourceResult};
+use crate::diag::{bail, At, SourceResult};
 use crate::engine::{Engine, Route, Sink, Traced};
 use crate::foundations::{
     Content, NativeElement, Packed, Resolve, Smart, StyleChain, Styles,
@@ -724,12 +724,7 @@ fn layout_fragment_impl(
         route: Route::extend(route),
     };
 
-    if !engine.route.within(Route::MAX_LAYOUT_DEPTH) {
-        bail!(
-            content.span(), "maximum layout depth exceeded";
-            hint: "try to reduce the amount of nesting in your layout",
-        );
-    }
+    engine.route.check_layout_depth().at(content.span())?;
 
     // If we are in a `PageElem`, this might already be a realized flow.
     let arenas = Arenas::default();
diff --git a/crates/typst/src/math/mod.rs b/crates/typst/src/math/mod.rs
index 383631a4..72b426b5 100644
--- a/crates/typst/src/math/mod.rs
+++ b/crates/typst/src/math/mod.rs
@@ -42,7 +42,7 @@ use self::fragment::*;
 use self::row::*;
 use self::spacing::*;
 
-use crate::diag::SourceResult;
+use crate::diag::{At, SourceResult};
 use crate::foundations::{
     category, Category, Content, Module, Resolve, Scope, SequenceElem, StyleChain,
     StyledElem,
@@ -239,6 +239,9 @@ impl LayoutMath for Content {
         if let Some((tag, realized)) =
             process(ctx.engine, &mut ctx.locator, self, styles)?
         {
+            ctx.engine.route.increase();
+            ctx.engine.route.check_show_depth().at(self.span())?;
+
             if let Some(tag) = &tag {
                 ctx.push(MathFragment::Tag(tag.clone()));
             }
@@ -246,6 +249,8 @@ impl LayoutMath for Content {
             if let Some(tag) = tag {
                 ctx.push(MathFragment::Tag(tag.with_kind(TagKind::End)));
             }
+
+            ctx.engine.route.decrease();
             return Ok(());
         }
 
diff --git a/crates/typst/src/realize/mod.rs b/crates/typst/src/realize/mod.rs
index 5f89e51d..a154fce1 100644
--- a/crates/typst/src/realize/mod.rs
+++ b/crates/typst/src/realize/mod.rs
@@ -15,8 +15,8 @@ pub use self::process::process;
 
 use std::mem;
 
-use crate::diag::{bail, SourceResult};
-use crate::engine::{Engine, Route};
+use crate::diag::{bail, At, SourceResult};
+use crate::engine::Engine;
 use crate::foundations::{
     Content, ContextElem, NativeElement, Packed, SequenceElem, Smart, StyleChain,
     StyleVec, StyledElem, Styles,
@@ -139,12 +139,7 @@ impl<'a, 'v> Builder<'a, 'v> {
             process(self.engine, self.locator, content, styles)?
         {
             self.engine.route.increase();
-            if !self.engine.route.within(Route::MAX_SHOW_RULE_DEPTH) {
-                bail!(
-                    content.span(), "maximum show rule depth exceeded";
-                    hint: "check whether the show rule matches its own output"
-                );
-            }
+            self.engine.route.check_show_depth().at(content.span())?;
 
             if let Some(tag) = &tag {
                 self.accept(self.arenas.store(TagElem::packed(tag.clone())), styles)?;
diff --git a/tests/suite/scripting/recursion.typ b/tests/suite/scripting/recursion.typ
index 43fe848e..6be96c1e 100644
--- a/tests/suite/scripting/recursion.typ
+++ b/tests/suite/scripting/recursion.typ
@@ -53,3 +53,9 @@
 // Hint: 22-25 check whether the show rule matches its own output
 #show math.equation: $x$
 $ x $
+
+--- recursion-show-math-realize ---
+// Error: 22-33 maximum show rule depth exceeded
+// Hint: 22-33 check whether the show rule matches its own output
+#show heading: it => heading[it]
+$ #heading[hi] $