update: fix some bugs

Signed-off-by: yjhmelody <[email protected]>

Author: yjhmelody

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> {