Add functions

6 files changed, 432 insertions, 142 deletions

diff --git a/src/compiler.rs b/src/compiler.rs
index c59de4d..18599ae 100644
--- a/src/compiler.rs
+++ b/src/compiler.rs
@@ -47,7 +47,7 @@ nextable_enum!(Prec {
     Factor
 });
 
-#[derive(Copy, Clone)]
+#[derive(Debug, Copy, Clone)]
 enum Type {
     NoType,
     UnkownType,
@@ -74,7 +74,14 @@ impl TryFrom<&str> for Type {
 struct Variable {
     name: String,
     typ: Type,
-    level: usize,
+    scope: usize,
+    active: bool,
+}
+
+struct Frame {
+    stack: Vec<Variable>,
+    scope: usize,
+    variables_below: usize,
 }
 
 struct Compiler {
@@ -82,13 +89,48 @@ struct Compiler {
     tokens: TokenStream,
     current_file: PathBuf,
 
-    level: usize,
-    stack: Vec<Variable>,
+    frames: Vec<Frame>,
 
     panic: bool,
     errors: Vec<Error>,
 }
 
+macro_rules! push_frame {
+    ($compiler:expr, $block:expr, $code:tt) => {
+        {
+            $compiler.frames.push(Frame {
+                stack: Vec::new(),
+                scope: 0,
+                variables_below: $compiler.frame().variables_below + $compiler.stack().len(),
+            });
+
+            // Return value stored as a variable
+            $compiler.define_variable("", Type::UnkownType, &mut $block).unwrap();
+            $code
+
+            $compiler.frames.pop().unwrap();
+            // The 0th slot is the return value, which is passed out
+            // from functions, and should not be popped.
+            0
+        }
+    };
+}
+
+macro_rules! push_scope {
+    ($compiler:expr, $block:expr, $code:tt) => {
+        let ss = $compiler.stack().len();
+        $compiler.frame_mut().scope += 1;
+
+        $code;
+
+        $compiler.frame_mut().scope -= 1;
+        for _ in ss..$compiler.stack().len() {
+            $block.add(Op::Pop, $compiler.line());
+        }
+        $compiler.stack_mut().truncate(ss);
+    };
+}
+
 impl Compiler {
     pub fn new(current_file: &Path, tokens: TokenStream) -> Self {
         Self {
@@ -96,14 +138,35 @@ impl Compiler {
             tokens,
             current_file: PathBuf::from(current_file),
 
-            level: 0,
-            stack: vec![],
+            frames: vec![Frame {
+                stack: Vec::new(),
+                scope: 0,
+                variables_below: 0,
+            }],
 
             panic: false,
             errors: vec![],
         }
     }
 
+    fn frame(&self) -> &Frame {
+        let last = self.frames.len() - 1;
+        &self.frames[last]
+    }
+
+    fn frame_mut(&mut self) -> &mut Frame {
+        let last = self.frames.len() - 1;
+        &mut self.frames[last]
+    }
+
+    fn stack(&self) -> &[Variable] {
+        &self.frame().stack.as_ref()
+    }
+
+    fn stack_mut(&mut self) -> &mut Vec<Variable> {
+        &mut self.frame_mut().stack
+    }
+
     fn clear_panic(&mut self) {
         if self.panic {
             self.panic = false;
@@ -122,10 +185,10 @@ impl Compiler {
         if self.panic { return }
         self.panic = true;
         self.errors.push(Error {
-            kind: kind,
+            kind,
             file: self.current_file.clone(),
             line: self.line(),
-            message: message,
+            message,
         });
     }
 
@@ -216,6 +279,8 @@ impl Compiler {
                 | Token::NotEqual
                 => self.binary(block),
 
+            Token::LeftParen => self.call(block),
+
             _ => { return false; },
         }
         return true;
@@ -260,7 +325,7 @@ impl Compiler {
             Token::Minus => &[Op::Sub],
             Token::Star => &[Op::Mul],
             Token::Slash => &[Op::Div],
-            Token::AssertEqual => &[Op::AssertEqual],
+            Token::AssertEqual => &[Op::Equal, Op::Assert],
             Token::EqualEqual => &[Op::Equal],
             Token::Less => &[Op::Less],
             Token::Greater => &[Op::Greater],
@@ -273,7 +338,10 @@ impl Compiler {
     }
 
     fn expression(&mut self, block: &mut Block) {
-        self.parse_precedence(block, Prec::No);
+        match self.peek_four() {
+            (Token::Fn, ..) => self.function(block),
+            _ => self.parse_precedence(block, Prec::No),
+        }
     }
 
     fn parse_precedence(&mut self, block: &mut Block, precedence: Prec) {
@@ -288,11 +356,105 @@ impl Compiler {
         }
     }
 
-    fn find_local(&self, name: &str, block: &Block) -> Option<(usize, Type, usize)> {
-        self.stack.iter()
-                  .enumerate()
-                  .rev()
-                  .find_map(|x| if x.1.name == name { Some((x.0, x.1.typ, x.1.level)) } else { None} )
+    fn find_local(&self, name: &str, _block: &Block) -> Option<(usize, Type, usize)> {
+        let frame = self.frame();
+        for (slot, var) in frame.stack.iter().enumerate().rev() {
+            if var.name == name && var.active {
+                return Some((slot, var.typ, var.scope));
+            }
+        }
+        None
+    }
+
+    fn call(&mut self, block: &mut Block) {
+        expect!(self, Token::LeftParen, "Expected '(' at start of function call.");
+
+        let mut arity = 0;
+        loop {
+            match self.peek() {
+                Token::EOF => {
+                    error!(self, "Unexpected EOF in function call.");
+                    break;
+                }
+                Token::RightParen => {
+                    self.eat();
+                    break;
+                }
+                _ => {
+                    self.expression(block);
+                    arity += 1;
+                    if !matches!(self.peek(), Token::RightParen) {
+                        expect!(self, Token::Comma, "Expected ',' after argument.");
+                    }
+                }
+            }
+        }
+
+        block.add(Op::Call(arity), self.line());
+
+        for _ in 0..arity {
+            block.add(Op::Pop, self.line());
+        }
+    }
+
+    fn function(&mut self, block: &mut Block) {
+        expect!(self, Token::Fn, "Expected 'fn' at start of function.");
+
+        let name = if !self.stack()[self.stack().len() - 1].active {
+            &self.stack()[self.stack().len() - 1].name
+        } else {
+            "anonumus function"
+        };
+
+        let mut _return_type = None;
+        let mut function_block = Block::new(name, &self.current_file);
+        let arity;
+        let _ret = push_frame!(self, function_block, {
+            loop {
+                match self.peek() {
+                    Token::Identifier(name) => {
+                        self.eat();
+                        expect!(self, Token::Colon, "Expected ':' after parameter name.");
+                        if let Ok(typ) = self.type_ident() {
+                            if let Ok(slot) = self.define_variable(&name, typ, &mut function_block) {
+                                self.stack_mut()[slot].active = true;
+                            }
+                        } else {
+                            error!(self, "Failed to parse parameter type.");
+                        }
+                        if !matches!(self.peek(), Token::Arrow | Token::LeftBrace) {
+                            expect!(self, Token::Comma, "Expected ',' after parameter.");
+                        }
+                    }
+                    Token::LeftBrace => {
+                        _return_type = Some(Type::NoType);
+                        break;
+                    }
+                    Token::Arrow => {
+                        self.eat();
+                        if let Ok(typ) = self.type_ident() {
+                            _return_type = Some(typ);
+                        } else {
+                            error!(self, "Failed to parse return type.");
+                        }
+                        break;
+                    }
+                    _ => {
+                        error!(self, "Expected '->' or more paramters in function definition.");
+                        break;
+                    }
+                }
+            }
+            arity = self.frame().stack.len() - 1;
+
+            self.scope(&mut function_block);
+        });
+
+
+        function_block.add(Op::Constant(Value::Bool(true)), self.line());
+        function_block.add(Op::Return, self.line());
+
+        block.add(Op::Constant(Value::Function(arity, Rc::new(function_block))), self.line());
     }
 
     fn variable_expression(&mut self, block: &mut Block) {
@@ -307,17 +469,32 @@ impl Compiler {
         }
     }
 
-    fn define_variable(&mut self, name: &str, typ: Type, block: &mut Block) {
+    fn define_variable(&mut self, name: &str, typ: Type, block: &mut Block) -> Result<usize, ()> {
         if let Some((_, _, level)) = self.find_local(&name, block) {
-            if level == self.level {
+            if level == self.frame().scope {
                 error!(self, format!("Multiple definitions of {} in this block.", name));
-                return;
+                return Err(());
             }
         }
 
+        let slot = self.stack().len();
+        let scope = self.frame().scope;
+        self.stack_mut().push(Variable {
+            name: String::from(name),
+            typ,
+            scope,
+            active: false
+        });
+        Ok(slot)
+    }
+
+    fn definition_statement(&mut self, name: &str, typ: Type, block: &mut Block) {
+        let slot = self.define_variable(name, typ, block);
         self.expression(block);
 
-        self.stack.push(Variable { name: String::from(name), level: self.level, typ });
+        if let Ok(slot) = slot {
+            self.stack_mut()[slot].active = true;
+        }
     }
 
     fn assign(&mut self, name: &str, block: &mut Block) {
@@ -334,25 +511,16 @@ impl Compiler {
             return;
         }
 
-        self.level += 1;
-        let h = self.stack.len();
-
-        while !matches!(self.peek(), Token::RightBrace | Token::EOF) {
-            self.statement(block);
-            match self.peek() {
-                Token::Newline => { self.eat(); },
-                Token::RightBrace => { break; },
-                _ => { error!(self, "Expect newline after statement."); },
+        push_scope!(self, block, {
+            while !matches!(self.peek(), Token::RightBrace | Token::EOF) {
+                self.statement(block);
+                match self.peek() {
+                    Token::Newline => { self.eat(); },
+                    Token::RightBrace => { break; },
+                    _ => { error!(self, "Expect newline after statement."); },
+                }
             }
-        }
-
-        self.level -= 1;
-
-        for _ in h..self.stack.len() {
-            block.add(Op::Pop, self.line());
-        }
-
-        self.stack.truncate(h);
+        });
 
         expect!(self, Token::RightBrace, "Expected '}' at end of block.");
     }
@@ -384,66 +552,53 @@ impl Compiler {
     fn for_loop(&mut self, block: &mut Block) {
         expect!(self, Token::For, "Expected 'for' at start of for-loop.");
 
-        // push outer scope for loop variable
-        self.level += 1;
-        let h = self.stack.len();
-
-        // Definition
-        match self.peek_four() {
-            (Token::Identifier(name), Token::Identifier(typ), Token::ColonEqual, ..) => {
-                self.eat();
-                self.eat();
-                self.eat();
-                if let Ok(typ) = Type::try_from(typ.as_ref()) {
-                    self.define_variable(&name, typ, block);
-                } else {
-                    error!(self, format!("Failed to parse type '{}'.", typ));
+        push_scope!(self, block, {
+            // Definition
+            match self.peek_four() {
+                (Token::Identifier(name), Token::Identifier(typ), Token::ColonEqual, ..) => {
+                    self.eat();
+                    self.eat();
+                    self.eat();
+                    if let Ok(typ) = Type::try_from(typ.as_ref()) {
+                        self.definition_statement(&name, typ, block);
+                    } else {
+                        error!(self, format!("Failed to parse type '{}'.", typ));
+                    }
                 }
-            }
 
-            (Token::Identifier(name), Token::ColonEqual, ..) => {
-                self.eat();
-                self.eat();
-                self.define_variable(&name, Type::UnkownType, block);
-            }
+                (Token::Identifier(name), Token::ColonEqual, ..) => {
+                    self.eat();
+                    self.eat();
+                    self.definition_statement(&name, Type::UnkownType, block);
+                }
 
-            (Token::Comma, ..) => {}
+                (Token::Comma, ..) => {}
 
-            _ => { error!(self, "Expected definition at start of for-loop."); }
-        }
+                _ => { error!(self, "Expected definition at start of for-loop."); }
+            }
 
-        expect!(self, Token::Comma, "Expect ',' between initalizer and loop expression.");
+            expect!(self, Token::Comma, "Expect ',' between initalizer and loop expression.");
 
-        let cond = block.curr();
-        self.expression(block);
-        let cond_out = block.add(Op::Illegal, self.line());
-        let cond_cont = block.add(Op::Illegal, self.line());
-        expect!(self, Token::Comma, "Expect ',' between initalizer and loop expression.");
+            let cond = block.curr();
+            self.expression(block);
+            let cond_out = block.add(Op::Illegal, self.line());
+            let cond_cont = block.add(Op::Illegal, self.line());
+            expect!(self, Token::Comma, "Expect ',' between initalizer and loop expression.");
 
-        let inc = block.curr();
-        {
-            let h = self.stack.len();
-            self.statement(block);
-            for _ in h..self.stack.len() {
-                block.add(Op::Pop, self.line());
-            }
-            self.stack.truncate(h);
-        }
-        block.add(Op::Jmp(cond), self.line());
+            let inc = block.curr();
+            push_scope!(self, block, {
+                self.statement(block);
+            });
+            block.add(Op::Jmp(cond), self.line());
 
-        // patch_jmp!(Op::Jmp, cond_cont => block.curr());
-        block.patch(Op::Jmp(block.curr()), cond_cont);
-        self.scope(block);
-        block.add(Op::Jmp(inc), self.line());
+            // patch_jmp!(Op::Jmp, cond_cont => block.curr());
+            block.patch(Op::Jmp(block.curr()), cond_cont);
+            self.scope(block);
+            block.add(Op::Jmp(inc), self.line());
 
-        block.patch(Op::JmpFalse(block.curr()), cond_out);
+            block.patch(Op::JmpFalse(block.curr()), cond_out);
 
-        // pop outer scope
-        self.level -= 1;
-        for _ in h..self.stack.len() {
-            block.add(Op::Pop, self.line());
-        }
-        self.stack.truncate(h);
+        });
     }
 
     fn type_ident(&mut self) -> Result<Type, ()> {
@@ -476,7 +631,7 @@ impl Compiler {
                 self.eat();
                 if let Ok(typ) = self.type_ident() {
                     expect!(self, Token::Equal, "Expected assignment.");
-                    self.define_variable(&name, typ, block);
+                    self.definition_statement(&name, typ, block);
                 } else {
                     error!(self, format!("Expected type found '{:?}'.", self.peek()));
                 }
@@ -485,7 +640,7 @@ impl Compiler {
             tokens!(Token::Identifier(name), Token::ColonEqual) => {
                 self.eat();
                 self.eat();
-                self.define_variable(&name, Type::UnkownType, block);
+                self.definition_statement(&name, Type::UnkownType, block);
             }
 
             tokens!(Token::Identifier(name), Token::Equal) => {
@@ -502,6 +657,12 @@ impl Compiler {
                 self.for_loop(block);
             }
 
+            tokens!(Token::Ret) => {
+                self.eat();
+                self.expression(block);
+                block.add(Op::Return, self.line());
+            }
+
             tokens!(Token::Unreachable) => {
                 self.eat();
                 block.add(Op::Unreachable, self.line());
diff --git a/src/error.rs b/src/error.rs
index abf8810..979ca07 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -6,7 +6,7 @@ use crate::tokenizer::Token;
 #[derive(Debug, Clone)]
 pub enum ErrorKind {
     TypeError(Op, Vec<Value>),
-    AssertFailed(Value, Value),
+    Assert,
     InvalidProgram,
     Unreachable,
 
@@ -30,8 +30,8 @@ impl fmt::Display for ErrorKind {
                     .fold(String::new(), |a, v| { format!("{}, {:?}", a, v) });
                 write!(f, "Cannot apply {:?} to values {}", op, values)
             }
-            ErrorKind::AssertFailed(a, b) => {
-                write!(f, "Assertion failed, {:?} != {:?}.", a, b)
+            ErrorKind::Assert => {
+                write!(f, "Assertion failed.")
             }
             ErrorKind::SyntaxError(line, token) => {
                 write!(f, "Syntax error on line {} at token {:?}", line, token)
diff --git a/src/lib.rs b/src/lib.rs
index 5bef0ae..fd9c094 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,4 +1,5 @@
 use std::path::Path;
+use std::rc::Rc;
 
 pub mod compiler;
 pub mod tokenizer;
@@ -19,7 +20,7 @@ pub fn run_string(s: &str) -> Result<(), Vec<Error>> {
 
 pub fn run(tokens: TokenStream, path: &Path) -> Result<(), Vec<Error>> {
     match compiler::compile("main", path, tokens) {
-        Ok(block) => vm::run_block(&block).or_else(|e| Err(vec![e])),
+        Ok(block) => vm::run_block(Rc::new(block)).or_else(|e| Err(vec![e])),
         Err(errors) => Err(errors),
     }
 }
@@ -60,9 +61,11 @@ mod tests {
             }
         };
     }
+
     test_file!(order_of_operations, "tests/order-of-operations.tdy");
     test_file!(variables, "tests/variables.tdy");
     test_file!(scoping, "tests/scoping.tdy");
     test_file!(if_, "tests/if.tdy");
     test_file!(for_, "tests/for.tdy");
+    test_file!(fun, "tests/fun.tdy");
 }
diff --git a/src/tokenizer.rs b/src/tokenizer.rs
index 2e28cd2..54712d0 100644
--- a/src/tokenizer.rs
+++ b/src/tokenizer.rs
@@ -33,6 +33,9 @@ pub enum Token {
     #[token("print")]
     Print,
 
+    #[token("ret")]
+    Ret,
+
     #[token("+")]
     Plus,
     #[token("++")]
@@ -96,8 +99,8 @@ pub enum Token {
     #[token("<=")]
     LessEqual,
 
-    #[token("|")]
-    Pipe,
+    #[token("fn")]
+    Fn,
 
     #[token("&&")]
     And,
diff --git a/src/vm.rs b/src/vm.rs
index 44a5829..46b4105 100644
--- a/src/vm.rs
+++ b/src/vm.rs
@@ -1,3 +1,4 @@
+use std::fmt::Debug;
 use std::path::{Path, PathBuf};
 use std::rc::Rc;
 use std::collections::HashMap;
@@ -14,12 +15,25 @@ macro_rules! error {
     };
 }
 
-#[derive(Debug, Clone, PartialEq, PartialOrd)]
+#[derive(Clone)]
 pub enum Value {
     Float(f64),
     Int(i64),
     Bool(bool),
     String(Rc<String>),
+    Function(usize, Rc<Block>),
+}
+
+impl Debug for Value {
+    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Value::Float(f) => write!(fmt, "(float {})", f),
+            Value::Int(i) => write!(fmt, "(int {})", i),
+            Value::Bool(b) => write!(fmt, "(bool {})", b),
+            Value::String(s) => write!(fmt, "(string \"{}\")", s),
+            Value::Function(arity, block) => write!(fmt, "(func {}-{})", block.name, arity),
+        }
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -46,12 +60,13 @@ pub enum Op {
     Less,    // <
     Greater, // >
 
-    AssertEqual,
+    Assert,
     Unreachable,
 
     ReadLocal(usize),
     Assign(usize),
 
+    Call(usize),
 
     Print,
     Return,
@@ -130,25 +145,32 @@ impl Block {
 }
 
 #[derive(Debug)]
-pub struct VM<'a> {
-    stack: Vec<Value>,
-
-    block: &'a Block,
+struct Frame {
+    stack_offset: usize,
+    block: Rc<Block>,
     ip: usize,
 }
 
-pub fn run_block<'a> (block: &'a Block) -> Result<(), Error> {
+#[derive(Debug)]
+pub struct VM {
+    stack: Vec<Value>,
+    frames: Vec<Frame>,
+}
+
+pub fn run_block(block: Rc<Block>) -> Result<(), Error> {
     let mut vm = VM {
         stack: Vec::new(),
-
-        block,
-        ip: 0,
+        frames: vec![Frame {
+            stack_offset: 0,
+            block,
+            ip: 0
+        }],
     };
 
     vm.run()
 }
 
-impl VM<'_> {
+impl VM {
     fn pop_twice(&mut self) -> (Value, Value) {
         let len = self.stack.len();
         let res = (self.stack[len-2].clone(), self.stack[len-1].clone());
@@ -160,41 +182,54 @@ impl VM<'_> {
         self.stack.get(self.stack.len() - amount)
     }
 
+    fn frame(&self) -> &Frame {
+        let last = self.frames.len() - 1;
+        &self.frames[last]
+    }
+
+    fn frame_mut(&mut self) -> &mut Frame {
+        let last = self.frames.len() - 1;
+        &mut self.frames[last]
+    }
+
     fn error(&self, kind: ErrorKind, message: Option<String>) -> Error {
+        let frame = self.frames.last().unwrap();
         Error {
             kind,
-            file: self.block.file.clone(),
-            line: self.block.line(self.ip),
+            file: frame.block.file.clone(),
+            line: frame.block.line(frame.ip),
             message,
         }
     }
 
     pub fn run(&mut self) -> Result<(), Error>{
+        //TODO better system so tests dont print
         const PRINT_WHILE_RUNNING: bool = true;
         const PRINT_BLOCK: bool = true;
 
         if PRINT_BLOCK {
-            self.block.debug_print();
+            self.frame().block.debug_print();
         }
 
         loop {
             if PRINT_WHILE_RUNNING {
-                print!("    [");
-                for (i, s) in self.stack.iter().enumerate() {
+                let start = self.frame().stack_offset;
+                print!("    {:3} [", start);
+                for (i, s) in self.stack.iter().skip(start).enumerate() {
                     if i != 0 {
                         print!(" ");
                     }
-                    match s {
-                        Value::String(rc) => print!("{:?}<{}>", rc.green(), Rc::strong_count(rc)),
-                        _ => print!("{:?}", s.green()),
-                    }
+                    print!("{:?}", s.green());
                 }
                 println!("]");
 
-                println!("{:5} {:05} {:?}", self.block.line(self.ip).red(), self.ip.blue(), self.block.ops[self.ip]);
+                println!("{:5} {:05} {:?}",
+                         self.frame().block.line(self.frame().ip).red(),
+                         self.frame().ip.blue(),
+                         self.frame().block.ops[self.frame().ip]);
             }
 
-            let op = self.block.ops[self.ip].clone();
+            let op = self.frame().block.ops[self.frame().ip].clone();
             match op {
                 Op::Illegal => {
                     error!(self, ErrorKind::InvalidProgram);
@@ -222,8 +257,8 @@ impl VM<'_> {
 
                 Op::Add => {
                     match self.pop_twice() {
-                        (Value::Float(a), Value::Float(b)) => self.stack.push(Value::Float(b + a)),
-                        (Value::Int(a), Value::Int(b)) => self.stack.push(Value::Int(b + a)),
+                        (Value::Float(a), Value::Float(b)) => self.stack.push(Value::Float(a + b)),
+                        (Value::Int(a), Value::Int(b)) => self.stack.push(Value::Int(a + b)),
                         (Value::String(a), Value::String(b)) => {
                             self.stack.push(Value::String(Rc::from(format!("{}{}", a, b))))
                         }
@@ -233,41 +268,56 @@ impl VM<'_> {
 
                 Op::Sub => {
                     match self.pop_twice() {
-                        (Value::Float(a), Value::Float(b)) => self.stack.push(Value::Float(b - a)),
-                        (Value::Int(a), Value::Int(b)) => self.stack.push(Value::Int(b - a)),
+                        (Value::Float(a), Value::Float(b)) => self.stack.push(Value::Float(a - b)),
+                        (Value::Int(a), Value::Int(b)) => self.stack.push(Value::Int(a - b)),
                         (a, b) => error!(self, ErrorKind::TypeError(op.clone(), vec![a, b])),
                     }
                 }
 
                 Op::Mul => {
                     match self.pop_twice() {
-                        (Value::Float(a), Value::Float(b)) => self.stack.push(Value::Float(b * a)),
-                        (Value::Int(a), Value::Int(b)) => self.stack.push(Value::Int(b * a)),
+                        (Value::Float(a), Value::Float(b)) => self.stack.push(Value::Float(a * b)),
+                        (Value::Int(a), Value::Int(b)) => self.stack.push(Value::Int(a * b)),
                         (a, b) => error!(self, ErrorKind::TypeError(op.clone(), vec![a, b])),
                     }
                 }
 
                 Op::Div => {
                     match self.pop_twice() {
-                        (Value::Float(a), Value::Float(b)) => self.stack.push(Value::Float(b / a)),
-                        (Value::Int(a), Value::Int(b)) => self.stack.push(Value::Int(b / a)),
+                        (Value::Float(a), Value::Float(b)) => self.stack.push(Value::Float(a / b)),
+                        (Value::Int(a), Value::Int(b)) => self.stack.push(Value::Int(a / b)),
                         (a, b) => error!(self, ErrorKind::TypeError(op.clone(), vec![a, b])),
                     }
                 }
 
                 Op::Equal => {
-                    let (a, b) = self.pop_twice();
-                    self.stack.push(Value::Bool(a == b));
+                    match self.pop_twice() {
+                        (Value::Float(a), Value::Float(b)) => self.stack.push(Value::Bool(a == b)),
+                        (Value::Int(a), Value::Int(b)) => self.stack.push(Value::Bool(a == b)),
+                        (Value::String(a), Value::String(b)) => self.stack.push(Value::Bool(a == b)),
+                        (Value::Bool(a), Value::Bool(b)) => self.stack.push(Value::Bool(a == b)),
+                        (a, b) => error!(self, ErrorKind::TypeError(op.clone(), vec![a, b])),
+                    }
                 }
 
                 Op::Less => {
-                    let (a, b) = self.pop_twice();
-                    self.stack.push(Value::Bool(a < b));
+                    match self.pop_twice() {
+                        (Value::Float(a), Value::Float(b)) => self.stack.push(Value::Bool(a < b)),
+                        (Value::Int(a), Value::Int(b)) => self.stack.push(Value::Bool(a < b)),
+                        (Value::String(a), Value::String(b)) => self.stack.push(Value::Bool(a < b)),
+                        (Value::Bool(a), Value::Bool(b)) => self.stack.push(Value::Bool(a < b)),
+                        (a, b) => error!(self, ErrorKind::TypeError(op.clone(), vec![a, b])),
+                    }
                 }
 
                 Op::Greater => {
-                    let (a, b) = self.pop_twice();
-                    self.stack.push(Value::Bool(a > b));
+                    match self.pop_twice() {
+                        (Value::Float(a), Value::Float(b)) => self.stack.push(Value::Bool(a > b)),
+                        (Value::Int(a), Value::Int(b)) => self.stack.push(Value::Bool(a > b)),
+                        (Value::String(a), Value::String(b)) => self.stack.push(Value::Bool(a > b)),
+                        (Value::Bool(a), Value::Bool(b)) => self.stack.push(Value::Bool(a > b)),
+                        (a, b) => error!(self, ErrorKind::TypeError(op.clone(), vec![a, b])),
+                    }
                 }
 
                 Op::And => {
@@ -292,42 +342,74 @@ impl VM<'_> {
                 }
 
                 Op::Jmp(line) => {
-                    self.ip = line;
+                    self.frame_mut().ip = line;
                     continue;
                 }
 
                 Op::JmpFalse(line) => {
-                    if Some(Value::Bool(false)) == self.stack.pop() {
-                        self.ip = line;
+                    if matches!(self.stack.pop(), Some(Value::Bool(false))) {
+                        self.frame_mut().ip = line;
                         continue;
                     }
                 }
 
-                Op::AssertEqual => {
-                    let (a, b) = self.pop_twice();
-                    if a != b {
-                        error!(self, ErrorKind::AssertFailed(a, b));
+                Op::Assert => {
+                    if matches!(self.stack.pop(), Some(Value::Bool(false))) {
+                        error!(self, ErrorKind::Assert);
                     }
-                    self.stack.push(Value::Bool(a == b));
+                    self.stack.push(Value::Bool(true));
                 }
 
                 Op::ReadLocal(slot) => {
+                    let slot = self.frame().stack_offset + slot;
                     self.stack.push(self.stack[slot].clone());
                 }
 
                 Op::Assign(slot) => {
+                    let slot = self.frame().stack_offset + slot;
                     self.stack[slot] = self.stack.pop().unwrap();
                 }
 
+                Op::Call(num_args) => {
+                    let new_base = self.stack.len() - 1 - num_args;
+                    match &self.stack[new_base] {
+                        Value::Function(arity, block) => {
+                            if arity != &num_args {
+                                error!(self,
+                                       ErrorKind::InvalidProgram,
+                                       format!("Invalid number of arguments, got {} expected {}.",
+                                               num_args, arity));
+                            }
+                            if PRINT_BLOCK {
+                                block.debug_print();
+                            }
+                            self.frames.push(Frame {
+                                stack_offset: new_base,
+                                block: Rc::clone(&block),
+                                ip: 0,
+                            });
+                            continue;
+                        },
+                        _ => {
+                            unreachable!()
+                        }
+                    }
+                }
+
                 Op::Print => {
-                    println!("PRINT: {:?}", self.stack.pop());
+                    println!("PRINT: {:?}", self.stack.pop().unwrap());
                 }
 
                 Op::Return => {
-                    return Ok(());
+                    let last = self.frames.pop().unwrap();
+                    if self.frames.is_empty() {
+                        return Ok(());
+                    } else {
+                        self.stack[last.stack_offset] = self.stack.pop().unwrap();
+                    }
                 }
             }
-            self.ip += 1;
+            self.frame_mut().ip += 1;
         }
     }
 }
diff --git a/tests/fun.tdy b/tests/fun.tdy
new file mode 100644
index 0000000..8b94cad
--- /dev/null
+++ b/tests/fun.tdy
@@ -0,0 +1,41 @@
+// Simplest
+f := fn {
+    print 1
+}
+f() <=> true
+
+// Simple
+f2 := fn a: int {
+    print a
+}
+f2(2) <=> true
+
+// Return value
+f3 := fn -> int {
+    ret 3
+}
+print f3()
+f3() <=> 3
+
+// Empty function
+f4 := fn {}
+print f4
+print f4()
+
+// Multiple arguments
+adder := fn a: int, b: int -> int {
+    ret a + b
+}
+adder(1, 2) <=> 3
+
+// Passing functions
+h := fn {
+    print "h"
+    ret 1
+}
+
+g := fn f: int {
+    ret f()
+}
+
+g(h) <=> 1