Implement functions

- Implement function declaration and call
- Change the precalculated variable stack positions to contain the
  offset from the end instead of the absolute position. This is
  important for passing fun args on the stack
- Add the ability to offset the stackframes. This is used to delete the
  stack where the fun args have been stored before the block executes
- Implement exit type for blocks in interpreter. This is used to get the
  return values and propagate them where needed
- Add recursive fibonacci examples

examples/recursive_fib.nek

@@ -0,0 +1,9 @@
+fun fib(n) {
+    if n <= 1 {
+        return n;
+    } else {
+        return fib(n-1) + fib(n-2);
+    }
+}
+
+print fib(30);

examples/recursive_fib.py

@@ -0,0 +1,6 @@
+def fib(n):
+    if n <= 1:
+        return n
+    return fib(n-1) + fib(n-2)
+
+print(fib(30))

src/ast.rs

@@ -1,3 +1,5 @@
+use std::rc::Rc;
+
 use crate::stringstore::{StringStore, Sid};
 
 /// Types for binary operators
@@ -86,6 +88,8 @@ pub enum Expression {
     /// Array access with name, stackpos and position
     ArrayAccess(Sid, usize, Box<Expression>),
 
+    FunCall(Sid, usize, Vec<Expression>),
+
     /// Variable
     Var(Sid, usize),
     /// Binary operation. Consists of type, left hand side and right hand side
@@ -114,9 +118,19 @@ pub struct If {
     pub body_false: BlockScope,
 }
 
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub struct FunDecl {
+    pub name: Sid,
+    pub fun_stackpos: usize,
+    pub argnames: Vec<Sid>,
+    pub body: Rc<BlockScope>,
+}
+
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub enum Statement {
+    Return(Expression),
     Declaration(Sid, usize, Expression),
+    FunDeclare(FunDecl),
     Expr(Expression),
     Block(BlockScope),
     Loop(Loop),

src/astoptimizer.rs

@@ -41,6 +41,8 @@ impl SimpleAstOptimizer {
                 }
                 Statement::Print(expr) => Self::optimize_expr(expr),
                 Statement::Declaration(_, _, expr) => Self::optimize_expr(expr),
+                Statement::FunDeclare(_) => (),
+                Statement::Return(expr) => Self::optimize_expr(expr),
             }
         }
     }

src/interpreter.rs

