Add an optimization pass which removes SetAttrs.

If a Lua program doesn't use _ENV directly to manipulate a global variable,
then we can simply treat the variable just like any other local variable.

luacompiler/src/lib/bytecode/instructions.rs

@@ -28,7 +28,7 @@ pub fn make_instr(opcode: Opcode, arg1: u8, arg2: u8, arg3: u8) -> u32 {
 /// Represents a high level instruction whose operands have a size of usize.
 /// This is used by the frontend to create an SSA IR, which later gets translated
 /// into smaller instructions that fit in 32 bits.
-#[derive(PartialEq, Eq, Debug)]
+#[derive(PartialEq, Eq, Debug, Clone, Copy)]
 pub struct HLInstr(pub Opcode, pub usize, pub usize, pub usize);
 
 impl HLInstr {

luacompiler/src/lib/bytecodegen/mod.rs

@@ -1,7 +1,12 @@
 use bytecode::LuaBytecode;
 use irgen::lua_ir::LuaIR;
 
-pub fn compile_to_bytecode(ir: LuaIR) -> LuaBytecode {
+pub fn compile_to_bytecode(mut ir: LuaIR) -> LuaBytecode {
+    ir.optimize_env_lookups();
+    LuaIRToLuaBc::new(ir).compile()
+}
+
+pub fn compile_to_bytecode_no_opt(ir: LuaIR) -> LuaBytecode {
     LuaIRToLuaBc::new(ir).compile()
 }
 

luacompiler/src/lib/irgen/lua_ir.rs

@@ -1,24 +1,146 @@
-use bytecode::instructions::HLInstr;
+use bytecode::instructions::{HLInstr, Opcode};
 use irgen::{constants_map::ConstantsMap, register_map::Lifetime};
 
 /// Represents an IR in which all instructions are in SSA form.
 pub struct LuaIR {
     pub instrs: Vec<HLInstr>,
     pub const_map: ConstantsMap,
-    pub lifetimes: Vec<Lifetime>,
+    pub lifetimes: Vec<Option<Lifetime>>,
+    pub non_env_lookups: Vec<usize>,
 }
 
 impl LuaIR {
     pub fn new(
         instrs: Vec<HLInstr>,
         const_map: ConstantsMap,
         mut lifetimes: Vec<Lifetime>,
+        non_env_lookups: Vec<usize>,
     ) -> LuaIR {
         lifetimes.sort_by(|x, y| x.start_point().cmp(&y.start_point()));
         LuaIR {
             instrs,
             const_map,
-            lifetimes,
+            lifetimes: lifetimes.into_iter().map(|l| Some(l)).collect(),
+            non_env_lookups,
         }
     }
+
+    /// Optimizes the IR and removes unnecessary `SetAttr` instructions.
+    /// For example consider the following Lua program: `x = 2`
+    /// The compiler generates:
+    /// -------------
+    /// LDI     1 0 0 -- load IntTable[0] into register 1, R(1) = 2
+    /// LDS     2 0 0 -- load StringTable[0] into register 2, R(2) = "x"
+    /// SetAttr 0 2 1 -- _ENV["x"] = 2
+    /// -------------
+    /// If the user does not query _ENV["x"] directly, that means we can optimize
+    /// away the `LDS` and `SetAttr` instructions, and we can treat "x" like a local
+    /// variable instead.
+    pub fn optimize_env_lookups(&mut self) {
+        let mut new_instrs = vec![];
+        let mut new_lifetimes: Vec<Option<Lifetime>> = Vec::with_capacity(self.lifetimes.len());
+        // push _ENV
+        new_lifetimes.push(Some(Lifetime::new(0)));
+        for _ in 1..self.lifetimes.len() {
+            new_lifetimes.push(None);
+        }
+        let len = self.instrs.len();
+        let mut i = 0;
+        while i < len {
+            let instr = self.instrs[i];
+            if i < len - 2 && self.non_env_lookups.binary_search(&instr.1).is_ok() {
+                // skip the attribute load
+                i += 1;
+                let attr = self.instrs[i];
+                self.lifetimes[attr.1] = None;
+                // skip the SetAttr
+                i += 1;
+            }
+            new_instrs.push(instr);
+            let regs_to_update = Self::get_register_operands(instr);
+            for r in regs_to_update {
+                Self::update_end_point_to(&mut new_lifetimes[r], new_instrs.len());
+            }
+            i += 1;
+        }
+        self.instrs = new_instrs;
+        self.lifetimes = new_lifetimes;
+    }
+
+    /// Gets all the register numbers that are used by the given instruction.
+    fn get_register_operands(instr: HLInstr) -> Vec<usize> {
+        match instr.0 {
+            Opcode::LDI | Opcode::LDS | Opcode::LDF => vec![instr.1],
+            Opcode::MOV => vec![instr.1, instr.2],
+            _ => vec![instr.1, instr.2, instr.3],
+        }
+    }
+
+    /// Update the endpoint of a lifetime. If it doesn't exist, then create it.
+    fn update_end_point_to(lifetime: &mut Option<Lifetime>, ep: usize) {
+        match lifetime {
+            Some(ref mut l) => l.set_end_point(ep),
+            None => *lifetime = Some(Lifetime::new(ep - 1)),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use bytecode::instructions::Opcode;
+
+    #[test]
+    fn optimize_env_lookups_generates_correct_code() {
+        let instrs = vec![
+            HLInstr(Opcode::LDI, 1, 0, 0),
+            HLInstr(Opcode::LDS, 2, 0, 0),
+            HLInstr(Opcode::SetAttr, 0, 2, 1),
+            HLInstr(Opcode::LDI, 3, 0, 0),
+        ];
+        let lifetimes = vec![
+            Lifetime::with_end_point(0, 3),
+            Lifetime::with_end_point(0, 3),
+            Lifetime::with_end_point(1, 3),
+            Lifetime::with_end_point(3, 4),
+        ];
+        let non_env_lookups = vec![1];
+        let mut ir = LuaIR::new(instrs, ConstantsMap::new(), lifetimes, non_env_lookups);
+        ir.optimize_env_lookups();
+        let expected_instrs = vec![HLInstr(Opcode::LDI, 1, 0, 0), HLInstr(Opcode::LDI, 3, 0, 0)];
+        assert_eq!(ir.instrs, expected_instrs);
+        let expected_lifetimes = vec![
+            Some(Lifetime::with_end_point(0, 1)),
+            Some(Lifetime::with_end_point(0, 1)),
+            None,
+            Some(Lifetime::with_end_point(1, 2)),
+        ];
+        assert_eq!(ir.lifetimes, expected_lifetimes);
+    }
+
+    #[test]
+    fn optimize_env_lookups_empty_non_env_lookups() {
+        let instrs = vec![
+            HLInstr(Opcode::LDI, 1, 0, 0),
+            HLInstr(Opcode::LDS, 2, 0, 0),
+            HLInstr(Opcode::SetAttr, 0, 2, 1),
+            HLInstr(Opcode::LDI, 3, 0, 0),
+        ];
+        let lifetimes = vec![
+            Lifetime::with_end_point(0, 3),
+            Lifetime::with_end_point(0, 3),
+            Lifetime::with_end_point(1, 3),
+            Lifetime::with_end_point(3, 4),
+        ];
+        let mut ir = LuaIR::new(instrs.clone(), ConstantsMap::new(), lifetimes, vec![]);
+        ir.optimize_env_lookups();
+        assert_eq!(ir.instrs, instrs);
+        let expected_lifetimes = vec![
+            Some(Lifetime::with_end_point(0, 3)),
+            Some(Lifetime::with_end_point(0, 3)),
+            Some(Lifetime::with_end_point(1, 3)),
+            Some(Lifetime::with_end_point(3, 4)),
+        ];
+        assert_eq!(ir.lifetimes, expected_lifetimes);
+    }
 }

luacompiler/src/lib/irgen/mod.rs

@@ -21,6 +21,7 @@ struct LuaToIR<'a> {
     pt: &'a LuaParseTree,
     reg_map: RegisterMap<'a>,
     const_map: ConstantsMap,
+    non_env_lookups: Vec<usize>,
     instrs: Vec<HLInstr>,
 }
 
@@ -30,6 +31,7 @@ impl<'a> LuaToIR<'a> {
             pt,
             reg_map: RegisterMap::new(),
             const_map: ConstantsMap::new(),
+            non_env_lookups: vec![],
             instrs: vec![],
         }
     }
@@ -56,7 +58,12 @@ impl<'a> LuaToIR<'a> {
                 }
             }
         }
-        LuaIR::new(self.instrs, self.const_map, self.reg_map.get_lifetimes())
+        LuaIR::new(
+            self.instrs,
+            self.const_map,
+            self.reg_map.get_lifetimes(),
+            self.non_env_lookups,
+        )
     }
 
     fn compile_stat(&mut self, nodes: &Vec<Node<u8>>) {
@@ -80,6 +87,7 @@ impl<'a> LuaToIR<'a> {
                 self.reg_map.set_reg(name, value);
                 self.instrs
                     .push(HLInstr(Opcode::SetAttr, env_reg, attr_reg, value));
+                self.non_env_lookups.push(value);
             }
             _ => {}
         }
