Implement return control flow

schala-lang/language/src/ast/visitor.rs

@@ -47,9 +47,12 @@ pub fn walk_block<V: ASTVisitor>(v: &mut V, block: &Block) {
                 if let Recursion::Continue = v.module(module_spec) {
                     walk_block(v, &module_spec.contents);
                 },
-            Flow(ref _flow_control) => {
-                //TODO do something with flow control in Visitor
+            Flow(ref flow_control) => match flow_control {
+                FlowControl::Return(Some(ref retval)) => {
+                    walk_expression(v, retval);
                 }
+                _ => (),
+            },
         }
     }
 }

schala-lang/language/src/tree_walk_eval/evaluator.rs

@@ -10,12 +10,23 @@ use crate::{
     util::ScopeStack,
 };
 
+#[derive(Debug)]
+enum StatementOutput {
+    Primitive(Primitive),
+    Nothing,
+}
+
 pub struct Evaluator<'a, 'b> {
-    pub type_context: &'b TypeContext,
-    pub state: &'b mut State<'a>,
+    type_context: &'b TypeContext,
+    state: &'b mut State<'a>,
+    early_returning: bool,
 }
 
 impl<'a, 'b> Evaluator<'a, 'b> {
+    pub(crate) fn new(state: &'b mut State<'a>, type_context: &'b TypeContext) -> Self {
+        Self { state, type_context, early_returning: false }
+    }
+
     pub fn evaluate(&mut self, reduced: ReducedIR, repl: bool) -> Vec<Result<String, String>> {
         let mut acc = vec![];
         for (def_id, function) in reduced.functions.into_iter() {
@@ -25,7 +36,8 @@ impl<'a, 'b> Evaluator<'a, 'b> {
 
         for statement in reduced.entrypoint.into_iter() {
             match self.statement(statement) {
-                Ok(Some(output)) if repl => acc.push(Ok(output.to_repl(self.type_context))),
+                Ok(StatementOutput::Primitive(output)) if repl =>
+                    acc.push(Ok(output.to_repl(self.type_context))),
                 Ok(_) => (),
                 Err(error) => {
                     acc.push(Err(error.msg));
@@ -38,27 +50,38 @@ impl<'a, 'b> Evaluator<'a, 'b> {
 
     fn block(&mut self, statements: Vec<Statement>) -> EvalResult<Primitive> {
         //TODO need to handle breaks, returns, etc.
-        let mut ret = None;
+        let mut retval = None;
         for stmt in statements.into_iter() {
-            if let Some(prim) = self.statement(stmt)? {
-                ret = Some(prim);
+            match self.statement(stmt)? {
+                StatementOutput::Nothing => (),
+                StatementOutput::Primitive(prim) => {
+                    retval = Some(prim);
+                }
+            };
+            if self.early_returning {
+                break;
             }
         }
-        Ok(if let Some(ret) = ret { ret } else { self.expression(Expression::unit())? })
+        Ok(if let Some(ret) = retval { ret } else { self.expression(Expression::unit())? })
     }
 
-    fn statement(&mut self, stmt: Statement) -> EvalResult<Option<Primitive>> {
+    fn statement(&mut self, stmt: Statement) -> EvalResult<StatementOutput> {
+        println!("Statement: {:?}", stmt);
         match stmt {
             Statement::Binding { ref id, expr, constant: _ } => {
                 let evaluated = self.expression(expr)?;
                 self.state.environments.insert(id.into(), evaluated.into());
-                Ok(None)
+                Ok(StatementOutput::Nothing)
             }
             Statement::Expression(expr) => {
                 let evaluated = self.expression(expr)?;
-                Ok(Some(evaluated))
+                Ok(StatementOutput::Primitive(evaluated))
+            }
+            Statement::Return(expr) => {
+                let evaluated = self.expression(expr)?;
+                self.early_returning = true;
+                Ok(StatementOutput::Primitive(evaluated))
             }
-            Statement::Return(expr) => unimplemented!(),
             Statement::Break => unimplemented!(),
             Statement::Continue => unimplemented!(),
         }
@@ -199,9 +222,10 @@ impl<'a, 'b> Evaluator<'a, 'b> {
             let mut new_scope = self.state.environments.new_scope(None);
             if matches(&cond, &alt.pattern, &mut new_scope) {
                 let mut new_state = State { environments: new_scope };
-                let mut evaluator = Evaluator { state: &mut new_state, type_context: self.type_context };
-
-                return evaluator.block(alt.item);
+                let mut evaluator = Evaluator::new(&mut new_state, self.type_context);
+                let output = evaluator.block(alt.item);
+                self.early_returning = evaluator.early_returning;
+                return output;
             }
         }
         Err("No valid match in match expression".into())
@@ -363,7 +387,7 @@ impl<'a, 'b> Evaluator<'a, 'b> {
         }
 
         let mut frame_state = State { environments: self.state.environments.new_scope(None) };
-        let mut evaluator = Evaluator { state: &mut frame_state, type_context: self.type_context };
+        let mut evaluator = Evaluator::new(&mut frame_state, self.type_context);
 
         for (n, evaled) in evaluated_args.into_iter().enumerate() {
             let n = n as u8;

schala-lang/language/src/tree_walk_eval/mod.rs

@@ -172,7 +172,7 @@ impl<'a> State<'a> {
         type_context: &TypeContext,
         repl: bool,
     ) -> Vec<Result<String, String>> {
-        let mut evaluator = evaluator::Evaluator { state: self, type_context };
+        let mut evaluator = evaluator::Evaluator::new(self, type_context);
         evaluator.evaluate(reduced, repl)
     }
 }

schala-lang/language/src/tree_walk_eval/test.rs

@@ -21,7 +21,7 @@ fn evaluate_input(input: &str) -> Result<String, String> {
     symbol_table.debug();
 
     let mut state = State::new();
-    let mut evaluator = Evaluator { state: &mut state, type_context: &type_context };
+    let mut evaluator = Evaluator::new(&mut state, &type_context);
     let mut outputs = evaluator.evaluate(reduced_ir, true);
     outputs.pop().unwrap()
 }
@@ -378,3 +378,18 @@ let value = Klewos::Klewos { a: 50, b: "nah" }
 
     eval_assert(source, r#"(50, "nah")"#);
 }
+
+#[test]
+fn early_return() {
+    let source = r#"
+fn chnurmek(a: Int): Int {
+    if a == 5 then {
+        return 9999;
+    }
+    return (a + 2);
+}
+
+(chnurmek(5), chnurmek(0))
+"#;
+    eval_assert(source, r#"(9999, 2)"#);
+}