Make loops and functions react to control flow

2 files changed, 20 insertions, 4 deletions

diff --git a/src/eval/func.rs b/src/eval/func.rs
index cd54a140..b870fd9d 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 be168d63..cb80968c 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())?;
                 }