Fix test.lua

examples/test.lua

@@ -95,3 +95,12 @@ for k, v in (ipairs(table)) do
 end
 
 print("hello " .. "there")
+
+
+SETMETATABLE(table, {
+    __add = function ()
+        return 1
+    end
+})
+
+print(table + 5)

src/ast.rs

@@ -889,6 +889,7 @@ impl Parse for TableConstructorEntry {
                 (stream.peek(), stream.peek2())
             {
                 let key = stream.parse::<Node<String>>()?;
+                dbg!(&key);
                 let key = key.with(Expression::Literal(Literal::String(key.kind.clone())));
                 if let Some(Token::Symbol('=')) = stream.peek() {
                     stream.next();

src/vm/mod.rs

@@ -5,11 +5,44 @@ use std::{cell::RefCell, collections::HashMap, fmt::Debug, i64, rc::Rc};
 use crate::{
     CompilationUnit,
     ast::{BinaryOperator, UnaryOperator},
-    vm::value::{Constant, IndexableValue, LuaInteger, Table, Value},
+    vm::value::{Constant, IndexableValue, LuaBool, LuaInteger, RustFunction, Table, Value},
 };
 
 pub mod value;
 
+#[derive(Debug)]
+pub struct SetMetatable;
+impl RustFunction for SetMetatable {
+    fn execute(&self, parameters: Vec<Value>) -> Result<Vec<Value>, RuntimeError> {
+        let table = parameters
+            .get(0)
+            .ok_or(RuntimeError::NotTable(Value::Nil))?;
+        let metatable = parameters
+            .get(1)
+            .ok_or(RuntimeError::NotTable(Value::Nil))?;
+
+        match table {
+            Value::Table {
+                metatable: inner, ..
+            } => match metatable {
+                Value::Table {
+                    contents: metatable,
+                    ..
+                } => {
+                    *inner.borrow_mut() = metatable.borrow().clone();
+                    Ok(Vec::new())
+                }
+                _ => Err(RuntimeError::NotTable(metatable.clone())),
+            },
+            _ => Err(RuntimeError::NotTable(table.clone())),
+        }
+    }
+
+    fn as_indexable(&self) -> String {
+        "SETMETATABLE".to_owned()
+    }
+}
+
 #[derive(Clone, Copy)]
 pub enum Instruction {
     /// R(A) := R(B)
@@ -44,6 +77,8 @@ pub enum Instruction {
 
     /// R(A) := R(B) + R(C)
     Add(u16, u16, u16),
+    /// R(A) := R(B) - R(C)
+    Sub(u16, u16, u16),
     /// R(A) := R(B) * R(C)
     Mult(u16, u16, u16),
     /// R(A) := R(B) / R(C)
@@ -152,6 +187,7 @@ impl Debug for Instruction {
             }
 
             Instruction::Add(arg0, arg1, arg2) => write!(f, "ADD {} {} {}", arg0, arg1, arg2),
+            Instruction::Sub(arg0, arg1, arg2) => write!(f, "SUB {} {} {}", arg0, arg1, arg2),
             Instruction::Mult(arg0, arg1, arg2) => write!(f, "MUL {} {} {}", arg0, arg1, arg2),
             Instruction::Div(arg0, arg1, arg2) => write!(f, "DIV {} {} {}", arg0, arg1, arg2),
             Instruction::IDiv(arg0, arg1, arg2) => write!(f, "IDIV {} {} {}", arg0, arg1, arg2),
@@ -182,6 +218,16 @@ impl Debug for Instruction {
 pub enum RuntimeError {
     #[error("Unable to perform {0:?} operator between {1:?} and {2:?}")]
     InvalidOperands(BinaryOperator, Value, Value),
+    #[error("Unable to call metamethod {0} for values {1:?} and {2:?}")]
+    InvalidBinop(String, Value, Value),
+    #[error("Unable to call metamethod {0} for value {1:?}")]
+    InvalidUnop(String, Value),
+    #[error("Metatable is not a table: {0:?}")]
+    MetatableNotTable(Value),
+    #[error("Metafunction not found: {0}")]
+    MetafunctionNotFound(String),
+    #[error("Metafunction is not callable: {0:?}")]
+    MetafunctionNotCallable(Value),
     #[error("Unable to perform {0:?} operator to {1:?}")]
     InvalidOperand(UnaryOperator, Value),
     #[error("Tried calling a non-function: {0:?}")]
@@ -208,7 +254,7 @@ pub struct Prototype {
     pub parameters: usize,
 }
 
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone)]
 pub struct VirtualMachine {
     pub(super) environment: Table,
     pub(super) constants: Vec<Constant>,
@@ -216,6 +262,22 @@ pub struct VirtualMachine {
     pub(super) proto_counter: u32,
 }
 
+impl Default for VirtualMachine {
+    fn default() -> Self {
+        let environment: Table = Default::default();
+        environment.borrow_mut().insert(
+            IndexableValue::String("SETMETATABLE".into()),
+            SetMetatable.into(),
+        );
+        Self {
+            environment,
+            constants: Default::default(),
+            prototypes: Default::default(),
+            proto_counter: Default::default(),
+        }
+    }
+}
+
 impl VirtualMachine {
     pub fn new_prototype(&mut self, instructions: Prototype) -> u32 {
         let proto_id = self.proto_counter;
@@ -223,9 +285,7 @@ impl VirtualMachine {
         self.prototypes.insert(proto_id, instructions);
         proto_id
     }
-}
 
-impl VirtualMachine {
     pub(crate) fn create_closure(&self, prototype: u32) -> Closure {
         Closure {
             vm: self.clone(),
@@ -474,7 +534,10 @@ impl ClosureRunner {
                     {
                         self.set_stack(
                             *reg,
-                            StackValue::Value(Value::Table(self.closure.environment.clone())),
+                            StackValue::Value(Value::Table {
+                                contents: self.closure.environment.clone(),
+                                metatable: Rc::new(RefCell::new(HashMap::new())),
+                            }),
                         );
                     } else {
                         let environment = self.closure.environment.borrow_mut().clone();
@@ -507,7 +570,7 @@ impl ClosureRunner {
                     match table {
                         Some(value) => {
                             let mut table = value.borrow_mut();
-                            if let Value::Table(table) = &mut *table {
+                            if let Value::Table { contents, .. } = &mut *table {
                                 let index_value = self
                                     .stack
                                     .get(indexreg)
@@ -521,10 +584,10 @@ impl ClosureRunner {
                                     .unwrap_or(Value::Nil);
                                 match value {
                                     Value::Nil => {
-                                        table.borrow_mut().remove(&index_value);
+                                        contents.borrow_mut().remove(&index_value);
                                     }
                                     _ => {
-                                        table.borrow_mut().insert(index_value, value);
+                                        contents.borrow_mut().insert(index_value, value);
                                     }
                                 }
                             } else {
@@ -538,8 +601,8 @@ impl ClosureRunner {
                     let value = match self.stack.get(tablereg).cloned() {
                         Some(value) => {
                             let table = value.borrow();
-                            if let Value::Table(table) = &*table {
-                                let table = table.borrow();
+                            if let Value::Table { contents, .. } = &*table {
+                                let table = contents.borrow();
                                 let index_value = self
                                     .stack
                                     .get(indexreg)
@@ -567,12 +630,12 @@ impl ClosureRunner {
                     match table {
                         Some(value) => {
                             let mut table = value.borrow_mut();
-                            if let Value::Table(table) = &mut *table {
+                            if let Value::Table { contents, .. } = &mut *table {
                                 let mut counter = *start_idx as i64;
                                 for i in self.function_register..self.top {
                                     let value = self.get_stack(i + 1);
                                     match value {
-                                        StackValue::Value(value) => table.borrow_mut().insert(
+                                        StackValue::Value(value) => contents.borrow_mut().insert(
                                             IndexableValue::Integer(LuaInteger(counter)),
                                             value.clone(),
                                         ),
@@ -589,7 +652,10 @@ impl ClosureRunner {
                 Instruction::NewTable(reg) => {
                     self.set_stack(
                         *reg,
-                        StackValue::Value(Value::Table(Rc::new(RefCell::new(HashMap::new())))),
+                        StackValue::Value(Value::Table {
+                            contents: Rc::new(RefCell::new(HashMap::new())),
+                            metatable: Rc::new(RefCell::new(HashMap::new())),
+                        }),
                     );
                 }
                 Instruction::Jmp(b) => {
@@ -740,66 +806,213 @@ impl ClosureRunner {
 
                 Instruction::Concat(res, lhs, rhs) => {
                     let (lhs, rhs) = self.lhs_and_rhs(lhs, rhs);
-                    self.set_stack(*res, StackValue::Value(lhs.concat(&rhs)?));
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(self.result_or_metamethod_binop(
+                            lhs.concat(&rhs),
+                            "__concat",
+                            lhs,
+                            rhs,
+                        )?),
+                    );
                 }
 
                 Instruction::Add(res, lhs, rhs) => {
                     let (lhs, rhs) = self.lhs_and_rhs(lhs, rhs);
-                    self.set_stack(*res, StackValue::Value(lhs.add(&rhs)?));
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(self.result_or_metamethod_binop(
+                            lhs.add(&rhs),
+                            "__add",
+                            lhs,
+                            rhs,
+                        )?),
+                    );
+                }
+                Instruction::Sub(res, lhs, rhs) => {
+                    let (lhs, rhs) = self.lhs_and_rhs(lhs, rhs);
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(self.result_or_metamethod_binop(
+                            rhs.unm().and_then(|rhs| lhs.add(&rhs)),
+                            "__sub",
+                            lhs,
+                            rhs,
+                        )?),
+                    );
                 }
                 Instruction::Mult(res, lhs, rhs) => {
                     let (lhs, rhs) = self.lhs_and_rhs(lhs, rhs);
-                    self.set_stack(*res, StackValue::Value(lhs.mult(&rhs)?));
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(self.result_or_metamethod_binop(
+                            lhs.mult(&rhs),
+                            "__mult",
+                            lhs,
+                            rhs,
+                        )?),
+                    );
                 }
                 Instruction::Div(res, lhs, rhs) => {
                     let (lhs, rhs) = self.lhs_and_rhs(lhs, rhs);
-                    self.set_stack(*res, StackValue::Value(lhs.div(&rhs)?));
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(self.result_or_metamethod_binop(
+                            lhs.div(&rhs),
+                            "__div",
+                            lhs,
+                            rhs,
+                        )?),
+                    );
                 }
                 Instruction::IDiv(res, lhs, rhs) => {
                     let (lhs, rhs) = self.lhs_and_rhs(lhs, rhs);
-                    self.set_stack(*res, StackValue::Value(lhs.idiv(&rhs)?));
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(self.result_or_metamethod_binop(
+                            lhs.idiv(&rhs),
+                            "__idiv",
+                            lhs,
+                            rhs,
+                        )?),
+                    );
                 }
                 Instruction::Mod(res, lhs, rhs) => {
                     let (lhs, rhs) = self.lhs_and_rhs(lhs, rhs);
-                    self.set_stack(*res, StackValue::Value(lhs.r#mod(&rhs)?));
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(self.result_or_metamethod_binop(
+                            lhs.r#mod(&rhs),
+                            "__mod",
+                            lhs,
+                            rhs,
+                        )?),
+                    );
                 }
                 Instruction::Exp(res, lhs, rhs) => {
                     let (lhs, rhs) = self.lhs_and_rhs(lhs, rhs);
-                    self.set_stack(*res, StackValue::Value(lhs.exp(&rhs)?));
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(self.result_or_metamethod_binop(
+                            lhs.exp(&rhs),
+                            "__exp",
+                            lhs,
+                            rhs,
+                        )?),
+                    );
                 }
 
                 Instruction::BitAnd(res, lhs, rhs) => {
                     let (lhs, rhs) = self.lhs_and_rhs(lhs, rhs);
-                    self.set_stack(*res, StackValue::Value(lhs.band(&rhs)?));
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(self.result_or_metamethod_binop(
+                            lhs.band(&rhs),
+                            "__band",
+                            lhs,
+                            rhs,
+                        )?),
+                    );
                 }
                 Instruction::BitOr(res, lhs, rhs) => {
                     let (lhs, rhs) = self.lhs_and_rhs(lhs, rhs);
-                    self.set_stack(*res, StackValue::Value(lhs.bor(&rhs)?));
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(self.result_or_metamethod_binop(
+                            lhs.bor(&rhs),
+                            "__bor",
+                            lhs,
+                            rhs,
+                        )?),
+                    );
                 }
                 Instruction::BitXOr(res, lhs, rhs) => {
                     let (lhs, rhs) = self.lhs_and_rhs(lhs, rhs);
-                    self.set_stack(*res, StackValue::Value(lhs.bxor(&rhs)?));
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(self.result_or_metamethod_binop(
+                            lhs.bxor(&rhs),
+                            "__bxor",
+                            lhs,
+                            rhs,
+                        )?),
+                    );
                 }
                 Instruction::BitSRight(res, lhs, rhs) => {
                     let (lhs, rhs) = self.lhs_and_rhs(lhs, rhs);
-                    self.set_stack(*res, StackValue::Value(lhs.bsright(&rhs)?));
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(self.result_or_metamethod_binop(
+                            lhs.bsright(&rhs),
+                            "__shr",
+                            lhs,
+                            rhs,
+                        )?),
+                    );
                 }
                 Instruction::BitSLeft(res, lhs, rhs) => {
                     let (lhs, rhs) = self.lhs_and_rhs(lhs, rhs);
-                    self.set_stack(*res, StackValue::Value(lhs.bsleft(&rhs)?));
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(self.result_or_metamethod_binop(
+                            lhs.bsleft(&rhs),
+                            "__shl",
+                            lhs,
+                            rhs,
+                        )?),
+                    );
                 }
 
                 Instruction::Equal(res, lhs, rhs) => {
                     let (lhs, rhs) = self.lhs_and_rhs(lhs, rhs);
+                    match (&lhs, &rhs) {
+                        (Value::Table { metatable, .. }, Value::Table { .. }) => {
+                            self.set_stack(
+                                *res,
+                                StackValue::Value(Value::Boolean(LuaBool(
+                                    self.call_metamethod(
+                                        metatable,
+                                        "__eq",
+                                        vec![lhs.clone(), rhs.clone()],
+                                    )?
+                                    .is_truthy(),
+                                ))),
+                            );
+                        }
+                        _ => {
                             self.set_stack(*res, StackValue::Value(lhs.eq(&rhs)?));
                         }
+                    }
+                }
                 Instruction::LessThan(res, lhs, rhs) => {
                     let (lhs, rhs) = self.lhs_and_rhs(lhs, rhs);
-                    self.set_stack(*res, StackValue::Value(lhs.lt(&rhs)?));
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(Value::Boolean(LuaBool(
+                            self.result_or_metamethod_binop(
+                                lhs.lt(&rhs),
+                                "__lt",
+                                lhs.clone(),
+                                rhs.clone(),
+                            )?
+                            .is_truthy(),
+                        ))),
+                    );
                 }
                 Instruction::LessThanOrEqual(res, lhs, rhs) => {
                     let (lhs, rhs) = self.lhs_and_rhs(lhs, rhs);
-                    self.set_stack(*res, StackValue::Value(lhs.lte(&rhs)?));
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(Value::Boolean(LuaBool(
+                            self.result_or_metamethod_binop(
+                                lhs.lte(&rhs),
+                                "__le",
+                                lhs.clone(),
+                                rhs.clone(),
+                            )?
+                            .is_truthy(),
+                        ))),
+                    );
                 }
 
                 Instruction::Or(res, lhs, rhs) => {
@@ -812,28 +1025,41 @@ impl ClosureRunner {
                 }
 
                 Instruction::Unm(res, reg) => {
-                    self.set_stack(
-                        *res,
-                        StackValue::Value(
-                            self.stack
+                    let value = self
+                        .stack
                         .get(reg)
                         .map(|v| v.borrow().clone())
-                                .unwrap_or(Value::Nil)
-                                .unm()?,
-                        ),
+                        .unwrap_or(Value::Nil);
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(self.result_or_metamethod_unop(
+                            value.unm(),
+                            "__unm",
+                            value.clone(),
+                        )?),
                     );
                 }
                 Instruction::Len(res, reg) => {
-                    self.set_stack(
-                        *res,
-                        StackValue::Value(
-                            self.stack
+                    let value = self
+                        .stack
                         .get(reg)
                         .map(|v| v.borrow().clone())
-                                .unwrap_or(Value::Nil)
-                                .len()?,
-                        ),
-                    );
+                        .unwrap_or(Value::Nil);
+                    match &value {
+                        Value::Table { metatable, .. } => {
+                            match self.call_metamethod(&metatable, "__len", vec![value.clone()]) {
+                                Ok(result) => {
+                                    self.set_stack(*res, StackValue::Value(result));
+                                }
+                                Err(_) => {
+                                    self.set_stack(*res, StackValue::Value(value.len()?));
+                                }
+                            }
+                        }
+                        _ => {
+                            self.set_stack(*res, StackValue::Value(value.len()?));
+                        }
+                    }
                 }
                 Instruction::Not(res, reg) => {
                     self.set_stack(
@@ -903,6 +1129,79 @@ impl ClosureRunner {
         }
     }
 
+    fn result_or_metamethod_unop(
+        &self,
+        value: Result<Value, RuntimeError>,
+        metamethod: &str,
+        param: Value,
+    ) -> Result<Value, RuntimeError> {
+        match value {
+            Ok(value) => Ok(value),
+            Err(_) => match &param {
+                Value::Table { metatable, .. } => {
+                    self.call_metamethod(metatable, metamethod, vec![param.clone()])
+                }
+                _ => Err(RuntimeError::InvalidUnop(
+                    metamethod.to_owned(),
+                    param.clone(),
+                )),
+            },
+        }
+    }
+
+    fn result_or_metamethod_binop(
+        &self,
+        value: Result<Value, RuntimeError>,
+        metamethod: &str,
+        lhs: Value,
+        rhs: Value,
+    ) -> Result<Value, RuntimeError> {
+        match value {
+            Ok(value) => Ok(value),
+            Err(_) => match (&lhs, &rhs) {
+                (Value::Table { metatable, .. }, _) => {
+                    self.call_metamethod(metatable, metamethod, vec![lhs.clone(), rhs.clone()])
+                }
+                (_, Value::Table { metatable, .. }) => {
+                    self.call_metamethod(metatable, metamethod, vec![lhs.clone(), rhs.clone()])
+                }
+                _ => Err(RuntimeError::InvalidBinop(metamethod.to_owned(), lhs, rhs)),
+            },
+        }
+    }
+
+    fn call_metamethod(
+        &self,
+        metatable: &Table,
+        metamethod: &str,
+        params: Vec<Value>,
+    ) -> Result<Value, RuntimeError> {
+        if let Some(value) = metatable
+            .borrow()
+            .get(&IndexableValue::String(metamethod.to_owned()))
+        {
+            match value {
+                Value::RustFunction(function) => {
+                    let result = function.borrow_mut().execute(params)?;
+                    Ok(result.into_iter().next().unwrap())
+                }
+                Value::Function(closure) => {
+                    let mut runnable = closure.run(params);
+                    #[allow(unused_assignments)]
+                    let mut return_value = None;
+                    while {
+                        return_value = runnable.next()?;
+                        return_value.is_none()
+                    } {}
+                    Ok(return_value.unwrap().into_iter().next().unwrap())
+                }
+                _ => Err(RuntimeError::MetafunctionNotCallable(value.clone())),
+            }
+        } else {
+            Err(RuntimeError::MetafunctionNotFound(metamethod.to_owned()))
+        }
+    }
+
     fn lhs_and_rhs(&self, lhs: &u16, rhs: &u16) -> (Value, Value) {
         (
             self.stack

src/vm/value.rs

@@ -136,7 +136,7 @@ pub enum Value {
     RustFunction(Rc<RefCell<dyn RustFunction>>),
     Function(Closure),
     Nil,
-    Table(Table),
+    Table { contents: Table, metatable: Table },
 }
 
 impl Value {
@@ -151,7 +151,7 @@ impl Value {
             }
             Value::Function(closure) => Ok(IndexableValue::Function(closure.prototype)),
             Value::Nil => Err(RuntimeError::InvalidTableIndex(self)),
-            Value::Table(_) => Err(RuntimeError::InvalidTableIndex(self)),
+            Value::Table { .. } => Err(RuntimeError::InvalidTableIndex(self)),
         }
     }
 
@@ -201,7 +201,7 @@ impl Display for Value {
             }
             Value::Function(closure) => write!(f, "<function({})>", closure.prototype),
             Value::Nil => write!(f, "nil"),
-            Value::Table(_) => write!(f, "<table>"),
+            Value::Table { .. } => write!(f, "<table>"),
         }
     }
 }
@@ -514,8 +514,8 @@ impl Value {
     pub fn len(&self) -> Result<Value, RuntimeError> {
         match self {
             Value::String(value) => Ok(Self::Integer(LuaInteger(value.len() as i64))),
-            Value::Table(table) => {
-                let table = table.borrow();
+            Value::Table { contents, .. } => {
+                let table = contents.borrow();
                 let mut int_indexes = table
                     .keys()
                     .filter(|v| match v {
@@ -585,7 +585,7 @@ impl Value {
             Value::RustFunction(_) => true,
             Value::Function(_) => true,
             Value::Nil => false,
-            Value::Table(_) => true,
+            Value::Table { .. } => true,
         }
     }
 }
@@ -602,9 +602,9 @@ impl Debug for Value {
                 .debug_tuple(&format!("Function({})", closure.prototype))
                 .finish(),
             Value::Nil => write!(f, "Nil"),
-            Value::Table(hash_map) => {
+            Value::Table { contents, .. } => {
                 let mut table = f.debug_struct("Table");
-                for (key, value) in hash_map.borrow().iter() {
+                for (key, value) in contents.borrow().iter() {
                     table.field(&format!("{:?}", key), value);
                 }
                 table.finish()