@@ -1,8 +1,8 @@
-use std::cell::RefCell;
+use std::{cell::RefCell, rc::Rc};
 use thiserror::Error;
 
 use crate::{
-    ast::{Ast, BinOpType, BlockScope, Expression, If, Statement, UnOpType},
+    ast::{Ast, BinOpType, BlockScope, Expression, FunDecl, If, Statement, UnOpType},
     astoptimizer::{AstOptimizer, SimpleAstOptimizer},
     lexer::lex,
     parser::parse,
@@ -25,6 +25,8 @@ pub enum RuntimeError {
     ArrayOutOfBounds(usize, usize),
     #[error("Division by zero")]
     DivideByZero,
+    #[error("Invalid number of arguments for function {0}. Expected {1}, got {2}")]
+    InvalidNumberOfArgs(String, usize, usize),
 }
 
 #[derive(Debug, PartialEq, Eq, Clone)]
@@ -32,6 +34,13 @@ pub enum Value {
     I64(i64),
     String(Sid),
     Array(RefCell<Vec<Value>>),
+    Void,
+}
+
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub enum BlockExit {
+    Normal,
+    Return(Value),
 }
 
 #[derive(Default)]
@@ -47,6 +56,8 @@ pub struct Interpreter {
     // Variable table stores the runtime values of variables
     vartable: Vec<Value>,
 
+    funtable: Vec<FunDecl>,
+
     stringstore: StringStore,
 }
 
@@ -63,10 +74,11 @@ impl Interpreter {
     }
 
     fn get_var(&self, idx: usize) -> Option<Value> {
-        self.vartable.get(idx).cloned()
+        self.vartable.get(self.vartable.len() - idx - 1).cloned()
     }
 
     fn get_var_mut(&mut self, idx: usize) -> Option<&mut Value> {
+        let idx = self.vartable.len() - idx - 1;
         self.vartable.get_mut(idx)
     }
 
@@ -92,14 +104,30 @@ impl Interpreter {
 
         self.stringstore = ast.stringstore;
 
-        self.run_block(&ast.main)
+        self.run_block(&ast.main)?;
+        Ok(())
     }
 
-    pub fn run_block(&mut self, prog: &BlockScope) -> Result<(), RuntimeError> {
-        let framepointer = self.vartable.len();
+    pub fn run_block(&mut self, prog: &BlockScope) -> Result<BlockExit, RuntimeError> {
+        self.run_block_fp_offset(prog, 0)
+    }
+
+    pub fn run_block_fp_offset(
+        &mut self,
+        prog: &BlockScope,
+        framepointer_offset: usize,
+    ) -> Result<BlockExit, RuntimeError> {
+        let framepointer = self.vartable.len() - framepointer_offset;
 
         for stmt in prog {
             match stmt {
+                Statement::Return(expr) => {
+                    let val = self.resolve_expr(expr)?;
+
+                    self.vartable.truncate(framepointer);
+                    return Ok(BlockExit::Return(val));
+                }
+
                 Statement::Expr(expr) => {
                     self.resolve_expr(expr)?;
                 }
@@ -110,7 +138,13 @@ impl Interpreter {
                 }
 
                 Statement::Block(block) => {
-                    self.run_block(block)?;
+                    match self.run_block(block)? {
+                        BlockExit::Return(val) => {
+                            self.vartable.truncate(framepointer);
+                            return Ok(BlockExit::Return(val));
+                        }
+                        _ => ()
+                    }
                 }
 
                 Statement::Loop(looop) => {
@@ -143,18 +177,29 @@ impl Interpreter {
                     body_true,
                     body_false,
                 }) => {
-                    if matches!(self.resolve_expr(condition)?, Value::I64(0)) {
-                        self.run_block(body_false)?;
+                    let exit = if matches!(self.resolve_expr(condition)?, Value::I64(0)) {
+                        self.run_block(body_false)?
                     } else {
-                        self.run_block(body_true)?;
+                        self.run_block(body_true)?
+                    };
+                    match exit {
+                        BlockExit::Return(val) => {
+                            self.vartable.truncate(framepointer);
+                            return Ok(BlockExit::Return(val));
+                        }
+                        _ => (),
                     }
                 }
+
+                Statement::FunDeclare(fundec) => {
+                    self.funtable.push(fundec.clone());
+                }
             }
         }
 
         self.vartable.truncate(framepointer);
 
-        Ok(())
+        Ok(BlockExit::Normal)
     }
 
     fn resolve_expr(&mut self, expr: &Expression) -> Result<Value, RuntimeError> {
@@ -174,6 +219,29 @@ impl Interpreter {
             Expression::ArrayAccess(name, idx, arr_idx) => {
                 self.resolve_array_access(*name, *idx, arr_idx)?
             }
+
+            Expression::FunCall(fun_name, fun_stackpos, args) => {
+                for arg in args {
+                    let arg = self.resolve_expr(arg)?;
+                    self.vartable.push(arg);
+                }
+
+                // Function existance has been verified in the parser, so unwrap here shouldn't fail
+                let num_args = self.funtable.get(*fun_stackpos).unwrap().argnames.len();
+
+                if num_args != args.len() {
+                    let fun_name = self.stringstore.lookup(*fun_name).cloned().unwrap_or("<unknown>".to_string());
+                    return Err(RuntimeError::InvalidNumberOfArgs(fun_name, num_args, args.len()));
+                }
+
+                match self.run_block_fp_offset(
+                    &Rc::clone(&self.funtable.get(*fun_stackpos).unwrap().body),
+                    num_args,
+                )? {
+                    BlockExit::Normal => Value::Void,
+                    BlockExit::Return(val) => val,
+                }
+            }
         };
 
         Ok(val)
@@ -310,8 +378,12 @@ impl Interpreter {
                 BinOpType::Add => Value::I64(lhs + rhs),
                 BinOpType::Mul => Value::I64(lhs * rhs),
                 BinOpType::Sub => Value::I64(lhs - rhs),
-                BinOpType::Div => Value::I64(lhs.checked_div(rhs).ok_or(RuntimeError::DivideByZero)?),
-                BinOpType::Mod => Value::I64(lhs.checked_rem(rhs).ok_or(RuntimeError::DivideByZero)?),
+                BinOpType::Div => {
+                    Value::I64(lhs.checked_div(rhs).ok_or(RuntimeError::DivideByZero)?)
+                }
+                BinOpType::Mod => {
+                    Value::I64(lhs.checked_rem(rhs).ok_or(RuntimeError::DivideByZero)?)
+                }
                 BinOpType::BOr => Value::I64(lhs | rhs),
                 BinOpType::BAnd => Value::I64(lhs & rhs),
                 BinOpType::BXor => Value::I64(lhs ^ rhs),
@@ -338,7 +410,8 @@ impl Interpreter {
         match val {
             Value::I64(val) => format!("{}", val),
             Value::Array(val) => format!("{:?}", val.borrow()),
-            Value::String(text) => format!("{}", self.stringstore.lookup(*text).unwrap()),
+            Value::String(text) => format!("{}", self.stringstore.lookup(*text).unwrap_or(&"<invalid string>".to_string())),
+            Value::Void => format!("void"),
         }
     }
 }

src/lexer.rs

@@ -70,6 +70,7 @@ impl<'a> Lexer<'a> {
                 ('|', '|') => self.push_tok_consume(T![||]),
 
                 // Single character tokens
+                (',', _) => self.push_tok(T![,]),
                 (';', _) => self.push_tok(T![;]),
                 ('+', _) => self.push_tok(T![+]),
                 ('-', _) => self.push_tok(T![-]),
@@ -199,6 +200,8 @@ impl<'a> Lexer<'a> {
             "print" => T![print],
             "if" => T![if],
             "else" => T![else],
+            "fun" => T![fun],
+            "return" => T![return],
 
             // If it doesn't match a keyword, it is a normal identifier
             _ => T![ident(ident)],

src/parser.rs

@@ -1,7 +1,7 @@
 use thiserror::Error;
 
 use crate::{
-    ast::{Ast, BlockScope, Expression, If, Loop, Statement},
+    ast::{Ast, BlockScope, Expression, FunDecl, If, Loop, Statement},
     stringstore::{Sid, StringStore},
     token::Token,
     util::{PutBackIter, PutBackableExt},
@@ -16,6 +16,10 @@ pub enum ParseErr {
     DeclarationOfNonVar,
     #[error("Use of undefined variable \"{0}\"")]
     UseOfUndeclaredVar(String),
+    #[error("Use of undefined function \"{0}\"")]
+    UseOfUndeclaredFun(String),
+    #[error("Redeclation of function \"{0}\"")]
+    RedeclarationFun(String),
 }
 
 type ResPE<T> = Result<T, ParseErr>;
@@ -39,6 +43,7 @@ struct Parser<T: Iterator<Item = Token>> {
     tokens: PutBackIter<T>,
     string_store: StringStore,
     var_stack: Vec<Sid>,
+    fun_stack: Vec<Sid>,
 }
 
 impl<T: Iterator<Item = Token>> Parser<T> {
@@ -47,10 +52,12 @@ impl<T: Iterator<Item = Token>> Parser<T> {
         let tokens = tokens.into_iter().putbackable();
         let string_store = StringStore::new();
         let var_stack = Vec::new();
+        let fun_stack = Vec::new();
         Self {
             tokens,
             string_store,
             var_stack,
+            fun_stack,
         }
     }
 
@@ -62,10 +69,14 @@ impl<T: Iterator<Item = Token>> Parser<T> {
         })
     }
 
+    fn parse_scoped_block(&mut self) -> ResPE<BlockScope> {
+        self.parse_scoped_block_fp_offset(0)
+    }
+
     /// Parse tokens into an abstract syntax tree. This will continuously parse statements until
     /// encountering end-of-file or a block end '}' .
-    fn parse_scoped_block(&mut self) -> ResPE<BlockScope> {
-        let framepointer = self.var_stack.len();
+    fn parse_scoped_block_fp_offset(&mut self, framepoint_offset: usize) -> ResPE<BlockScope> {
+        let framepointer = self.var_stack.len() - framepoint_offset;
         let mut prog = Vec::new();
 
         loop {
@@ -109,8 +120,80 @@ impl<T: Iterator<Item = Token>> Parser<T> {
                 Statement::Print(expr)
             }
 
+            T![return] => {
+                self.next();
+                let stmt = Statement::Return(self.parse_expr()?);
+                
+                // After a statement, there must be a semicolon
+                validate_next!(self, T![;], ";");
+
+                stmt
+            }
+
             T![if] => Statement::If(self.parse_if()?),
 
+            T![fun] => {
+                self.next();
+
+                let fun_name = match self.next() {
+                    T![ident(fun_name)] => fun_name,
+                    tok => return Err(ParseErr::UnexpectedToken(tok, "<ident>".to_string())),
+                };
+
+                let fun_name = self.string_store.intern_or_lookup(&fun_name);
+
+                if self.fun_stack.contains(&fun_name) {
+                    return Err(ParseErr::RedeclarationFun(
+                        self.string_store
+                            .lookup(fun_name)
+                            .cloned()
+                            .unwrap_or("<unknown>".to_string()),
+                    ));
+                }
+
+                let fun_stackpos = self.fun_stack.len();
+                self.fun_stack.push(fun_name);
+
+                let mut arg_names = Vec::new();
+
+                validate_next!(self, T!['('], "(");
+
+                while matches!(self.peek(), T![ident(_)]) {
+                    let var_name = match self.next() {
+                        T![ident(var_name)] => var_name,
+                        _ => unreachable!(),
+                    };
+
+                    let var_name = self.string_store.intern_or_lookup(&var_name);
+                    arg_names.push(var_name);
+
+                    // Push the variable onto the varstack
+                    self.var_stack.push(var_name);
+
+                    // If there are more args skip the comma so that the loop will read the argname
+                    if self.peek() == &T![,] {
+                        self.next();
+                    }
+                }
+
+                validate_next!(self, T![')'], ")");
+
+                validate_next!(self, T!['{'], "{");
+
+                // Create the scoped block with a stack offset. This will pop the args that are
+                // added to the stack while parsing args
+                let body = self.parse_scoped_block_fp_offset(arg_names.len())?;
+
+                validate_next!(self, T!['}'], "}");
+
+                Statement::FunDeclare(FunDecl {
+                    name: fun_name,
+                    fun_stackpos,
+                    argnames: arg_names,
+                    body: body.into(),
+                })
+            }
+
             _ => {
                 let first = self.next();
 
@@ -118,11 +201,12 @@ impl<T: Iterator<Item = Token>> Parser<T> {
                     (T![ident(name)], T![<-]) => {
                         self.next();
 
+                        let rhs = self.parse_expr()?;
+
                         let sid = self.string_store.intern_or_lookup(&name);
                         let sp = self.var_stack.len();
                         self.var_stack.push(sid);
 
-                        let rhs = self.parse_expr()?;
                         Statement::Declaration(sid, sp, rhs)
                     }
                     (first, _) => {
@@ -272,6 +356,31 @@ impl<T: Iterator<Item = Token>> Parser<T> {
                 Expression::ArrayAccess(sid, stackpos, index.into())
             }
 
+            T![ident(name)] if self.peek() == &T!['('] => {
+                // Skip the opening parenthesis
+                self.next();
+
+                let sid = self.string_store.intern_or_lookup(&name);
+
+                let mut args = Vec::new();
+
+                while !matches!(self.peek(), T![')']) {
+                    let arg = self.parse_expr()?;
+                    args.push(arg);
+
+                    // If there are more args skip the comma so that the loop will read the argname
+                    if self.peek() == &T![,] {
+                        self.next();
+                    }
+                }
+
+                validate_next!(self, T![')'], ")");
+
+                let fun_stackpos = self.get_fun_stackpos(sid)?;
+
+                Expression::FunCall(sid, fun_stackpos, args)
+            }
+
             T![ident(name)] => {
                 let sid = self.string_store.intern_or_lookup(&name);
                 let stackpos = self.get_stackpos(sid)?;
@@ -303,7 +412,7 @@ impl<T: Iterator<Item = Token>> Parser<T> {
             .iter()
             .rev()
             .position(|it| *it == varid)
-            .map(|it| self.var_stack.len() - it - 1)
+            .map(|it| it)
             .ok_or(ParseErr::UseOfUndeclaredVar(
                 self.string_store
                     .lookup(varid)
@@ -312,6 +421,20 @@ impl<T: Iterator<Item = Token>> Parser<T> {
             ))
     }
 
+    fn get_fun_stackpos(&self, varid: Sid) -> ResPE<usize> {
+        self.fun_stack
+            .iter()
+            .rev()
+            .position(|it| *it == varid)
+            .map(|it| self.fun_stack.len() - it - 1)
+            .ok_or(ParseErr::UseOfUndeclaredFun(
+                self.string_store
+                    .lookup(varid)
+                    .map(String::from)
+                    .unwrap_or("<unknown>".to_string()),
+            ))
+    }
+
     /// Get the next Token without removing it
     fn peek(&mut self) -> &Token {
         self.tokens.peek().unwrap_or(&T![EoF])

src/token.rs

@@ -1,4 +1,7 @@
-use crate::{ast::{BinOpType, UnOpType}, T};
+use crate::{
+    ast::{BinOpType, UnOpType},
+    T,
+};
 
 /// Language keywords
 #[derive(Debug, PartialEq, Eq)]
@@ -11,6 +14,10 @@ pub enum Keyword {
     If,
     /// Else keyword ("else")
     Else,
+    /// Function declaration keyword ("fun")
+    Fun,
+    /// Return keyword ("return")
+    Return,
 }
 
 /// Literal values
@@ -67,6 +74,9 @@ pub enum Token {
     /// Combined tokens consisting of multiple characters
     Combo(Combo),
 
+    /// Comma (",")
+    Comma,
+
     /// Equal Sign ("=")
     Equal,
 
@@ -199,6 +209,14 @@ macro_rules! T {
         crate::token::Token::Keyword(crate::token::Keyword::Else)
     };
 
+    [fun] => {
+        crate::token::Token::Keyword(crate::token::Keyword::Fun)
+    };
+
+    [return] => {
+        crate::token::Token::Keyword(crate::token::Keyword::Return)
+    };
+
     // Literals
     [i64($($val:tt)*)] => {
         crate::token::Token::Literal(crate::token::Literal::I64($($val)*))
@@ -251,6 +269,10 @@ macro_rules! T {
     };
 
     // Normal Tokens
+    [,] => {
+        crate::token::Token::Comma
+    };
+
     [=] => {
         crate::token::Token::Equal
     };