@@ -251,26 +259,23 @@ mod tests {
         }
         // check lifetimes
         let expected_lifetimes = vec![
-            Lifetime::with_end_point(0, 15),
-            Lifetime::with_end_point(1, 2),
-            Lifetime::with_end_point(2, 3),
-            Lifetime::with_end_point(3, 4),
-            Lifetime::with_end_point(4, 5),
-            Lifetime::with_end_point(5, 6),
-            Lifetime::with_end_point(6, 7),
-            Lifetime::with_end_point(7, 8),
-            Lifetime::with_end_point(8, 9),
-            Lifetime::with_end_point(9, 10),
-            Lifetime::with_end_point(10, 11),
-            Lifetime::with_end_point(11, 12),
-            Lifetime::with_end_point(12, 13),
-            Lifetime::with_end_point(13, 14),
-            Lifetime::with_end_point(14, 15),
+            Some(Lifetime::with_end_point(0, 15)),
+            Some(Lifetime::with_end_point(1, 2)),
+            Some(Lifetime::with_end_point(2, 3)),
+            Some(Lifetime::with_end_point(3, 4)),
+            Some(Lifetime::with_end_point(4, 5)),
+            Some(Lifetime::with_end_point(5, 6)),
+            Some(Lifetime::with_end_point(6, 7)),
+            Some(Lifetime::with_end_point(7, 8)),
+            Some(Lifetime::with_end_point(8, 9)),
+            Some(Lifetime::with_end_point(9, 10)),
+            Some(Lifetime::with_end_point(10, 11)),
+            Some(Lifetime::with_end_point(11, 12)),
+            Some(Lifetime::with_end_point(12, 13)),
+            Some(Lifetime::with_end_point(13, 14)),
+            Some(Lifetime::with_end_point(14, 15)),
         ];
