From 4f09233bdae8f79ebafed43e8135f1a0285bd370 Mon Sep 17 00:00:00 2001
From: Martin Haug <mhaug@live.de>
Date: Mon, 28 Feb 2022 14:36:02 +0100
Subject: [PATCH] Enable join collection for control flow constructs

---
 src/eval/control.rs               |  63 +++++++++++++++++++++++
 src/eval/func.rs                  |   2 +-
 src/eval/mod.rs                   |  80 ++++++++++--------------------
 tests/ref/code/return.png         | Bin 0 -> 2946 bytes
 tests/typ/code/break-continue.typ |  37 ++++++++++++--
 tests/typ/code/return.typ         |  40 +++++++++++++++
 6 files changed, 163 insertions(+), 59 deletions(-)
 create mode 100644 src/eval/control.rs
 create mode 100644 tests/ref/code/return.png

diff --git a/src/eval/control.rs b/src/eval/control.rs
new file mode 100644
index 000000000..b310bfb84
--- /dev/null
+++ b/src/eval/control.rs
@@ -0,0 +1,63 @@
+use super::{ops, EvalResult, Value};
+use crate::diag::{At, Error, TypError};
+use crate::syntax::Span;
+
+/// A control flow event that occurred during evaluation.
+#[derive(Clone, Debug, PartialEq)]
+pub enum Control {
+    /// Stop iteration in a loop.
+    Break(Value, Span),
+    /// Skip the remainder of the current iteration in a loop.
+    Continue(Value, Span),
+    /// Stop execution of a function early, returning a value. The bool
+    /// indicates whether this was an explicit return with value.
+    Return(Value, bool, Span),
+    /// Stop the execution because an error occurred.
+    Err(TypError),
+}
+
+impl From<TypError> for Control {
+    fn from(error: TypError) -> Self {
+        Self::Err(error)
+    }
+}
+
+impl From<Control> for TypError {
+    fn from(control: Control) -> Self {
+        match control {
+            Control::Break(_, span) => Error::boxed(span, "cannot break outside of loop"),
+            Control::Continue(_, span) => {
+                Error::boxed(span, "cannot continue outside of loop")
+            }
+            Control::Return(_, _, span) => {
+                Error::boxed(span, "cannot return outside of function")
+            }
+            Control::Err(e) => e,
+        }
+    }
+}
+
+/// Join a value with an evaluated result.
+pub(super) fn join_result(
+    prev: Value,
+    result: EvalResult<Value>,
+    result_span: Span,
+) -> EvalResult<Value> {
+    match result {
+        Ok(value) => Ok(ops::join(prev, value).at(result_span)?),
+        Err(Control::Break(value, span)) => Err(Control::Break(
+            ops::join(prev, value).at(result_span)?,
+            span,
+        )),
+        Err(Control::Continue(value, span)) => Err(Control::Continue(
+            ops::join(prev, value).at(result_span)?,
+            span,
+        )),
+        Err(Control::Return(value, false, span)) => Err(Control::Return(
+            ops::join(prev, value).at(result_span)?,
+            false,
+            span,
+        )),
+        other => other,
+    }
+}
diff --git a/src/eval/func.rs b/src/eval/func.rs
index a7f2a209c..451dcbbb7 100644
--- a/src/eval/func.rs
+++ b/src/eval/func.rs
@@ -139,7 +139,7 @@ impl Closure {
 
         // Evaluate the body.
         let value = match self.body.eval(ctx, &mut scp) {
-            Err(Control::Return(value, _)) => value.unwrap_or_default(),
+            Err(Control::Return(value, _, _)) => value,
             other => other?,
         };
 
diff --git a/src/eval/mod.rs b/src/eval/mod.rs
index 3686d665e..380e9e9d1 100644
--- a/src/eval/mod.rs
+++ b/src/eval/mod.rs
@@ -11,6 +11,7 @@ mod styles;
 mod capture;
 mod class;
 mod collapse;
+mod control;
 mod func;
 mod layout;
 mod module;
@@ -23,6 +24,7 @@ pub use array::*;
 pub use capture::*;
 pub use class::*;
 pub use collapse::*;
+pub use control::*;
 pub use dict::*;
 pub use func::*;
 pub use layout::*;
@@ -35,7 +37,7 @@ pub use value::*;
 
 use unicode_segmentation::UnicodeSegmentation;
 
-use crate::diag::{At, Error, StrResult, Trace, Tracepoint, TypError, TypResult};
+use crate::diag::{At, Error, StrResult, Trace, Tracepoint, TypResult};
 use crate::geom::{Angle, Fractional, Length, Relative};
 use crate::library;
 use crate::syntax::ast::*;
@@ -55,40 +57,6 @@ pub trait Eval {
 /// The result type for evaluating a syntactic construct.
 pub type EvalResult<T> = Result<T, Control>;
 
-/// A control flow event that occurred during evaluation.
-#[derive(Clone, Debug, PartialEq)]
-pub enum Control {
-    /// Stop iteration in a loop.
-    Break(Span),
-    /// Skip the remainder of the current iteration in a loop.
-    Continue(Span),
-    /// Stop execution of a function early, optionally returning a value.
-    Return(Option<Value>, Span),
-    /// Stop the execution because an error occurred.
-    Err(TypError),
-}
-
-impl From<TypError> for Control {
-    fn from(error: TypError) -> Self {
-        Self::Err(error)
-    }
-}
-
-impl From<Control> for TypError {
-    fn from(control: Control) -> Self {
-        match control {
-            Control::Break(span) => Error::boxed(span, "cannot break outside of loop"),
-            Control::Continue(span) => {
-                Error::boxed(span, "cannot continue outside of loop")
-            }
-            Control::Return(_, span) => {
-                Error::boxed(span, "cannot return outside of function")
-            }
-            Control::Err(e) => e,
-        }
-    }
-}
-
 impl Eval for Markup {
     type Output = Template;
 
@@ -337,12 +305,10 @@ impl Eval for BlockExpr {
 
         let mut output = Value::None;
         for expr in self.exprs() {
-            let value = expr.eval(ctx, scp)?;
-            output = ops::join(output, value).at(expr.span())?;
+            output = join_result(output, expr.eval(ctx, scp), expr.span())?;
         }
 
         scp.exit();
-
         Ok(output)
     }
 }
@@ -637,12 +603,14 @@ impl Eval for WhileExpr {
         let condition = self.condition();
         while condition.eval(ctx, scp)?.cast::<bool>().at(condition.span())? {
             let body = self.body();
-            let value = match body.eval(ctx, scp) {
-                Err(Control::Break(_)) => break,
-                Err(Control::Continue(_)) => continue,
-                other => other?,
-            };
-            output = ops::join(output, value).at(body.span())?;
+            match join_result(output, body.eval(ctx, scp), body.span()) {
+                Err(Control::Break(value, _)) => {
+                    output = value;
+                    break;
+                }
+                Err(Control::Continue(value, _)) => output = value,
+                other => output = other?,
+            }
         }
 
         Ok(output)
@@ -663,13 +631,14 @@ impl Eval for ForExpr {
                     $(scp.top.def_mut(&$binding, $value);)*
 
                     let body = self.body();
-                    let value = match body.eval(ctx, scp) {
-                        Err(Control::Break(_)) => break,
-                        Err(Control::Continue(_)) => continue,
-                        other => other?,
-                    };
-
-                    output = ops::join(output, value).at(body.span())?;
+                    match join_result(output, body.eval(ctx, scp), body.span()) {
+                        Err(Control::Break(value, _)) => {
+                            output = value;
+                            break;
+                        }
+                        Err(Control::Continue(value, _)) => output = value,
+                        other => output = other?,
+                    }
                 }
 
                 scp.exit();
@@ -783,7 +752,7 @@ impl Eval for BreakExpr {
     type Output = Value;
 
     fn eval(&self, _: &mut Context, _: &mut Scopes) -> EvalResult<Self::Output> {
-        Err(Control::Break(self.span()))
+        Err(Control::Break(Value::default(), self.span()))
     }
 }
 
@@ -791,7 +760,7 @@ impl Eval for ContinueExpr {
     type Output = Value;
 
     fn eval(&self, _: &mut Context, _: &mut Scopes) -> EvalResult<Self::Output> {
-        Err(Control::Continue(self.span()))
+        Err(Control::Continue(Value::default(), self.span()))
     }
 }
 
@@ -799,8 +768,11 @@ impl Eval for ReturnExpr {
     type Output = Value;
 
     fn eval(&self, ctx: &mut Context, scp: &mut Scopes) -> EvalResult<Self::Output> {
+        let value = self.body().map(|body| body.eval(ctx, scp)).transpose()?;
+        let explicit = value.is_some();
         Err(Control::Return(
-            self.body().map(|body| body.eval(ctx, scp)).transpose()?,
+            value.unwrap_or_default(),
+            explicit,
             self.span(),
         ))
     }
diff --git a/tests/ref/code/return.png b/tests/ref/code/return.png
new file mode 100644
index 0000000000000000000000000000000000000000..504108872c14991cffc2f4912c9f0ace5089b615
GIT binary patch
literal 2946
zcmZuzX*8690{zC!Fw9^EV^1@(Co0O4oszO;%bKz8d-jpBk7X={5<(>*5@j396tb2E
z8Ecko*|#D~UhkcE-g)PpcYoY-@45f)y@`hU+RRM6OaK5d>s;3|0ss)?uTF;jCE)3H
zy$=Aub#%1UO+v=jCPG6}Xabb2oxE<L4??onwTL6M8yT)r@3N%oGY2l8D_yHq?_j%8
ze$F0Fa*{AYTS*#cT3SxHCQ1}%e`IrSu-D1ZN`UETIbWPQ?3{}bSV){#p{?b&|31!G
z><j-kmHVK7r)8sQqgz26hL$jSX;`EY$Ap#;kO}G0z=}cQ3XvSq8ip_?_5U%kCvqQ1
zQBkB<4bq*;8P4atBaP1vgs0bF)XI=&*NSynGp5bqC{w{RUT)QU>~3aP(j>)>z9P&g
zRPRc}T0#K?dk6L{6McQ&iPqz+EJl9z@Th0@BpK}62XBB<@3kaN7*oqhfb&!k!wI?e
z+`v4L+@av2L0mQ$W(X-65Xqoi9lQkG8=<E#8dq3~Jm_{*grp0|g}op7WhcIB1@;}h
zJN)PS=9snEh?3n{>S0<0H4~n9X%aG^=RBY*cind2d9@oeVoSnWH7qQ?2BWo>Wjy^-
zN>N8>Gr^lzE!9p47<@MY!@R-&lp?hYVfuh{!CB1J_Bf`Dk8!w&!T3gjx^?AW&VHOr
zV)sDtzpWl&Ys|&sGD$^FCF=8hOJ$3Sa609l+#lA*`uBASEF_RwVjEz-^H{FGx3`xQ
zmA<iK4@wy-2vQ3PT3S#JcC%R#I$BaW`sN^-=-_L((ZX1*<lM#=2-M_3&3k0|ov>!!
z{wxrl4mVm?!Q93Bj&zbJ?WoVBb+V|r7KES8mX4uagCTYX7@y-L%v46YCsORT&h9c*
z@UhCAort#d0VUdVrAyk<r}l=ps1va3dURevv7UQxeDzMR5djx1$tnZMpeBAmWZ54u
z>#?yE^fg<ZeY=iJE;TPpi(A<az&)1YGZzX5t9E;Rjs|j`Ttr==^EuFbCG`mfddcGt
zUjg=pW|0S7PtD7O>3Hn<V#c;+VJ$&qI@P3nwSIbF<_>TzXI&CTnS1yJc{zD9_AUhl
zlAp0HD)wn`;z&3@GG90MPdL_O!{!owq{+f;V-K^Bqg+GTH#fk@Ax2CetBX3|A4x-2
znegKErNmig9tyCRNVOu6m1JiAK_o$8`~LR~P_x@zI5#!0Ko=B7gI=7sE-dWR21t!p
z*rSh~rPzUZvs|;w<0veh_d}cKUPsoR&}?@Y<@ZGgK!t}DIxh3>xXE6L*0yWr6GZ=#
zhshHiPH6K#<vT9iV}bgWm{cS$s6{yyW6)n!#p>+z;q*Ly!K+|cT@&DzKJ*zqZ<rdZ
zFbh4WP-v;iC%HAS0oiuidv{*D&jJFU51l69Bgw{?JULxj5}Ene?NUYi<XXnDx1d+?
zR0<b*Ww>IcWJd{VRd#*U?=@CHpF62#e&KdTwZCkYVJ-E(6>m0p9DxOzY>W7gz72ih
z7AynJy_3Z+alleRO=}M4eks@gVkUqaGl0d{io*fFqtO(jACY1APcFfNo3C+@n**s9
z>xb{#cDRNbECx_SKB+WLE2{0s{F?^jwEp>ogoQ<fve`g*Hf7w$e+ylA5DLROcEy*F
z)2RN!^BUedSi@5l|7n5`wsF8p%x9%fuo4#VZh&6IW7JqH$d`5`Ys-?oO)Wlt0t`P-
zM0t1u76EOdHr+%6aQ3l(_&Ux(pvpU)n>X^Y#}*f8;5JLL3M*w>S<X40Z|5e1hBsHl
zAn#K$vy@3N{M8RRt0(byFDQK3v#?2r=GK+Va^F9n$T1zxx{>apbPv+E{r%wQa`&BL
z@C>@Zao=ctI&6)T%Py{@gt>(a+lx_`1Rt0%aGO>lYa8*dSaRT`#`R3SI(6wBD_-QU
z9#fzbYk49KYy?q+E^tvCV1})tYU5@&j{y{}MhQMRr-ZroB8xt@u6JP+UQ1vB%=V~i
zbv>NGhGEyR9@z5<YBS%Oxk}qRfp(S?_-+lq0oK5+JpBL-pE|to`;y^L8-nMg`G_!E
zl(D^nZgviE(}`)n6})+)#>mfL^Gl;%a1RN1vyHGu?EpbZ9p`U8jxwEvG#R4_z7O%0
z-3ma~pIZ{hCP7eTn@9uS^zUs-uss8%%U=%jR)iv^iQV-;;L69KwTY7y=$k?l%v)yq
z>N35>mnqPvy|oM?Nv0Hqv^_#&Gia=q-Z@)`L4_vV#xd!IYnF#uPxzL^p7G3xIB6)<
zd53Vl<3|1164PYasZKkHvK8Wuh^~IWyq&ZK+mCOL9xWf0ON{w>K`kvFxMly*&(-ck
zLcF`5v9~F16%UdhG+92FGuZWjo6@iDvNzHy!6{oi`msEJDhOb`Vvnvmpg#`G@4nj1
zekZ`fY5+Cuu73;rY^i|yumjH|J<n0MK^pJyk|aqR=63ukX*Pzc3{HiPp$ZN7zBa!&
z;70%70s8keg(yLCM?KF<e3PAu81jos{+4EGNOu_;#_DAA;JYvsYf5~*S}eOao6S<c
z*FR-%Knxx7X*F&++}$8j*TlqOr5gLA+F}S~ni`QkyHiu%sQ?>X{@j}7dD3TpdWz0N
zsWi3ySwGZUw`j|%7Z%~7+Qj9BEa>j$q#&ETM-zfF_H~S<(`C0H3|22&Ui!$ZNvn+5
za8bpxEF!2WFmP_VLjfqi9pK*m8W2ihZzf%PMh#)ceec@igk01A;5-4{tf4Zg@>S4u
z`3RXg$0Y(mOtFrW+EoUhZ?Z8uE{OsYcTe9o@wr=mR3{;ItEz2UqK79aT|z{b@f8*?
z4d(Hhuxy>9ZX%EzO&X%B0ifJ$v^RZBQOb4tux0Y@Pul6h7rgA^+ZZE$NomFR_&P!>
zZBY9Pm@cBnrr0Y#S||g^UjZaDL|4Ba&Wpio#qHrpyH<Y37?~=gcyQ4-?P>c@W@_mM
z{T1R1;nE2q8Hn<^u&(mctHdv%Uo&J`Rs`WQXBs+rkW!5aKIo}*AvRo=Jka!#bTzCY
z6^*HHh*v8+f)w77kr!SBOJ=F68=984DP!*lY1TKv5rN%)YVyoqnrqQj`$2*9kdO7l
zEbmj(S1!;=k^98UD)&VI&7V+Sjj|Nb(gIi{$MJIRrz=jY$K{)@u<rBf&R}pahdgEu
zwBPLgH(99tqaT%SdTU)=1M}(z_QJu_j#J*m_W)kvvpi;ZJHC6Y2jw6TOSz}elFptL
zrg(=Xlyd5v&Fv)wtZzLkzRJe4WVHP1J2P4okn+G;8{+dwI=J{>a8ibWvdPNEntK`|
z=I!S@W!@v*4GaR3WTMDivJj>ShP+IXipHo2J(8Tn3~`!S^}1KsaSXESMAmcE=rBF4
zx*53U3Zn$9<;^Roy9P`oBe~~ya!M}S=kOeyq5SrDNYbpgs)@kj@IODtARld1k`87p
zTCH#%;+G(g%E#f35buO3y-@wQ&D`Ck5--74ow!7gl5kohV>>ZT^-~L16zO2G)sosj
zU%`;iSY<YhFAbo9{6eeL;`?kTN=fkmhtlo!?}6+qp-*DN48pYLB4d5UDKF7`&{yEX
zjG=FMLf)oC^)}f8_;CtMSLrW~!&0Uffbb{MXYeH)bd!y4JidxJXqDnLWR8KM&K+%Z
zv$onk*IM1o`m!+Q^b5$yqRt%4R@Q-(D8bu_Ex%lKroz4UP{l>rb1^!B)dqfqAQ-g!
z0I4!W!phY&u26Zj#*@Vn(XR-7bpHHQ@-yKw7jX&Q-~*jYF;~zY_y8fNlMl;Lal%YR
zcDRWL?oi+7WO8QpSeI(Sw#r8<9?q}SvasC$+Y6Z{kpH@|Nr{o*n-c~^w^L%6zuy|r
MxvH;Ksez0B2QuAAj{pDw

literal 0
HcmV?d00001

diff --git a/tests/typ/code/break-continue.typ b/tests/typ/code/break-continue.typ
index e54651f19..60dac44d3 100644
--- a/tests/typ/code/break-continue.typ
+++ b/tests/typ/code/break-continue.typ
@@ -4,8 +4,8 @@
 ---
 // Test break.
 
-#let error = false
 #let var = 0
+#let error = false
 
 #for i in range(10) {
   var += i
@@ -15,18 +15,32 @@
   }
 }
 
-#test(error, false)
 #test(var, 21)
+#test(error, false)
+
+---
+// Test joining with break.
+
+#let i = 0
+#let x = while true {
+  i += 1
+  str(i)
+  if i >= 5 {
+    "."
+    break
+  }
+}
+
+#test(x, "12345.")
 
 ---
 // Test continue.
 
-#let x = 0
 #let i = 0
+#let x = 0
 
 #while x < 8 {
   i += 1
-
   if mod(i, 3) == 0 {
     continue
   }
@@ -36,6 +50,20 @@
 // If continue did not work, this would equal 10.
 #test(x, 12)
 
+---
+// Test joining with continue.
+
+#let x = for i in range(5) {
+  "a"
+  if mod(i, 3) == 0 {
+    "_"
+    continue
+  }
+  str(i)
+}
+
+#test(x, "a_a1a2a_a4")
+
 ---
 // Test break outside of loop.
 
@@ -43,6 +71,7 @@
   // Error: 3-8 cannot break outside of loop
   break
 }
+
 #f()
 
 ---
diff --git a/tests/typ/code/return.typ b/tests/typ/code/return.typ
index 46ff190c0..9ee3aed75 100644
--- a/tests/typ/code/return.typ
+++ b/tests/typ/code/return.typ
@@ -2,12 +2,52 @@
 // Ref: false
 
 ---
+// Test return with value.
 #let f(x) = {
   return x + 1
 }
 
 #test(f(1), 2)
 
+---
+// Test return with joining.
+
+#let f(x) = {
+  "a"
+  if x == 0 {
+    return "b"
+  } else if x == 1 {
+    "c"
+  } else {
+    "d"
+    return
+    "e"
+  }
+}
+
+#test(f(0), "b")
+#test(f(1), "ac")
+#test(f(2), "ad")
+
+---
+// Test return with joining and template.
+// Ref: true
+
+#let f(text, caption: none) = {
+  text
+  if caption == none {
+    [\.]
+    return
+  }
+  [, ]
+  emph(caption)
+  [\.]
+}
+
+#f(caption: [with caption])[My figure]
+
+#f[My other figure]
+
 ---
 // Test return outside of function.