From e0c72003f9cd21b10640bf2d4d202b06897f9aea Mon Sep 17 00:00:00 2001
From: Daniel M <daniel.q.mueller@stud.h-da.de>
Date: Thu, 3 Feb 2022 14:50:55 +0100
Subject: [PATCH] Implement return values

---
 src/ast.rs         |  1 +
 src/interpreter.rs | 37 ++++++++++++++++++++++++++++---------
 src/lexer.rs       |  1 +
 src/parser.rs      | 13 +++++++++++++
 src/token.rs       |  2 ++
 5 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/src/ast.rs b/src/ast.rs
index f2e98f9..1bdd137 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -119,6 +119,7 @@ pub enum Statement {
     If(If),
     Print(Expression),
     FunDecl(String, Vec<String>, Ast),
+    Return(Expression),
 }
 
 #[derive(Debug, PartialEq, Eq, Clone, Default)]
diff --git a/src/interpreter.rs b/src/interpreter.rs
index ccc01ed..4914ede 100644
--- a/src/interpreter.rs
+++ b/src/interpreter.rs
@@ -8,6 +8,11 @@ pub enum Value {
     String(Rc<String>),
 }
 
+pub enum RunEnd {
+    Return(Value),
+    End,
+}
+
 pub struct Interpreter {
     // Variable table stores the runtime values of variables
     vartable: HashMap<String, Value>,
@@ -36,13 +41,17 @@ impl Interpreter {
         self.run(&ast);
     }
 
-    pub fn run(&mut self, prog: &Ast) {
+    pub fn run(&mut self, prog: &Ast) -> RunEnd {
         for stmt in &prog.prog {
             match stmt {
                 Statement::Expr(expr) => {
                     self.resolve_expr(expr);
                 }
 
+                Statement::Return(expr) => {
+                    return RunEnd::Return(self.resolve_expr(expr));
+                }
+
                 Statement::Loop(looop) => {
                     // loop runs as long condition != 0
                     loop {
@@ -50,7 +59,10 @@ impl Interpreter {
                             break;
                         }
 
-                        self.run(&looop.body);
+                        match self.run(&looop.body) {
+                            RunEnd::Return(val) => return RunEnd::Return(val),
+                            RunEnd::End => (),
+                        }
 
                         if let Some(adv) = &looop.advancement {
                             self.resolve_expr(&adv);
@@ -64,18 +76,23 @@ impl Interpreter {
                 }
 
                 Statement::If(If {condition, body_true, body_false}) => {
-                    if matches!(self.resolve_expr(condition), Value::I64(0)) {
-                        self.run(body_false);
+                    let end = if matches!(self.resolve_expr(condition), Value::I64(0)) {
+                        self.run(body_false)
                     } else {
-                        self.run(body_true);
+                        self.run(body_true)
+                    };
+                    match end {
+                        RunEnd::Return(val) => return RunEnd::Return(val),
+                        RunEnd::End => (),
                     }
                 }
                 Statement::FunDecl(name, args, body) => {
                     self.funtable.insert(name.clone(), (args.clone(), body.clone()).into());
                 }
             }
-
         }
+
+        RunEnd::End
     }
 
     fn resolve_expr(&mut self, expr: &Expression) -> Value {
@@ -96,9 +113,11 @@ impl Interpreter {
                     panic!("Invalid number of arguments for function");
                 }
 
-                self.run(&fun.borrow().1);
-
-                Value::I64(0)
+                let end = self.run(&fun.borrow().1);
+                match end {
+                    RunEnd::Return(val) => val,
+                    RunEnd::End => Value::I64(0),
+                }
             }
         }
     }
diff --git a/src/lexer.rs b/src/lexer.rs
index 1937d42..457c754 100644
--- a/src/lexer.rs
+++ b/src/lexer.rs
@@ -191,6 +191,7 @@ impl<'a> Lexer<'a> {
                         "if" => Token::If,
                         "else" => Token::Else,
                         "fun" => Token::Fun,
+                        "return" => Token::Return,
 
                         // If it doesn't match a keyword, it is a normal identifier
                         _ => Token::Ident(ident),
diff --git a/src/parser.rs b/src/parser.rs
index 713ee1c..58b0a3f 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -52,6 +52,19 @@ impl<T: Iterator<Item = Token>> Parser<T> {
                 Statement::Print(expr)
             }
 
+            Token::Return => {
+                self.next();
+
+                let expr = self.parse_expr();
+
+                // After a statement, there must be a semicolon
+                if !matches!(self.next(), Token::Semicolon) {
+                    panic!("Expected semicolon after statement");
+                }
+
+                Statement::Return(expr)
+            }
+
             Token::If => Statement::If(self.parse_if()),
 
             Token::Fun => {
diff --git a/src/token.rs b/src/token.rs
index 7796033..0576c28 100644
--- a/src/token.rs
+++ b/src/token.rs
@@ -27,6 +27,8 @@ pub enum Token {
 
     Comma,
 
+    Return,
+
     /// Left Parenthesis ('(')
     LParen,