-        assert_eq!(ir.lifetimes.len(), expected_lifetimes.len());
-        for (lhs, rhs) in ir.lifetimes.iter().zip(expected_lifetimes.iter()) {
-            assert_eq!(lhs, rhs);
-        }
+        assert_eq!(ir.lifetimes, expected_lifetimes);
         // check constats map
         let expected_ints = vec![1, 2, 3];
         let ints = ir.const_map.get_ints();
@@ -320,15 +325,12 @@ mod tests {
         }
         // check lifetimes
         let expected_lifetimes = vec![
-            Lifetime::with_end_point(0, 4),
-            Lifetime::with_end_point(1, 4),
-            Lifetime::with_end_point(2, 3),
-            Lifetime::with_end_point(3, 4),
+            Some(Lifetime::with_end_point(0, 4)),
+            Some(Lifetime::with_end_point(1, 4)),
+            Some(Lifetime::with_end_point(2, 3)),
+            Some(Lifetime::with_end_point(3, 4)),
         ];
-        assert_eq!(ir.lifetimes.len(), expected_lifetimes.len());
-        for (lhs, rhs) in ir.lifetimes.iter().zip(expected_lifetimes.iter()) {
-            assert_eq!(lhs, rhs);
-        }
+        assert_eq!(ir.lifetimes, expected_lifetimes);
         // check constats map
         let expected_ints = vec![1];
         let ints = ir.const_map.get_ints();

luacompiler/src/lib/irgen/register_map.rs

@@ -28,7 +28,7 @@ impl Lifetime {
         self.1
     }
 
-    fn set_end_point(&mut self, ep: usize) {
+    pub fn set_end_point(&mut self, ep: usize) {
         self.1 = ep
     }
 }

luacompiler/tests/integration_test.rs

