update: fix some bugs

Signed-off-by: yjhmelody <465402634@qq.com>

src/binary/mod.rs

@@ -4,6 +4,7 @@ pub mod chunk;
 pub mod reader;
 pub mod writer;
 
+
 /// decode Lua binary chunk to prototype structure
 pub fn decode(data: Vec<u8>) -> Rc<chunk::Prototype> {
  let mut r = reader::Reader::new(data);
@@ -12,11 +13,11 @@ pub fn decode(data: Vec<u8>) -> Rc<chunk::Prototype> {
  r.read_proto()
 }
 
-pub fn encode(proto: Rc<chunk::Prototype>) -> Vec<u8> {
+pub fn encode(proto: Rc<chunk::Prototype>, src: Option<String>) -> Vec<u8> {
  let mut writer = writer::Writer::new();
  writer.write_header();
- writer.write_byte(4);
- writer.write_proto(proto);
+ writer.write_byte(1);
+ writer.write_proto(proto, src);
  writer.as_bytes()
 }
 
@@ -46,12 +47,18 @@ mod tests {
 
  #[test]
  fn test_encode() {
- let s = fs::read("D:/code/Rust/lua-rs/tests/luac.out").expect("error");
- let proto = encode(decode(s));
-// println!("{:#?}", proto);
+ let chunk = fs::read("D:/code/Rust/lua-rs/tests/luac.out").expect("error");
+ let proto = encode(decode(chunk.clone()), Some("@example.lua".to_string()));
+ let proto = encode(decode(chunk.clone()), Some("@hello.lua".to_string()));
 
- let s = unsafe { String::from_utf8_unchecked(proto) };
+ let s = unsafe { String::from_utf8_unchecked(proto.clone()) };
  fs::write("D:/code/Rust/lua-rs/tests/test.out", s);
- assert_eq!(2, 0);
+ let mut i = 0;
+ for (a, b) in chunk.iter().zip(proto.iter()) {
+ println!("{}", i);
+ assert_eq!(*a, *b);
+ i += 1;
+ }
+ assert_eq!(chunk[33..], proto[33..]);
  }
 }

src/binary/reader.rs

@@ -64,11 +64,12 @@ impl Reader {
 
  fn read_string0(&mut self) -> Option<String> {
  let mut size = self.read_byte() as usize;
+ dbg!(&size);
  if size == 0 {
  return None;
  }
  if size == 0xFF {
- size = self.read_u64() as usize; // size_t
+ size = self.read_u32() as usize; // size_t
  }
  let bytes = self.read_bytes(size - 1);
  let string = String::from_utf8(bytes);
@@ -88,6 +89,7 @@ impl Reader {
  }
 
  pub fn check_header(&mut self) {
+ // 17 + 16 = 33
  assert_eq!(self.read_bytes(4), LUA_SIGNATURE, "not a precompiled chunk!");
  assert_eq!(self.read_byte(), LUAC_VERSION, "version mismatch!");
  assert_eq!(self.read_byte(), LUAC_FORMAT, "format mismatch!");

src/binary/writer.rs

@@ -68,9 +68,10 @@ impl Writer {
 
  fn write_string0(&mut self, s: &String) -> Option<()> {
  if s.len() == 0 {
+ self.write_byte(0);
  None
  } else if s.len() < 0xFF {
- self.write_byte(s.len() as u8 - 1);
+ self.write_byte(s.len() as u8 + 1);
  self.write_bytes(s.as_bytes().to_vec());
  Some(())
  } else {
@@ -93,8 +94,9 @@ impl Writer {
  self.write_lua_number(LUAC_NUM);
  }
 
- pub fn write_proto(&mut self, proto: Rc<Prototype>) {
- self.write_string0(&proto.source.clone().unwrap());
+ pub fn write_proto(&mut self, proto: Rc<Prototype>, parent_source: Option<String>) {
+ self.write_string0(&parent_source.clone().unwrap_or_default());
+
  self.write_u32(proto.line_defined);
  self.write_u32(proto.last_line_defined);
  self.write_byte(proto.num_params);
@@ -118,7 +120,7 @@ impl Writer {
 
  self.write_u32(proto.prototypes.len() as u32);
  for prototype in proto.prototypes.iter() {
- self.write_proto(prototype.clone());
+ self.write_proto(prototype.clone(), None);
  }
 
  self.write_u32(proto.line_info.len() as u32);

src/compiler/codegen2.rs

@@ -60,7 +60,12 @@ pub fn gen_prototype(block: Box<Block>) -> Result<Rc<Prototype>> {
  let last_line = block.last_line;
  let fn_def = FnDef::new(ParList::default(), block, 0, last_line);
  let mut fn_info = FnInfo::new(None, ParList::default(), 0, last_line);
- fn_info.add_local_var("_ENV".to_string(), 0)?;
+// fn_info.add_local_var("_ENV".to_string(), 0)?;
+ fn_info.up_values.insert(
+ "_ENV".to_string(),
+ UpValueInfo::new(Some(0), Some(0), 0),
+ );
+
  fn_info.codegen_fn_def_exp(&fn_def, 0)?;
 
  Ok(fn_info.to_prototype())
@@ -192,7 +197,7 @@ impl FnInfo {
  used_regs: 0,
  max_regs: 0,
  scope_level: 0,
- local_vars: Vec::new(),
+ local_vars: vec![HashMap::new()],
  breaks: Vec::new(),
 // parent_index,
 // current_index,
@@ -208,17 +213,6 @@ impl FnInfo {
  }
  }
 
- fn find_fn_info_by_index(mut fn_info: Rc<RefCell<FnInfo>>, index: usize) -> Option<Rc<RefCell<FnInfo>>> {
- unimplemented!();
-// if fn_info.borrow_mut().current_index == index {
-// Some(fn_info)
-// } else {
-// for sub_fn in fn_info.borrow_mut().sub_fns.iter() {
-// return Self::find_fn_info_by_index(sub_fn.0.clone(), index);
-// }
-// None
-// }
- }
 
  fn constant_index(&mut self, k: &Constant) -> usize {
  match self.constants.get(k) {
@@ -335,10 +329,10 @@ impl FnInfo {
  }));
  let slot = new_var.borrow_mut().slot;
  if self.local_vars.len() <= self.scope_level {
- self.local_vars.resize(self.local_vars.len() + 2, HashMap::new());
+ self.local_vars.resize(self.local_vars.len() * 2, HashMap::new());
  }
- self.local_vars[self.scope_level].insert(name, new_var);
 
+ self.local_vars[self.scope_level].insert(name, new_var);
  Ok(slot)
  }
 
@@ -455,6 +449,7 @@ impl FnInfo {
  })
  }
  }
+
  res
  }
 
@@ -744,7 +739,7 @@ impl FnInfo {
  // clear sBx Op
  ins = ins << 18 >> 18;
  // reset sBx op
- ins = ins | (sBx as u32 + MAXARG_SBX as u32) << 14;
+ ins = ins | ((sBx + MAXARG_SBX) as u32) << 14;
  self.instructions[pc] = ins;
  }
 }
@@ -984,9 +979,6 @@ impl FnInfo {
  self.exit_scope()
  }
 
- fn remove_tail_nils(exps: &Vec<Exp>) -> Result<Vec<Exp>> {
- unimplemented!()
- }
 
  fn codegen_assign_stat(&mut self, names: &Vec<Exp>, vals: &Vec<Exp>, line: Line) -> Result<()> {
  let old_regs = self.used_regs;
@@ -1065,7 +1057,6 @@ impl FnInfo {
  self.emit_set_table(last_line, a as isize, k_regs[i], v_regs[i]);
  }
  } else {
- dbg!();
  let a = self.up_value_index(&"_ENV".to_string())
  .ok_or(Error::NotUpValue {
  line: last_line,
@@ -1125,14 +1116,13 @@ impl FnInfo {
  let a = self.alloc_registers(n)?;
  self.emit_load_nil(last_line, a as isize, n as isize);
  }
+ }
 
- self.used_regs = old_regs;
- let start_pc = self.pc() + 1;
- for name in names {
- self.add_local_var(name.clone(), start_pc)?;
- }
+ self.used_regs = old_regs;
+ let start_pc = self.pc() + 1;
+ for name in names {
+ self.add_local_var(name.clone(), start_pc)?;
  }
- dbg!();
 
  Ok(())
  }
@@ -1172,12 +1162,11 @@ impl FnInfo {
  let a = self.alloc_registers(n)?;
  self.emit_load_nil(last_line, a as isize, n as isize);
  }
-
- self.used_regs = old_regs;
- let start_pc = self.pc() + 1;
- for name in names {
- self.add_local_var(name.clone(), start_pc)?;
- }
+ }
+ self.used_regs = old_regs;
+ let start_pc = self.pc() + 1;
+ for name in names {
+ self.add_local_var(name.clone(), start_pc)?;
  }
 
  Ok(())
@@ -1449,6 +1438,9 @@ fn is_vararg_or_fn_call(exp: &Exp) -> bool {
 
 
 mod tests {
+ use std::fs;
+
+ use crate::binary::encode;
  use crate::compiler::lexer::*;
  use crate::compiler::parser::*;
 
@@ -1507,4 +1499,31 @@ mod tests {
  let proto = gen_prototype(Box::new(block));
  println!("{:#?}", proto);
  }
+
+ #[test]
+ fn test_codegen2() {
+ let s = r##"
+ local g = {
+ a = 1,
+ b = {}
+ }
+
+ print(g.a)
+ print(g.b)
+
+ local a = 1
+ while a < 1 do
+ a = a + 1
+ end
+ print(a)
+ "##.to_string();
+
+ let mut lexer = Lexer::from_iter(s.into_bytes(), "test".to_string());
+ let block = parse_block(&mut lexer).expect("parse error");
+ let proto = gen_prototype(Box::new(block));
+
+ let bytes = encode(proto.unwrap(), Some("@hello2.lua".to_string()));
+ fs::write("D:/code/Rust/lua-rs/tests/test.out", bytes);
+ assert_eq!(1, 2);
+ }
 }

src/compiler/parser.rs

@@ -358,6 +358,7 @@ fn _parse_var_list(lexer: &mut impl Lex, var0: Exp) -> Result<Vec<Exp>> {
 fn _parse_name_list(lexer: &mut impl Lex, name0: String) -> Result<Vec<String>> {
  let mut name_list = vec![name0];
  while let Ok(Token::SepComma) = lexer.look_ahead() {
+ lexer.skip_next_token();
  let name = lexer.next_ident()?;
  name_list.push(name);
  }
@@ -628,6 +629,7 @@ fn parse_table_constructor_exp(lexer: &mut impl Lex) -> Result<Exp> {
  if !lexer.check_next_token(Token::SepLcurly) {
  return Err(Error::IllegalExpression { line: lexer.current_line() });
  }
+ let line = lexer.current_line();
  // [fieldlist]
  let fields = _parse_field_list(lexer)?;
 
@@ -636,8 +638,7 @@ fn parse_table_constructor_exp(lexer: &mut impl Lex) -> Result<Exp> {
  return Err(Error::IllegalExpression { line: lexer.current_line() });
  }
 
- let last_line = lexer.current_line();
- Ok(Exp::TableConstructor(fields, last_line))
+ Ok(Exp::TableConstructor(fields, line))
 }
 
 fn parse_fn_def_exp(lexer: &mut impl Lex) -> Result<Exp> {