Fix control flow bug

3 files changed, 37 insertions, 5 deletions

diff --git a/src/eval/mod.rs b/src/eval/mod.rs
index d39c3c2c..b35cf1ef 100644
--- a/src/eval/mod.rs
+++ b/src/eval/mod.rs
@@ -111,6 +111,7 @@ fn eval_markup(
     scp: &mut Scopes,
     nodes: &mut impl Iterator<Item = MarkupNode>,
 ) -> TypResult<Content> {
+    let flow = ctx.flow.take();
     let mut seq = Vec::with_capacity(nodes.size_hint().1.unwrap_or_default());
 
     while let Some(node) = nodes.next() {
@@ -146,6 +147,10 @@ fn eval_markup(
         }
     }
 
+    if flow.is_some() {
+        ctx.flow = flow;
+    }
+
     Ok(Content::sequence(seq))
 }
 
@@ -343,6 +348,7 @@ fn eval_code(
     scp: &mut Scopes,
     exprs: &mut impl Iterator<Item = Expr>,
 ) -> TypResult<Value> {
+    let flow = ctx.flow.take();
     let mut output = Value::None;
 
     while let Some(expr) = exprs.next() {
@@ -383,6 +389,10 @@ fn eval_code(
         }
     }
 
+    if flow.is_some() {
+        ctx.flow = flow;
+    }
+
     Ok(output)
 }
 
@@ -785,6 +795,7 @@ impl Eval for WhileExpr {
     type Output = Value;
 
     fn eval(&self, ctx: &mut Context, scp: &mut Scopes) -> TypResult<Self::Output> {
+        let flow = ctx.flow.take();
         let mut output = Value::None;
 
         let condition = self.condition();
@@ -804,6 +815,10 @@ impl Eval for WhileExpr {
             }
         }
 
+        if flow.is_some() {
+            ctx.flow = flow;
+        }
+
         Ok(output)
     }
 }
@@ -812,11 +827,12 @@ impl Eval for ForExpr {
     type Output = Value;
 
     fn eval(&self, ctx: &mut Context, scp: &mut Scopes) -> TypResult<Self::Output> {
+        let flow = ctx.flow.take();
+        let mut output = Value::None;
+        scp.enter();
+
         macro_rules! iter {
             (for ($($binding:ident => $value:ident),*) in $iter:expr) => {{
-                let mut output = Value::None;
-                scp.enter();
-
                 #[allow(unused_parens)]
                 for ($($value),*) in $iter {
                     $(scp.top.def_mut(&$binding, $value);)*
@@ -836,8 +852,6 @@ impl Eval for ForExpr {
                     }
                 }
 
-                scp.exit();
-                return Ok(output);
             }};
         }
 
@@ -878,6 +892,13 @@ impl Eval for ForExpr {
                 bail!(self.iter().span(), "cannot loop over {}", iter.type_name());
             }
         }
+
+        if flow.is_some() {
+            ctx.flow = flow;
+        }
+
+        scp.exit();
+        Ok(output)
     }
 }
 
diff --git a/tests/ref/code/break-continue.png b/tests/ref/code/break-continue.png
index 848f4fa6..401b3940 100644
--- a/tests/ref/code/break-continue.png
+++ b/tests/ref/code/break-continue.png
diff --git a/tests/typ/code/break-continue.typ b/tests/typ/code/break-continue.typ
index 02c221a4..2415cb8f 100644
--- a/tests/typ/code/break-continue.typ
+++ b/tests/typ/code/break-continue.typ
@@ -134,3 +134,14 @@
   set text(blue, ..break)
   [Not happening]
 }
+
+---
+// Test second block during break flow.
+// Ref: true
+
+#for i in range(10) {
+  table(
+    { [A]; break },
+    for _ in range(3) [B]
+  )
+}