diff --git a/lib/src/backend/lua/bytecode.rs b/lib/src/backend/lua/bytecode.rs index 22c1492..94a2944 100644 --- a/lib/src/backend/lua/bytecode.rs +++ b/lib/src/backend/lua/bytecode.rs @@ -1,8 +1,8 @@ use indexmap::IndexSet; use std::fmt::{Display, Formatter, Write}; use std::hash::{Hash, Hasher}; +use smallvec::SmallVec; -use crate::ast::SmallVec8; use crate::parser::StString; use super::register::Register; @@ -14,7 +14,7 @@ macro_rules! excess_k { }; } -macro_rules! excess_sBx { +macro_rules! excess_sbx { ($v: expr) => { excess_k!($v, 17) }; @@ -223,7 +223,7 @@ pub enum LuaByteCode { /// Call k, v: k is callee symbol position, v is argument count, return value not included /// A B C: R[A], ... ,R[A+C-2] := R[A](R[A+1], ... ,R[A+B-1]) - Call(u8, u8, u8), + Call(Register, u8, u8), /// A (adjust vararg parameters) VarArgPrep(u8), @@ -275,15 +275,15 @@ impl LuaByteCode { // A sB k LuaByteCode::GetTabUp(a, sb, c) | LuaByteCode::Gei(a, sb, c) - | LuaByteCode::Gti(a, sb, c) => (c as u32) << 17 | (sb as u32) << 9 | a.num() as u32, + | LuaByteCode::Gti(a, sb, c) + | LuaByteCode::Call(a, sb, c) => (c as u32) << 17 | (sb as u32) << 9 | a.num() as u32, // A B C all literal LuaByteCode::Return(a, b, c) - | LuaByteCode::SetTabUp(a, b, c) - | LuaByteCode::Call(a, b, c) => (c as u32) << 17 | (b as u32) << 9 | a as u32, + | LuaByteCode::SetTabUp(a, b, c) => (c as u32) << 17 | (b as u32) << 9 | a as u32, // ABx LuaByteCode::LoadK(a, bx) => (bx as u32) << 8 | a.num() as u32, // AsBx - LuaByteCode::LoadI(a, sbx) => excess_sBx!(sbx) << 8 | a.num() as u32, + LuaByteCode::LoadI(a, sbx) => excess_sbx!(sbx) << 8 | a.num() as u32, // A B LuaByteCode::Move(a, b) => (b.num() as u32) << 9 | a.num() as u32, // A only @@ -299,7 +299,7 @@ impl LuaByteCode { pub struct LuaCompiledCode { pub byte_codes: Vec, pub constants: IndexSet, - pub upvalues: SmallVec8, + pub upvalues: SmallVec<[LuaUpValue; 32]>, } impl LuaCompiledCode { @@ -331,13 +331,10 @@ impl LuaCompiledCode { // A sB k LuaByteCode::GetTabUp(a, b, c) | LuaByteCode::Gti(a, b, c) - | LuaByteCode::Gei(a, b, c) => { + | LuaByteCode::Gei(a, b, c) + | LuaByteCode::Call(a, b, c) => { write!(s, "R{} {b} {c}", a.num()).unwrap(); } - // A B C all literal - LuaByteCode::Call(a, b, c) => { - write!(s, "{a} {b} {c}").unwrap(); - } // ABC with k LuaByteCode::SetTabUp(a, b, c) | LuaByteCode::Return(a, b, c) => { write!(s, "{a} {b} {c}").unwrap(); @@ -368,7 +365,17 @@ impl LuaCompiledCode { write!(s, " ; _ENV \"{}\"", self.constants[*c as usize]).unwrap(); } LuaByteCode::Call(a, b, c) => { - write!(s, " ; {b} in {c} out").unwrap(); + if *b == 0 { + write!(s, " ; all in ").unwrap(); + } else { + write!(s, " ; {} in ", b - 1).unwrap(); + } + + if *c == 0 { + write!(s, "all out").unwrap(); + } else { + write!(s, "{} out", c - 1).unwrap(); + } } _ => {} } @@ -379,18 +386,29 @@ impl LuaCompiledCode { impl Display for LuaCompiledCode { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - // constants - writeln!(f, "Constants:")?; - for (idx, constant) in self.constants.iter().enumerate() { - writeln!(f, "{:<3} {:?}", idx, constant)? - } + // upvalues + if !self.upvalues.is_empty() { + writeln!(f, "UpValues:")?; - writeln!(f)?; + for (idx, up_val) in self.upvalues.iter().enumerate() { + writeln!(f, "{:<4} {:?}", idx, up_val)? + } + writeln!(f)?; + } // constants + if !self.constants.is_empty() { + writeln!(f, "Constants:")?; + for (idx, constant) in self.constants.iter().enumerate() { + writeln!(f, "{:<4} {:?}", idx, constant)? + } + writeln!(f)?; + } + + // ByteCodes writeln!(f, "ByteCodes:")?; for (idx, bc) in self.byte_codes.iter().enumerate() { - writeln!(f, "{:<6} {:<20}", idx + 1, self.disassembly_code(bc))? + writeln!(f, "{:<4} {:<20}", idx + 1, self.disassembly_code(bc))? } Ok(()) diff --git a/lib/src/backend/lua/dump.rs b/lib/src/backend/lua/dump.rs index da7bb3e..bc57bc4 100644 --- a/lib/src/backend/lua/dump.rs +++ b/lib/src/backend/lua/dump.rs @@ -1,12 +1,13 @@ -use byteorder::{LittleEndian, WriteBytesExt}; +use std::io; +use std::io::Write; -use super::{utils::*, LuaConstants}; -use super::{Function, LuaBackend, Prototype}; +use byteorder::{LittleEndian, WriteBytesExt}; -use crate::backend::lua::bytecode::LuaCompiledCode; use crate::backend::CompiledCode; -use std::io; -use std::io::Write; +use crate::backend::lua::bytecode::LuaCompiledCode; + +use super::{LuaConstants, utils::*}; +use super::{Function, LuaBackend, Prototype}; /// Lua signature const LUA_SIGNATURE: &str = "\x1bLua"; @@ -37,8 +38,8 @@ pub fn lua_dump_module(backend: &LuaBackend, w: &mut dyn Write) -> io::Result<() // LUAC_NUMBER lua_dump_bytes(w, &LUAC_NUMBER.to_le_bytes())?; - // size of UpValues in 1 byte - lua_dump_byte(w, backend.module_upvalues().len() as u8)?; + // size of UpValues in 1 byte, TODO: hard-coded 1 + lua_dump_byte(w, 1)?; // Start to dump functions // get main function @@ -183,7 +184,11 @@ fn lua_dump_bytes(w: &mut dyn Write, bytes: &[u8]) -> io::Result<()> { #[inline] fn lua_dump_string(w: &mut dyn Write, opt_str: Option<&String>) -> io::Result<()> { match opt_str { - Some(s) => todo!(), + Some(s) => { + let bytes = s.as_bytes(); + lua_dump_size(w, (bytes.len() + 1) as u64)?; + w.write_all(bytes) + }, None => lua_dump_size(w, 0), } } diff --git a/lib/src/backend/lua/mod.rs b/lib/src/backend/lua/mod.rs index 1f310ef..a8bb650 100644 --- a/lib/src/backend/lua/mod.rs +++ b/lib/src/backend/lua/mod.rs @@ -84,9 +84,11 @@ pub struct LuaBackend { states: SmallVec<[LuaBackendStates; 32]>, local_function: Option, local_proto: Option, - constants: IndexSet, reg_mgr: RegisterManager, - module_upvalues: SmallVec<[LuaConstants; 32]>, + + // tmp values for generating function + upvalues: SmallVec<[LuaUpValue; 32]>, + constants: IndexSet, } impl LuaBackend { @@ -136,11 +138,6 @@ impl LuaBackend { self.app.clone() } - #[inline] - fn module_upvalues(&self) -> &SmallVec<[LuaConstants; 32]> { - &self.module_upvalues - } - fn push_attribute_with_scope(&mut self, scope: Scope) { let attr = LuaBackendStates { scope: Some(scope), @@ -220,14 +217,10 @@ impl CodeGenBackend for LuaBackend { local_proto: None, constants: IndexSet::new(), reg_mgr: RegisterManager::new(), - module_upvalues: smallvec![], + upvalues: smallvec![], } } - fn get_module_bytes(&mut self, w: &mut dyn Write) -> io::Result<()> { - lua_dump_module(self, w) - } - fn gen_function(&mut self, func: usize) -> Result, CodeGenError> { let app = self.app.read(); let f = app @@ -250,6 +243,14 @@ impl CodeGenBackend for LuaBackend { } } + // create upvalue table + self.upvalues.push(LuaUpValue { + name: None, + stack: 1, + index: 0, + kind: 0, + }); + let mut fun = f.write(); self.push_attribute_with_scope(fun_scope); self.visit_statement_mut(fun.parse_tree_mut()); @@ -260,17 +261,15 @@ impl CodeGenBackend for LuaBackend { let byte_codes = mem::take(&mut self.byte_codes); let constants = mem::replace(&mut self.constants, IndexSet::new()); + let upvalues = mem::replace(&mut self.upvalues, smallvec![]); + + // reset RegMan and check register is balance + assert!(self.reg_mgr.check_and_reset()); Ok(Box::new(LuaCompiledCode { byte_codes, constants, - // TODO:: upvalues - upvalues: smallvec![LuaUpValue { - name: None, - stack: 1, - index: 0, - kind: 0 - }], + upvalues, })) } @@ -287,6 +286,10 @@ impl CodeGenBackend for LuaBackend { self.visit_expression_mut(&mut operands[0]); } + + fn get_module_bytes(&mut self, w: &mut dyn Write) -> io::Result<()> { + lua_dump_module(self, w) + } } impl AstVisitorMut for LuaBackend { @@ -328,8 +331,14 @@ impl AstVisitorMut for LuaBackend { match access_mode & (LuaAccessMode::READ | LuaAccessMode::WRITE | LuaAccessMode::CALL) { // Callee process LuaAccessMode::CALL => { + assert!(self.top_attribute().register.is_none()); + + let dest = self.reg_mgr.alloc_hard(); + self.top_attribute().register = Some(dest); self.top_attribute().constant_index = - Some(self.add_string_constant(var_expr.org_name())) + Some(self.add_string_constant(var_expr.org_name())); + + self.push_code(LuaByteCode::GetTabUp(dest, 0, 0)); } // Write register into stack LuaAccessMode::WRITE => {} @@ -337,11 +346,13 @@ impl AstVisitorMut for LuaBackend { LuaAccessMode::READ => { let scope = self.top_attribute().scope.as_ref().unwrap(); if let Some(variable) = scope.find_variable(var_expr.name()) { - let reg = self.reg_mgr.alloc_hard(); - self.top_attribute().register = Some(reg); + self.top_attribute().register = Some(self.reg_mgr.alloc_local_variable(variable.name())); - // TODO: initialize - self.push_code(LuaByteCode::LoadI(reg, 0)); + // let reg = self.reg_mgr.alloc_hard(); + // self.top_attribute().register = Some(reg); + // + // // TODO: initialize + // self.push_code(LuaByteCode::LoadI(reg, 0)); } else { // TODO: variable not found error } @@ -355,30 +366,35 @@ impl AstVisitorMut for LuaBackend { self.push_access_attribute(LuaAccessMode::CALL); self.visit_expression_mut(call.callee_mut()); - let callee_index = self.top_attribute().constant_index; - self.pop_attribute(); + let callee_attr = self.pop_attribute(); + let callee_reg = callee_attr.register.unwrap(); // TODO: // self.push_code(LuaByteCode::GetTabUp(0, 0, 0)); // visit all arguments for arg in call.arguments_mut() { + self.push_access_attribute(LuaAccessMode::READ); self.visit_expression_mut(arg); - let arg_value_index = self.top_attribute().constant_index; - self.pop_attribute(); + let arg_attr = self.pop_attribute(); + let arg_value_index = arg_attr.constant_index; // Load argument if let Some(idx) = arg_value_index { // TODO: // self.push_code(LuaByteCode::LoadK(0, idx as u32)); } + + self.reg_mgr.free(&arg_attr.register.unwrap()); } self.push_code(LuaByteCode::Call( - callee_index.unwrap(), - call.arguments().len() as u8, - 0, - )) + callee_reg, + call.arguments().len() as u8 + 1, + 1, + )); + + self.reg_mgr.free(&callee_reg); } fn visit_if_statement_mut(&mut self, ifst: &mut IfStatement) { @@ -456,15 +472,18 @@ impl AstVisitorMut for LuaBackend { let lhs_reg = self.pop_attribute().register.unwrap(); if let Some(constant_index) = rhs.constant_index { + assert!(rhs.register.is_none()); self.code_load_constant(lhs_reg, constant_index) } else if let Some(r) = rhs.register { - assert!(lhs_reg != r); + assert_ne!(lhs_reg, r); self.reg_mgr.free(&r); self.code_move(r, lhs_reg); } else { panic!() } - self.top_attribute().register = Some(lhs_reg); + // TODO: keep register? + // self.top_attribute().register = Some(lhs_reg); + self.reg_mgr.free(&lhs_reg); } } diff --git a/lib/src/backend/lua/register.rs b/lib/src/backend/lua/register.rs index 70c9be5..d5f708c 100644 --- a/lib/src/backend/lua/register.rs +++ b/lib/src/backend/lua/register.rs @@ -1,8 +1,12 @@ use std::collections::HashSet; +use log::warn; +use smallmap::{Map as SmallMap, smallmap}; + +use crate::parser::StString; const MAX_REGISTER_ID: u8 = 255; -#[derive(PartialEq, Eq, Clone, Debug, Copy)] +#[derive(PartialEq, Eq, Clone, Debug, Copy, Hash)] pub enum Register { VirtualRegister(usize), LuaRegister(u8), @@ -32,6 +36,8 @@ pub struct RegisterManager { virtual_register_cursor: usize, real_register_cursor: u8, used_real_registers: HashSet, + local_variable_register: SmallMap, + local_variable_register_reverse: SmallMap, } impl RegisterManager { @@ -41,7 +47,48 @@ impl RegisterManager { virtual_register_cursor: 0, real_register_cursor: 0, used_real_registers: HashSet::with_capacity(MAX_REGISTER_ID as usize), + local_variable_register: smallmap![], + local_variable_register_reverse: smallmap![], + } + } + + #[inline] + pub fn alloc_local_variable(&mut self, v: &StString) -> Register { + match self.local_variable_register.get(v) { + Some(r) => *r, + None => { + // TODO: allocate virtual + let r = self.alloc_hard(); + self.local_variable_register.insert(v.clone(), r); + self.local_variable_register_reverse.insert(r, v.clone()); + + r + } + } + } + + /// Reset RegMan, and return register usage is balance + #[inline] + pub fn check_and_reset(&mut self) -> bool { + // free all local variable registers + for (r, _) in self.local_variable_register_reverse.iter() { + if let Register::LuaRegister(x) = r { + self.used_real_registers.remove(x); + } + } + + let balanced = self.used_real_registers.is_empty(); + if !balanced { + warn!("Registers in use: {:?}", self.used_real_registers); } + + self.virtual_register_cursor = 0; + self.real_register_cursor = 0; + self.used_real_registers.clear(); + self.local_variable_register.clean(); + self.local_variable_register_reverse.clean(); + + balanced } pub fn alloc(&mut self) -> Register { @@ -67,6 +114,11 @@ impl RegisterManager { #[inline] pub fn free(&mut self, reg: &Register) { + // prevent free for local variable register + if self.local_variable_register_reverse.contains_key(®) { + return; + } + match *reg { Register::LuaRegister(x) => { self.used_real_registers.remove(&x); diff --git a/viewer/test_projects/example1/test_proj.xml b/viewer/test_projects/example1/test_proj.xml index e5e1ff9..9572173 100644 --- a/viewer/test_projects/example1/test_proj.xml +++ b/viewer/test_projects/example1/test_proj.xml @@ -22,8 +22,7 @@ end_program ]]> - a := 123; - a := b; + print(a);