Add table length

examples/test.lua

@@ -16,9 +16,15 @@ function min(x, y)
     return m
 end
 
+function f(x)
+    return x + 5
+end
+
 global sometable = {}
-sometable["hello"] = {}
+sometable["hello"] = { add(10)(15) }
+print(#sometable["hello"])
 sometable["hello"].there = "my dude"
+print(sometable.hello.there)
 
 print(max(11.12345, 9))
 print(add(10)(15))

src/ast.rs

@@ -272,7 +272,7 @@ impl Parse for Statement {
 
             Ok(Self::Assignment(
                 None,
-                vec![(access.0.0, access.0.1)],
+                vec![(access.0, access.1.0)],
                 ExpressionList(vec![expression]),
             ))
         } else if let Ok(expr) = stream.parse() {
@@ -290,12 +290,10 @@ pub enum AccessModifier {
 }
 
 #[derive(Debug, Clone)]
-pub struct IndexedAccess(Node<String>, Vec<Node<Expression>>);
+pub struct IndexedAccess(Vec<Node<Expression>>);
 
 impl Parse for IndexedAccess {
     fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
-        let name = stream.parse()?;
-
         let mut expressions = Vec::new();
         while let Some(Token::Symbol('[') | Token::Symbol('.')) = stream.peek() {
             match stream.next().unwrap() {
@@ -315,18 +313,19 @@ impl Parse for IndexedAccess {
             }
         }
 
-        Ok(IndexedAccess(name, expressions))
+        Ok(IndexedAccess(expressions))
     }
 }
 
 #[derive(Debug, Clone)]
-pub struct IndexedAssignment(IndexedAccess);
+pub struct IndexedAssignment(Node<String>, IndexedAccess);
 
 impl Parse for IndexedAssignment {
     fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
+        let name = stream.parse()?;
         let access = stream.parse()?;
         stream.expect_symbol('=')?;
-        Ok(IndexedAssignment(access))
+        Ok(IndexedAssignment(name, access))
     }
 }
 
@@ -338,7 +337,8 @@ pub enum Expression {
     FunctionDefinition(Vec<Node<String>>, Block),
     FunctionCall(Box<Node<Expression>>, Node<ExpressionList>),
     Literal(Literal),
-    TableConstructor,
+    TableConstructor(Vec<(Option<Node<Expression>>, Node<Expression>)>),
+    IndexedAccess(Box<Node<Expression>>, Box<Node<Expression>>),
 }
 
 impl Parse for Expression {
@@ -372,6 +372,7 @@ pub struct PrimaryExpression(Node<Expression>);
 #[derive(Debug, Clone, Copy)]
 pub enum UnaryOperator {
     Negation,
+    Length,
 }
 
 impl Parse for UnaryOperator {
@@ -379,6 +380,7 @@ impl Parse for UnaryOperator {
         if let Some(token) = stream.next() {
             match token {
                 Token::Symbol('-') => Ok(UnaryOperator::Negation),
+                Token::Symbol('#') => Ok(UnaryOperator::Length),
                 _ => Err(stream.expected_err("unop")),
             }
         } else {
@@ -495,15 +497,28 @@ impl Parse for PrimaryExpression {
             Expression::Literal(Literal::Nil)
         } else if let Some(Token::Symbol('{')) = peeked {
             stream.next();
+
+            let mut entries = Vec::new();
+            if let Ok(entry) = stream.parse::<TableConstructorEntry>() {
+                entries.push((entry.0, entry.1));
+                while let Some(Token::Symbol(',') | Token::Symbol(';')) = stream.peek() {
+                    stream.next();
+                    let entry = stream.parse::<TableConstructorEntry>()?;
+                    entries.push((entry.0, entry.1))
+                }
+            }
+
             stream.expect_symbol('}')?;
-            Expression::TableConstructor
+            Expression::TableConstructor(entries)
         } else {
             Expression::ValueRef(stream.parse()?)
         };
 
-        while let Some(Token::Symbol('(') | Token::Symbol('[')) = stream.peek() {
+        while let Some(Token::Symbol('(') | Token::Symbol('[') | Token::Symbol('.')) = stream.peek()
-            match stream.next().unwrap() {
+        {
+            match stream.peek().unwrap() {
                 Token::Symbol('(') => {
+                    stream.next();
                     let expression_list = stream.parse::<Node<ExpressionList>>()?;
                     stream.expect(Token::Symbol(')'))?;
                     expression = Expression::FunctionCall(
@@ -514,7 +529,19 @@ impl Parse for PrimaryExpression {
                         expression_list,
                     );
                 }
-                Token::Symbol('[') => todo!(),
+                Token::Symbol('[') | Token::Symbol('.') => {
+                    let meta = Metadata::produce(&mut stream, pre.clone());
+                    let IndexedAccess(accesses) = stream.parse()?;
+                    for access in accesses {
+                        expression = Expression::IndexedAccess(
+                            Box::new(Node {
+                                kind: expression,
+                                meta: meta.clone(),
+                            }),
+                            Box::new(access),
+                        )
+                    }
+                }
                 _ => panic!(),
             }
         }
@@ -573,6 +600,39 @@ fn parse_binop_rhs(
     Ok(lhs.0)
 }
 
+#[derive(Debug, Clone)]
+struct TableConstructorEntry(Option<Node<Expression>>, Node<Expression>);
+
+impl Parse for TableConstructorEntry {
+    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
+        if let Some(Token::Symbol('[')) = stream.peek() {
+            stream.next();
+            let key = stream.parse()?;
+            stream.expect_symbol(']')?;
+            stream.expect_symbol('=')?;
+            let value = stream.parse()?;
+            Ok(TableConstructorEntry(Some(key), value))
+        } else {
+            if let (Some(Token::Word(_)), Some(Token::Symbol('='))) =
+                (stream.peek(), stream.peek2())
+            {
+                let key = stream.parse::<Node<String>>()?;
+                let key = key.with(Expression::Literal(Literal::String(key.kind.clone())));
+                if let Some(Token::Symbol('=')) = stream.peek() {
+                    stream.next();
+                    let value = stream.parse()?;
+                    Ok(TableConstructorEntry(Some(key), value))
+                } else {
+                    Ok(TableConstructorEntry(None, key))
+                }
+            } else {
+                let value = stream.parse()?;
+                Ok(TableConstructorEntry(None, value))
+            }
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 pub enum Literal {
     Float(LuaFloat),

src/compile.rs

@@ -2,7 +2,7 @@ use std::collections::{HashMap, HashSet};
 
 use crate::{
     ast::{AccessModifier, BinaryOperator, Block, Expression, Literal, Statement, UnaryOperator},
-    vm::{Constant, Instruction, LuaBool},
+    vm::{Constant, Instruction, LuaBool, LuaInteger},
 };
 
 #[derive(Clone, Debug)]
@@ -401,8 +401,24 @@ impl Expression {
                     constants
                 }
             },
-            Expression::TableConstructor => {
+            Expression::TableConstructor(entries) => {
-                let constants = HashSet::new();
+                let mut constants = HashSet::new();
+                let mut counter = 1;
+                for (key, value) in entries {
+                    if let Some(key) = key {
+                        constants.extend(key.kind.find_constants(scope));
+                    } else {
+                        constants.insert(Constant::Integer(LuaInteger(counter)));
+                        counter += 1;
+                    }
+                    constants.extend(value.kind.find_constants(scope));
+                }
+                constants
+            }
+            Expression::IndexedAccess(expr, index) => {
+                let mut constants = HashSet::new();
+                constants.extend(expr.kind.find_constants(scope));
+                constants.extend(index.kind.find_constants(scope));
                 constants
             }
         }
@@ -509,6 +525,7 @@ impl Expression {
                 for reg in &registers {
                     match op {
                         UnaryOperator::Negation => instructions.push(Instruction::Unm(*reg, *reg)),
+                        UnaryOperator::Length => instructions.push(Instruction::Len(*reg, *reg)),
                     }
                 }
                 (instructions, registers)
@@ -624,12 +641,67 @@ impl Expression {
                 ));
                 (instructions, vec![reg])
             }
-            Expression::TableConstructor => {
+            Expression::TableConstructor(entries) => {
                 let mut instructions = Vec::new();
                 let reg = scope.register_counter.next();
                 instructions.push(Instruction::NewTable(reg));
+
+                let mut counter = 1;
+                for (key, value) in entries {
+                    if let Some(key) = key {
+                        let (instr, key_regs) = key.kind.compile(state, scope, Some(1));
+                        instructions.extend(instr);
+                        let (instr, value_regs) = value.kind.compile(state, scope, Some(1));
+                        instructions.extend(instr);
+                        instructions.push(Instruction::SetTable(
+                            reg,
+                            *key_regs.first().unwrap(),
+                            *value_regs.first().unwrap(),
+                        ));
+                    } else {
+                        let (instr, value_regs) = value.kind.compile(
+                            state,
+                            scope,
+                            if entries.len() == 1 { None } else { Some(1) },
+                        );
+                        instructions.extend(instr);
+
+                        if value_regs.len() > 0 {
+                            let key_reg = scope.register_counter.next();
+                            instructions.push(Instruction::LoadK(
+                                key_reg,
+                                state.get_constant(&Constant::Integer(LuaInteger(counter))),
+                            ));
+                            instructions.push(Instruction::SetTable(
+                                reg,
+                                key_reg,
+                                *value_regs.first().unwrap(),
+                            ));
+                            counter += 1;
+                        } else {
+                            instructions.push(Instruction::SetList(reg));
+                        }
+                    }
+                }
+
                 (instructions, vec![reg])
             }
+            Expression::IndexedAccess(expr, index) => {
+                let mut instructions = Vec::new();
+                let (instr, expr_regs) = expr.kind.compile(state, scope, Some(1));
+                instructions.extend(instr);
+                let (instr, index_regs) = index.kind.compile(state, scope, Some(1));
+                instructions.extend(instr);
+
+                let local = scope.register_counter.next();
+                instructions.push(Instruction::GetTable(
+                    local,
+                    *expr_regs.first().unwrap(),
+                    *index_regs.first().unwrap(),
+                ));
+
+                (instructions, vec![local])
+            }
         }
     }
 }

src/vm.rs

@@ -7,7 +7,7 @@ use crate::{
     ast::{BinaryOperator, UnaryOperator},
 };
 
-#[derive(Clone, Hash, PartialEq, Eq, Copy)]
+#[derive(Clone, Hash, PartialEq, Eq, Copy, PartialOrd, Ord)]
 pub struct VMFloat(u64);
 
 impl VMFloat {
@@ -37,7 +37,7 @@ impl Debug for LuaFloat {
     }
 }
 
-#[derive(Clone, Copy, Hash, PartialEq, Eq)]
+#[derive(Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord)]
 pub struct LuaInteger(pub i64);
 
 impl Debug for LuaInteger {
@@ -58,7 +58,7 @@ impl From<&LuaInteger> for LuaFloat {
     }
 }
 
-#[derive(Clone, Copy, Hash, PartialEq, Eq)]
+#[derive(Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord)]
 pub struct LuaBool(pub bool);
 
 impl Debug for LuaBool {
@@ -124,6 +124,8 @@ pub enum Instruction {
     SetUpVal(u16, u16),
     /// R(A)[R(B)] := R(C)
     SetTable(u16, u16, u16),
+    /// R(A) := all values returned from previous function
+    SetList(u16),
     /// R(A) := R(B)[R(C)]
     GetTable(u16, u16, u16),
     /// R(A) := {}
@@ -132,6 +134,8 @@ pub enum Instruction {
     Add(u16, u16, u16),
     /// R(A) := -R(B)
     Unm(u16, u16),
+    /// R(A) := #R(B) (length operator)
+    Len(u16, u16),
     /// R(A) := R(B) == R(C)
     Equal(u16, u16, u16),
     /// R(A) := R(B) < R(C)
@@ -172,6 +176,7 @@ impl Debug for Instruction {
             Instruction::GetTable(arg0, arg1, arg2) => {
                 write!(f, "GETTABLE {} {} {}", arg0, arg1, arg2)
             }
+            Instruction::SetList(arg0) => write!(f, "SETLIST {}", arg0),
             Instruction::NewTable(arg0) => write!(f, "NEWTABLE {}", arg0),
             Instruction::Jmp(arg0) => write!(f, "JMP {}", arg0),
             Instruction::Test(arg0, arg1, arg2) => write!(f, "TEST {} {} {}", arg0, arg1, arg2),
@@ -187,6 +192,7 @@ impl Debug for Instruction {
             Instruction::Add(arg0, arg1, arg2) => write!(f, "ADD {} {} {}", arg0, arg1, arg2),
             Instruction::LoadNil(arg0, arg1) => write!(f, "LOADNIL {} {}", arg0, arg1),
             Instruction::Unm(arg0, arg1) => write!(f, "UNM {} {}", arg0, arg1),
+            Instruction::Len(arg0, arg1) => write!(f, "LEN {} {}", arg0, arg1),
             Instruction::Or(arg0, arg1, arg2) => write!(f, "OR {} {} {}", arg0, arg1, arg2),
             Instruction::And(arg0, arg1, arg2) => write!(f, "AND {} {} {}", arg0, arg1, arg2),
         }
@@ -209,6 +215,8 @@ pub enum RuntimeError {
     NotTable(Value),
     #[error("Value is not coercable to a float: {0:?}")]
     NotFloatable(Value),
+    #[error("Value does not have a length: {0:?}")]
+    NotLengthable(Value),
     #[error("{0}")]
     Custom(String),
 }
@@ -259,7 +267,7 @@ impl Value {
     pub fn as_indexable(self) -> Result<IndexableValue, RuntimeError> {
         match self {
             Value::String(value) => Ok(IndexableValue::String(value)),
-            Value::Float(value) => Ok(IndexableValue::Number(value)),
+            Value::Float(value) => Ok(IndexableValue::Float(value)),
             Value::Integer(value) => Ok(IndexableValue::Integer(value)),
             Value::Boolean(value) => Ok(IndexableValue::Bool(value)),
             Value::RustFunction(value) => {
@@ -289,10 +297,10 @@ impl Value {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, Hash, Clone)]
+#[derive(Debug, PartialEq, Eq, Hash, Clone, Ord, PartialOrd)]
 pub enum IndexableValue {
     String(String),
-    Number(VMFloat),
+    Float(VMFloat),
     Integer(LuaInteger),
     Bool(LuaBool),
     RustFunction(String),
@@ -393,6 +401,47 @@ 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();
+                let mut int_indexes = table
+                    .keys()
+                    .filter(|v| match v {
+                        IndexableValue::Integer(_) => true,
+                        _ => false,
+                    })
+                    .map(|v| match v {
+                        IndexableValue::Integer(int) => int.0,
+                        _ => panic!(),
+                    })
+                    .collect::<Vec<_>>();
+                int_indexes.push(0);
+                int_indexes.sort();
+
+                for idx in int_indexes {
+                    let idx_value = table.get(&IndexableValue::Integer(LuaInteger(idx)));
+                    let idx_value_after = idx
+                        .checked_add_unsigned(1)
+                        .and_then(|i| table.get(&IndexableValue::Integer(LuaInteger(i))));
+
+                    // Directly from https://www.lua.org/manual/5.5/manual.html#3.4.7
+                    if (idx == 0 || idx_value.is_some()) && idx_value_after.is_none() {
+                        return Ok(Self::Integer(LuaInteger(idx)));
+                    }
+                }
+                panic!()
+            }
+            Value::Float(_) => Err(RuntimeError::NotLengthable(self.clone())),
+            Value::Integer(_) => Err(RuntimeError::NotLengthable(self.clone())),
+            Value::Boolean(_) => Err(RuntimeError::NotLengthable(self.clone())),
+            Value::RustFunction(_) => Err(RuntimeError::NotLengthable(self.clone())),
+            Value::Function(_) => Err(RuntimeError::NotLengthable(self.clone())),
+            Value::Nil => Err(RuntimeError::NotLengthable(self.clone())),
+        }
+    }
+
     pub fn and(&self, other: &Value) -> Result<Value, RuntimeError> {
         if self.is_truthy() {
             Ok(self.clone())
@@ -768,6 +817,30 @@ impl ClosureRunner {
 
                     self.set_stack(*res, value);
                 }
+                Instruction::SetList(reg) => {
+                    let table = self.stack.get(reg).cloned();
+                    match table {
+                        Some(value) => {
+                            let mut table = value.borrow_mut();
+                            if let Value::Table(table) = &mut *table {
+                                let mut counter = 1;
+                                for i in self.function_register..self.top {
+                                    let value = self.get_stack(i + 1);
+                                    match value {
+                                        StackValue::Value(value) => table.borrow_mut().insert(
+                                            IndexableValue::Integer(LuaInteger(counter)),
+                                            value.clone(),
+                                        ),
+                                    };
+                                    counter += 1;
+                                }
+                            } else {
+                                return Err(RuntimeError::NotTable(table.clone()));
+                            }
+                        }
+                        None => todo!(),
+                    }
+                }
                 Instruction::NewTable(reg) => {
                     self.set_stack(
                         *reg,
@@ -959,6 +1032,18 @@ impl ClosureRunner {
                         ),
                     );
                 }
+                Instruction::Len(res, reg) => {
+                    self.set_stack(
+                        *res,
+                        StackValue::Value(
+                            self.stack
+                                .get(reg)
+                                .map(|v| v.borrow().clone())
+                                .unwrap_or(Value::Nil)
+                                .len()?,
+                        ),
+                    );
+                }
                 Instruction::Or(res, lhs, rhs) => {
                     let lhs = self
                         .stack