From 8e0f5993f12a590c42dfebfbc99b75dba00daf15 Mon Sep 17 00:00:00 2001
From: Martin Haug <mhaug@live.de>
Date: Mon, 28 Feb 2022 13:41:15 +0100
Subject: [PATCH] Make loops and functions react to control flow

---
 src/eval/func.rs                  |  8 ++++--
 src/eval/mod.rs                   | 16 +++++++++--
 tests/typ/code/break-continue.typ | 44 ++++++++++++++++++++++++++++++-
 tests/typ/code/return.typ         | 11 ++++++--
 4 files changed, 72 insertions(+), 7 deletions(-)

diff --git a/src/eval/func.rs b/src/eval/func.rs
index cd54a140b..b870fd9d6 100644
--- a/src/eval/func.rs
+++ b/src/eval/func.rs
@@ -2,7 +2,7 @@ use std::fmt::{self, Debug, Formatter, Write};
 use std::hash::{Hash, Hasher};
 use std::sync::Arc;
 
-use super::{Cast, Eval, Scope, Scopes, Value};
+use super::{Cast, Control, Eval, Scope, Scopes, Value};
 use crate::diag::{At, TypResult};
 use crate::syntax::ast::Expr;
 use crate::syntax::{Span, Spanned};
@@ -138,7 +138,11 @@ impl Closure {
         }
 
         // Evaluate the body.
-        let value = self.body.eval(ctx, &mut scp)?;
+        let value = match self.body.eval(ctx, &mut scp) {
+            Ok(value) => value,
+            Err(Control::Return(value, _)) => return Ok(value.unwrap_or_default()),
+            other => other?,
+        };
 
         Ok(value)
     }
diff --git a/src/eval/mod.rs b/src/eval/mod.rs
index be168d633..cb80968c7 100644
--- a/src/eval/mod.rs
+++ b/src/eval/mod.rs
@@ -633,7 +633,12 @@ impl Eval for WhileExpr {
         let condition = self.condition();
         while condition.eval(ctx, scp)?.cast::<bool>().at(condition.span())? {
             let body = self.body();
-            let value = body.eval(ctx, scp)?;
+            let value = match body.eval(ctx, scp) {
+                Ok(value) => value,
+                Err(Control::Break(_)) => break,
+                Err(Control::Continue(_)) => continue,
+                other => other?,
+            };
             output = ops::join(output, value).at(body.span())?;
         }
 
@@ -654,7 +659,14 @@ impl Eval for ForExpr {
                 for ($($value),*) in $iter {
                     $(scp.top.def_mut(&$binding, $value);)*
 
-                    let value = self.body().eval(ctx, scp)?;
+                    let body = self.body();
+                    let value = match body.eval(ctx, scp) {
+                        Ok(value) => value,
+                        Err(Control::Break(_)) => break,
+                        Err(Control::Continue(_)) => continue,
+                        other => other?,
+                    };
+
                     output = ops::join(output, value)
                         .at(self.body().span())?;
                 }
diff --git a/tests/typ/code/break-continue.typ b/tests/typ/code/break-continue.typ
index 2b38cf77c..e54651f19 100644
--- a/tests/typ/code/break-continue.typ
+++ b/tests/typ/code/break-continue.typ
@@ -2,13 +2,55 @@
 // Ref: false
 
 ---
+// Test break.
+
+#let error = false
+#let var = 0
+
 #for i in range(10) {
+  var += i
   if i > 5 {
-    // Error: 5-10 break is not yet implemented
     break
+    error = true
   }
 }
 
+#test(error, false)
+#test(var, 21)
+
+---
+// Test continue.
+
+#let x = 0
+#let i = 0
+
+#while x < 8 {
+  i += 1
+
+  if mod(i, 3) == 0 {
+    continue
+  }
+  x += i
+}
+
+// If continue did not work, this would equal 10.
+#test(x, 12)
+
+---
+// Test break outside of loop.
+
+#let f() = {
+  // Error: 3-8 cannot break outside of loop
+  break
+}
+#f()
+
+---
+// Test continue outside of loop.
+
+// Error: 12-20 cannot continue outside of loop
+#let x = { continue }
+
 ---
 // Error: 1-10 unexpected keyword `continue`
 #continue
diff --git a/tests/typ/code/return.typ b/tests/typ/code/return.typ
index bd30c46f4..46ff190c0 100644
--- a/tests/typ/code/return.typ
+++ b/tests/typ/code/return.typ
@@ -3,8 +3,15 @@
 
 ---
 #let f(x) = {
-  // Error: 3-15 return is not yet implemented
   return x + 1
 }
 
-#f(1)
+#test(f(1), 2)
+
+---
+// Test return outside of function.
+
+#for x in range(5) {
+  // Error: 3-9 cannot return outside of function
+  return
+}