@@ -14,11 +14,7 @@ fn ldi_generation() {
     assert_eq!(bc.reg_count(), 3);
     assert_eq!(bc.get_int(0), 1);
     assert_eq!(bc.get_string(0), "x");
-    let expected_instrs = vec![
-        make_instr(Opcode::LDI, 1, 0, 0),
-        make_instr(Opcode::LDS, 2, 0, 0),
-        make_instr(Opcode::SetAttr, 0, 2, 1),
-    ];
+    let expected_instrs = vec![make_instr(Opcode::LDI, 1, 0, 0)];
     assert_eq!(bc.instrs_len(), expected_instrs.len());
     for i in 0..expected_instrs.len() {
         assert_eq!(bc.get_instr(i), expected_instrs[i]);
@@ -32,11 +28,7 @@ fn ldf_generation() {
     assert_eq!(bc.reg_count(), 3);
     assert_eq!(bc.get_float(0).to_string(), "2");
     assert_eq!(bc.get_string(0), "x");
-    let expected_instrs = vec![
-        make_instr(Opcode::LDF, 1, 0, 0),
-        make_instr(Opcode::LDS, 2, 0, 0),
-        make_instr(Opcode::SetAttr, 0, 2, 1),
-    ];
+    let expected_instrs = vec![make_instr(Opcode::LDF, 1, 0, 0)];
     assert_eq!(bc.instrs_len(), expected_instrs.len());
     for i in 0..expected_instrs.len() {
         assert_eq!(bc.get_instr(i), expected_instrs[i]);
@@ -50,11 +42,7 @@ fn lds_generation() {
     assert_eq!(bc.reg_count(), 3);
     assert_eq!(bc.get_string(0), "1.2");
     assert_eq!(bc.get_string(1), "x");
-    let expected_instrs = vec![
-        make_instr(Opcode::LDS, 1, 0, 0),
-        make_instr(Opcode::LDS, 2, 1, 0),
-        make_instr(Opcode::SetAttr, 0, 2, 1),
-    ];
+    let expected_instrs = vec![make_instr(Opcode::LDS, 1, 0, 0)];
     assert_eq!(bc.instrs_len(), expected_instrs.len());
     for i in 0..expected_instrs.len() {
         assert_eq!(bc.get_instr(i), expected_instrs[i]);
@@ -71,8 +59,6 @@ fn assert_bytecode(opcode: Opcode, operation: &str) {
         make_instr(Opcode::LDI, 1, 0, 0),
         make_instr(Opcode::LDI, 2, 1, 0),
         make_instr(opcode, 3, 1, 2),
-        make_instr(Opcode::LDS, 4, 0, 0),
-        make_instr(Opcode::SetAttr, 0, 4, 3),
     ];
     assert_eq!(bc.instrs_len(), expected_instrs.len());
     for i in 0..expected_instrs.len() {

luavm/src/lib/instructions/tables.rs

@@ -36,15 +36,15 @@ mod tests {
     use lua_values::LuaVal;
     use luacompiler::{
         bytecode::instructions::{make_instr, Opcode},
-        bytecodegen::compile_to_bytecode,
+        bytecodegen::compile_to_bytecode_no_opt,
         irgen::compile_to_ir,
         LuaParseTree,
     };
 
     fn get_vm_for(p: String) -> Vm {
         let pt = LuaParseTree::from_str(p).unwrap();
         let ir = compile_to_ir(&pt);
-        let bc = compile_to_bytecode(ir);
+        let bc = compile_to_bytecode_no_opt(ir);
         Vm::new(bc)
     }
 

luavm/src/lib/mod.rs

@@ -68,12 +68,14 @@ impl Vm {
 mod tests {
     use super::*;
     use luacompiler::irgen::register_map::ENV_REG;
-    use luacompiler::{bytecodegen::compile_to_bytecode, irgen::compile_to_ir, LuaParseTree};
+    use luacompiler::{
+        bytecodegen::compile_to_bytecode_no_opt, irgen::compile_to_ir, LuaParseTree,
+    };
 
     fn get_vm_for(p: String) -> Vm {
         let pt = LuaParseTree::from_str(p).unwrap();
         let ir = compile_to_ir(&pt);
-        let bc = compile_to_bytecode(ir);
+        let bc = compile_to_bytecode_no_opt(ir);
         Vm::new(bc)
     }