Add varargs

examples/test.lua

@@ -1,36 +1,7 @@
-global b = 5
 
-function add(x)
+function f(x, ...)
-    return function (y)
+    local b, c = ...
-        x = x + 1
+    return x + 5, b, c
-        b = b + 1
-        return x + y, 1, 2, b
-    end
 end
 
-function min(x, y)
+print(123, f(10, 11, 12))
-    local m = x
-    if y < x then
-        m = y
-    end
-    return m
-end
-
-function f(x)
-    return x + 5
-end
-
-global sometable = {}
-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))
-print(add(10)(15))
-print(b)
-print(min(11, 9))
-print(10 - 15)
-print("hello there!")
-print(true or 0)

src/ast.rs

@@ -129,10 +129,39 @@ impl Metadata {
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct Ellipsis;
+
+impl Parse for Ellipsis {
+    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
+        stream.expect_symbol('.')?;
+        stream.expect_symbol('.')?;
+        stream.expect_symbol('.')?;
+        Ok(Ellipsis)
+    }
+}
+
+#[derive(Debug, Clone)]
+pub enum IdentOrEllipsis {
+    Ident(String),
+    Ellipsis,
+}
+
+impl Parse for IdentOrEllipsis {
+    fn parse(mut stream: TokenStream) -> Result<Self, TokenStreamError> {
+        if let Ok(ident) = stream.parse() {
+            Ok(IdentOrEllipsis::Ident(ident))
+        } else {
+            stream.parse::<Ellipsis>()?;
+            Ok(IdentOrEllipsis::Ellipsis)
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct Function {
     pub name: Option<Node<String>>,
-    pub params: Vec<Node<String>>,
+    pub params: Vec<Node<IdentOrEllipsis>>,
     pub block: Block,
     pub _meta: Metadata,
 }
@@ -147,11 +176,23 @@ impl Parse for Function {
 
         let mut params = Vec::new();
 
-        if let Ok(param) = stream.parse() {
+        if let Ok(param) = stream.parse::<Node<IdentOrEllipsis>>() {
-            params.push(param);
+            params.push(param.clone());
-            while stream.peek() == Some(Token::Symbol(',')) {
+            match param.kind {
-                stream.next();
+                IdentOrEllipsis::Ident(_) => {
-                params.push(stream.parse()?);
+                    while stream.peek() == Some(Token::Symbol(',')) {
+                        stream.next();
+                        let param = stream.parse::<Node<IdentOrEllipsis>>()?;
+                        params.push(param.clone());
+                        match param.kind {
+                            IdentOrEllipsis::Ident(_) => {}
+                            IdentOrEllipsis::Ellipsis => {
+                                break;
+                            }
+                        }
+                    }
+                }
+                IdentOrEllipsis::Ellipsis => {}
             }
         }
 
@@ -334,11 +375,12 @@ pub enum Expression {
     ValueRef(String),
     UnOp(UnaryOperator, Box<Node<Expression>>),
     BinOp(BinaryOperator, Box<Node<Expression>>, Box<Node<Expression>>),
-    FunctionDefinition(Vec<Node<String>>, Block),
+    FunctionDefinition(Vec<Node<IdentOrEllipsis>>, Block),
     FunctionCall(Box<Node<Expression>>, Node<ExpressionList>),
     Literal(Literal),
     TableConstructor(Vec<(Option<Node<Expression>>, Node<Expression>)>),
     IndexedAccess(Box<Node<Expression>>, Box<Node<Expression>>),
+    Ellipsis,
 }
 
 impl Parse for Expression {
@@ -510,6 +552,8 @@ impl Parse for PrimaryExpression {
 
             stream.expect_symbol('}')?;
             Expression::TableConstructor(entries)
+        } else if let Ok(_) = stream.parse::<Ellipsis>() {
+            Expression::Ellipsis
         } else {
             Expression::ValueRef(stream.parse()?)
         };

src/compile.rs

@@ -1,14 +1,17 @@
 use std::collections::{HashMap, HashSet};
 
 use crate::{
-    ast::{AccessModifier, BinaryOperator, Block, Expression, Literal, Statement, UnaryOperator},
+    ast::{
-    vm::{Constant, Instruction, LuaBool, LuaInteger},
+        AccessModifier, BinaryOperator, Block, Expression, IdentOrEllipsis, Literal, Statement,
+        UnaryOperator,
+    },
+    vm::{Constant, Instruction, LuaBool, LuaInteger, Prototype},
 };
 
 #[derive(Clone, Debug)]
 pub struct State {
     pub constants: Vec<Constant>,
-    pub prototypes: Vec<Vec<Instruction>>,
+    pub prototypes: Vec<Prototype>,
 }
 
 impl State {
@@ -28,6 +31,7 @@ pub struct Scope {
     pub register_counter: LocalCounter,
     pub highest_upvalue: u16,
     pub upvalues: HashMap<String, u16>,
+    pub is_vararg: bool,
 }
 
 #[derive(Clone, Debug, Default)]
@@ -43,22 +47,6 @@ impl LocalCounter {
     }
 
     pub fn consecutive(&mut self, amount: usize) -> Vec<u16> {
-        'outer: for free_num in self.1.clone() {
-            let mut potentials = vec![free_num];
-            let mut curr = free_num;
-            for _ in 0..amount {
-                if let Some(next) = self.1.iter().find(|v| **v == curr + 1) {
-                    potentials.push(*next);
-                    curr = *next;
-                } else {
-                    continue 'outer;
-                }
-            }
-
-            self.1
-                .retain_mut(|v| potentials.iter().find(|p| v != *p).is_none());
-            return potentials;
-        }
         let mut returned = Vec::new();
         for _ in 0..amount {
             returned.push(self.new());
@@ -190,11 +178,19 @@ impl Statement {
 
                 match access_modifier {
                     Some(AccessModifier::Local) => {
+                        let mut vararg_applied = false;
                         for (i, (name, indexes)) in names.iter().enumerate() {
-                            instructions.push(Instruction::Move(
+                            if let Some(expr_reg) = expr_regs.get(i) {
-                                *new_registers.get(i).unwrap(),
+                                instructions.push(Instruction::Move(
-                                expr_regs.get(i).cloned().unwrap(),
+                                    *new_registers.get(i).unwrap(),
-                            ));
+                                    *expr_reg,
+                                ));
+                            } else if !vararg_applied {
+                                instructions.push(Instruction::MoveRetValues(
+                                    *new_registers.get(i).unwrap(),
+                                ));
+                                vararg_applied = true;
+                            }
                             if indexes.len() > 0 {
                                 todo!()
                             }
@@ -283,17 +279,21 @@ impl Statement {
             }
             Statement::Return(expr_list) => {
                 let mut ret_registers = Vec::new();
-                for expr in &expr_list.0 {
+                let mut vararg = false;
+                for (i, expr) in expr_list.0.iter().enumerate() {
                     let (instr, registers) = expr.kind.compile(
                         state,
                         scope,
-                        if expr_list.0.len() == 1 {
+                        if i == expr_list.0.len() - 1 {
                             None
                         } else {
                             Some(1)
                         },
                     );
                     instructions.extend(instr);
+                    if registers.len() == 0 {
+                        vararg = true;
+                    }
                     ret_registers.extend(registers);
                 }
 
@@ -310,9 +310,21 @@ impl Statement {
                     *ret_register = new_reg;
                 }
 
+                if vararg {
+                    instructions.push(Instruction::MoveRetValues(
+                        first_ret_register + ret_registers.len() as u16,
+                    ));
+                }
+
+                dbg!(&first_ret_register);
+
                 instructions.push(Instruction::Return(
-                    *ret_registers.first().unwrap_or(&scope.register_counter.0),
+                    first_ret_register,
-                    *ret_registers.last().unwrap_or(&0),
+                    if vararg {
+                        0
+                    } else {
+                        *ret_registers.last().unwrap_or(&0)
+                    },
                 ));
             }
             Statement::If(expr, block) => {
@@ -421,6 +433,7 @@ impl Expression {
                 constants.extend(index.kind.find_constants(scope));
                 constants
             }
+            Expression::Ellipsis => HashSet::new(),
         }
     }
 
@@ -533,9 +546,16 @@ impl Expression {
             Expression::FunctionDefinition(params, block) => {
                 let mut inner_scope = Scope::default();
                 for param in params {
-                    inner_scope
+                    match &param.kind {
-                        .locals
+                        IdentOrEllipsis::Ident(name) => {
-                        .insert(param.kind.clone(), inner_scope.register_counter.next());
+                            inner_scope
+                                .locals
+                                .insert(name.clone(), inner_scope.register_counter.next());
+                        }
+                        IdentOrEllipsis::Ellipsis => {
+                            inner_scope.is_vararg = true;
+                        }
+                    }
                 }
 
                 inner_scope.highest_upvalue = scope.highest_upvalue + scope.register_counter.0;
@@ -547,7 +567,14 @@ impl Expression {
                 }
 
                 let instructions = block.compile(state, &mut inner_scope);
-                state.prototypes.push(instructions);
+                state.prototypes.push(Prototype {
+                    instructions,
+                    parameters: if inner_scope.is_vararg {
+                        params.len() - 1
+                    } else {
+                        params.len()
+                    },
+                });
 
                 let mut instructions = Vec::new();
                 instructions.push(Instruction::Close(scope.register_counter.0));
@@ -564,26 +591,32 @@ impl Expression {
 
                 let old_function_reg = registers.first().unwrap();
 
+                let mut registers = scope
+                    .register_counter
+                    .consecutive(params.kind.0.len() + 1)
+                    .into_iter();
+
                 let mut param_scope = scope.clone();
                 let mut original_param_regs = Vec::new();
-                for param in params.kind.0.iter() {
+                let mut vararg = false;
+                for (i, param) in params.kind.0.iter().enumerate() {
                     let (instr, registers) = param.kind.compile(
                         state,
                         &mut param_scope,
-                        if params.kind.0.len() == 1 {
+                        if i == params.kind.0.len() - 1 {
                             None
                         } else {
                             Some(1)
                         },
                     );
                     instructions.extend(instr);
-                    original_param_regs.extend(registers);
+                    if registers.len() > 0 {
+                        original_param_regs.push(*registers.first().unwrap());
+                    } else {
+                        vararg = true;
+                    }
                 }
 
-                let mut registers = scope
-                    .register_counter
-                    .consecutive(original_param_regs.len() + 1)
-                    .into_iter();
                 let function_reg = registers.next().unwrap();
                 let mut param_regs = Vec::new();
                 for _ in &original_param_regs {
@@ -600,6 +633,11 @@ impl Expression {
                     instructions.push(Instruction::Move(function_reg, *old_function_reg));
                 }
 
+                if vararg {
+                    let last_reg = param_regs.last().unwrap_or(&function_reg) + 1;
+                    instructions.push(Instruction::MoveRetValues(last_reg));
+                }
+
                 let last_param_reg = param_regs.last().unwrap_or(&function_reg);
 
                 let mut return_regs = Vec::new();
@@ -616,7 +654,7 @@ impl Expression {
 
                 instructions.push(Instruction::Call(
                     *&function_reg,
-                    param_regs.len() as u16,
+                    if vararg { 0 } else { param_regs.len() as u16 },
                     if return_regs.len() == 0 {
                         0
                     } else {
@@ -647,7 +685,7 @@ impl Expression {
                 instructions.push(Instruction::NewTable(reg));
 
                 let mut counter = 1;
-                for (key, value) in entries {
+                for (i, (key, value)) in entries.iter().enumerate() {
                     if let Some(key) = key {
                         let (instr, key_regs) = key.kind.compile(state, scope, Some(1));
                         instructions.extend(instr);
@@ -662,7 +700,11 @@ impl Expression {
                         let (instr, value_regs) = value.kind.compile(
                             state,
                             scope,
-                            if entries.len() == 1 { None } else { Some(1) },
+                            if i == entries.len() - 1 {
+                                None
+                            } else {
+                                Some(1)
+                            },
                         );
                         instructions.extend(instr);
 
@@ -679,7 +721,7 @@ impl Expression {
                             ));
                             counter += 1;
                         } else {
-                            instructions.push(Instruction::SetList(reg));
+                            instructions.push(Instruction::SetList(reg, counter as u32));
                         }
                     }
                 }
@@ -702,6 +744,15 @@ impl Expression {
 
                 (instructions, vec![local])
             }
+            Expression::Ellipsis => {
+                if !scope.is_vararg {
+                    panic!("Function is not vararg!");
+                }
+                let mut instructions = Vec::new();
+                let new_reg = scope.register_counter.new();
+                instructions.push(Instruction::Vararg(new_reg));
+                (instructions, Vec::new())
+            }
         }
     }
 }

src/lib.rs

@@ -70,7 +70,10 @@ impl Debug for CompilationUnit {
 
 impl CompilationUnit {
     pub fn with_virtual_machine<'a>(&self, vm: &'a mut VirtualMachine) -> ExecutionUnit<'a> {
-        let chunk_id = vm.new_prototype(self.instructions.clone());
+        let chunk_id = vm.new_prototype(vm::Prototype {
+            instructions: self.instructions.clone(),
+            parameters: 0,
+        });
         for prototype in &self.state.prototypes {
             vm.new_prototype(prototype.clone());
         }

src/vm.rs

@@ -110,6 +110,8 @@ impl Debug for Constant {
 pub enum Instruction {
     /// R(A) := R(B)
     Move(u16, u16),
+    /// R(A) ... R(A + func_register - top) := previous function return values
+    MoveRetValues(u16),
     /// R(A) := K(Bx)
     LoadK(u16, u32),
     /// R(A), ..., R(B) := nil
@@ -124,8 +126,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
+    /// R(A)[B...] := all values returned from previous function
-    SetList(u16),
+    SetList(u16, u32),
     /// R(A) := R(B)[R(C)]
     GetTable(u16, u16, u16),
     /// R(A) := {}
@@ -159,12 +161,15 @@ pub enum Instruction {
     Close(u16),
     /// R(A) := closure(KPROTO[Bx], R(A), ..., R(A+n))
     Closure(u16, u32),
+    /// R(A), ... := varargs
+    Vararg(u16),
 }
 
 impl Debug for Instruction {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Instruction::Move(arg0, arg1) => write!(f, "MOVE {} {}", arg0, arg1),
+            Instruction::MoveRetValues(arg0) => write!(f, "MOVERETVALS {}", arg0),
             Instruction::LoadK(arg0, arg1) => write!(f, "LOADK {} {}", arg0, arg1),
             Instruction::SetGlobal(arg0, arg1) => write!(f, "SETGLOBAL {} {}", arg0, arg1),
             Instruction::GetGlobal(arg0, arg1) => write!(f, "GETGLOBAL {} {}", arg0, arg1),
@@ -176,7 +181,7 @@ impl Debug for Instruction {
             Instruction::GetTable(arg0, arg1, arg2) => {
                 write!(f, "GETTABLE {} {} {}", arg0, arg1, arg2)
             }
-            Instruction::SetList(arg0) => write!(f, "SETLIST {}", arg0),
+            Instruction::SetList(arg0, arg1) => write!(f, "SETLIST {} {}", arg0, arg1),
             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),
@@ -195,6 +200,7 @@ impl Debug for Instruction {
             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),
+            Instruction::Vararg(arg0) => write!(f, "VARARG {}", arg0),
         }
     }
 }
@@ -506,16 +512,22 @@ impl<T: RustFunction + 'static> From<T> for Value {
     }
 }
 
+#[derive(Debug, Clone, Default)]
+pub struct Prototype {
+    pub instructions: Vec<Instruction>,
+    pub parameters: usize,
+}
+
 #[derive(Debug, Clone, Default)]
 pub struct VirtualMachine {
     pub(super) environment: Rc<RefCell<Environment>>,
     pub(super) constants: Vec<Constant>,
-    pub(super) prototypes: HashMap<u32, Vec<Instruction>>,
+    pub(super) prototypes: HashMap<u32, Prototype>,
     pub(super) proto_counter: u32,
 }
 
 impl VirtualMachine {
-    pub fn new_prototype(&mut self, instructions: Vec<Instruction>) -> u32 {
+    pub fn new_prototype(&mut self, instructions: Prototype) -> u32 {
         let proto_id = self.proto_counter;
         self.proto_counter += 1;
         self.prototypes.insert(proto_id, instructions);
@@ -571,6 +583,7 @@ impl Closure {
             function_register: 0,
             return_registers: Vec::new(),
             top: 0,
+            parameters: params,
         }
     }
 
@@ -594,6 +607,7 @@ pub struct ClosureRunner {
     function_register: u16,
     return_registers: Vec<u16>,
     top: u16,
+    parameters: Vec<Value>,
 }
 
 #[derive(Clone, Debug)]
@@ -648,7 +662,10 @@ impl ClosureRunner {
     pub fn execute(&mut self, unit: &CompilationUnit) -> ClosureRunner {
         let mut vm = self.closure.vm.clone();
         vm.constants = unit.constants.clone();
-        let proto_id = vm.new_prototype(unit.instructions.clone());
+        let proto_id = vm.new_prototype(Prototype {
+            instructions: unit.instructions.clone(),
+            parameters: 0,
+        });
         for prototype in &unit.state.prototypes {
             vm.new_prototype(prototype.clone());
         }
@@ -693,7 +710,7 @@ impl ClosureRunner {
             }
         }
 
-        let instructions = self
+        let prototype = self
             .closure
             .vm
             .prototypes
@@ -702,12 +719,26 @@ impl ClosureRunner {
             .clone();
         let constants = self.closure.vm.constants.clone();
 
-        if let Some(instr) = instructions.get(self.program_counter) {
+        if let Some(instr) = prototype.instructions.get(self.program_counter) {
             match instr {
                 Instruction::Move(a, b) => {
                     let b = self.get_stack(*b);
                     self.set_stack(*a, b);
                 }
+                Instruction::MoveRetValues(res) => {
+                    let length = self.top - self.function_register;
+
+                    let mut values = Vec::new();
+                    for i in 0..length {
+                        let b = self.get_stack(self.function_register + i + 1);
+                        values.push(b);
+                    }
+                    for (i, val) in values.into_iter().enumerate() {
+                        self.set_stack(*res + i as u16, val);
+                    }
+
+                    self.top = *res + length;
+                }
                 Instruction::LoadK(reg, constant) => {
                     self.set_stack(
                         *reg,
@@ -817,13 +848,13 @@ impl ClosureRunner {
 
                     self.set_stack(*res, value);
                 }
-                Instruction::SetList(reg) => {
+                Instruction::SetList(reg, start_idx) => {
                     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;
+                                let mut counter = *start_idx as i64;
                                 for i in self.function_register..self.top {
                                     let value = self.get_stack(i + 1);
                                     match value {
@@ -868,14 +899,10 @@ impl ClosureRunner {
                     }
                 }
                 Instruction::Call(func_reg, param_len, ret_len) => {
-                    let param_start_func_reg = if *param_len == 0 {
+                    let param_start_func_reg = *func_reg;
-                        self.function_register
-                    } else {
-                        *func_reg
-                    };
 
                     let param_len = if *param_len == 0 {
-                        self.top - self.top.min(param_start_func_reg)
+                        self.top - self.top.min(param_start_func_reg + 1)
                     } else {
                         *param_len
                     };
@@ -950,8 +977,8 @@ impl ClosureRunner {
                     self.program_counter += 1;
                     let mut ret_values = Vec::new();
                     let (reg_start, reg_end) = if *reg_end == 0 {
-                        if self.function_register > 0 && self.top > 0 {
+                        if self.top > 0 {
-                            (self.function_register + 1, self.top)
+                            (*reg_start, self.top - 1)
                         } else {
                             (*reg_start, *reg_end)
                         }
@@ -1070,6 +1097,21 @@ impl ClosureRunner {
                         .unwrap_or(Value::Nil);
                     self.set_stack(*res, StackValue::Value(lhs.and(&rhs)?));
                 }
+                Instruction::Vararg(reg) => {
+                    self.function_register = reg - 1;
+                    self.top = self.function_register + self.parameters.len() as u16
+                        - prototype.parameters as u16;
+
+                    for (i, param) in self
+                        .parameters
+                        .clone()
+                        .iter()
+                        .skip(prototype.parameters)
+                        .enumerate()
+                    {
+                        self.set_stack(*reg + i as u16, StackValue::Value(param.clone()));
+                    }
+                }
             };
 
             self.program_counter += 1;