Skip to content

Commit aafb8f0

Browse files
rbartlenskyrbartlensky
committed
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.
1 parent 2fe812c commit aafb8f0

File tree

6 files changed

+162
-51
lines changed

6 files changed

+162
-51
lines changed

‎luacompiler/src/lib/bytecode/instructions.rs

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

3434
impl HLInstr {

‎luacompiler/src/lib/bytecodegen/mod.rs

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
use bytecode::LuaBytecode;
22
use irgen::lua_ir::LuaIR;
33

4-
pub fn compile_to_bytecode(ir: LuaIR) -> LuaBytecode {
4+
pub fn compile_to_bytecode(mut ir: LuaIR) -> LuaBytecode {
5+
ir.optimize_env_lookups();
56
LuaIRToLuaBc::new(ir).compile()
67
}
78

‎luacompiler/src/lib/irgen/lua_ir.rs

+125-3
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,146 @@
1-
use bytecode::instructions::HLInstr;
1+
use bytecode::instructions::{HLInstr, Opcode};
22
use irgen::{constants_map::ConstantsMap, register_map::Lifetime};
33

44
/// Represents an IR in which all instructions are in SSA form.
55
pub struct LuaIR {
66
pub instrs: Vec<HLInstr>,
77
pub const_map: ConstantsMap,
8-
pub lifetimes: Vec<Lifetime>,
8+
pub lifetimes: Vec<Option<Lifetime>>,
9+
pub non_env_lookups: Vec<usize>,
910
}
1011

1112
impl LuaIR {
1213
pub fn new(
1314
instrs: Vec<HLInstr>,
1415
const_map: ConstantsMap,
1516
mut lifetimes: Vec<Lifetime>,
17+
non_env_lookups: Vec<usize>,
1618
) -> LuaIR {
1719
lifetimes.sort_by(|x, y| x.start_point().cmp(&y.start_point()));
1820
LuaIR {
1921
instrs,
2022
const_map,
21-
lifetimes,
23+
lifetimes: lifetimes.into_iter().map(|l| Some(l)).collect(),
24+
non_env_lookups,
2225
}
2326
}
27+
28+
/// Optimizes the IR and removes unnecessary `SetAttr` instructions.
29+
/// For example consider the following Lua program: `x = 2`
30+
/// The compiler generates:
31+
/// -------------
32+
/// LDI 1 0 0 -- load IntTable[0] into register 1, R(1) = 2
33+
/// LDS 2 0 0 -- load StringTable[0] into register 2, R(2) = "x"
34+
/// SetAttr 0 2 1 -- _ENV["x"] = 2
35+
/// -------------
36+
/// If the user does not query _ENV["x"] directly, that means we can optimize
37+
/// away the `LDS` and `SetAttr` instructions, and we can treat "x" like a local
38+
/// variable instead.
39+
pub fn optimize_env_lookups(&mut self) {
40+
let mut new_instrs = vec![];
41+
let mut new_lifetimes: Vec<Option<Lifetime>> = Vec::with_capacity(self.lifetimes.len());
42+
// push _ENV
43+
new_lifetimes.push(Some(Lifetime::new(0)));
44+
for _ in 1..self.lifetimes.len() {
45+
new_lifetimes.push(None);
46+
}
47+
let len = self.instrs.len();
48+
let mut i = 0;
49+
while i < len {
50+
let instr = self.instrs[i];
51+
if i < len - 2 && self.non_env_lookups.binary_search(&instr.1).is_ok() {
52+
// skip the attribute load
53+
i += 1;
54+
let attr = self.instrs[i];
55+
self.lifetimes[attr.1] = None;
56+
// skip the SetAttr
57+
i += 1;
58+
}
59+
new_instrs.push(instr);
60+
let regs_to_update = Self::get_register_operands(instr);
61+
for r in regs_to_update {
62+
Self::update_end_point_to(&mut new_lifetimes[r], new_instrs.len());
63+
}
64+
i += 1;
65+
}
66+
self.instrs = new_instrs;
67+
self.lifetimes = new_lifetimes;
68+
}
69+
70+
/// Gets all the register numbers that are used by the given instruction.
71+
fn get_register_operands(instr: HLInstr) -> Vec<usize> {
72+
match instr.0 {
73+
Opcode::LDI | Opcode::LDS | Opcode::LDF => vec![instr.1],
74+
Opcode::MOV => vec![instr.1, instr.2],
75+
_ => vec![instr.1, instr.2, instr.3],
76+
}
77+
}
78+
79+
/// Update the endpoint of a lifetime. If it doesn't exist, then create it.
80+
fn update_end_point_to(lifetime: &mut Option<Lifetime>, ep: usize) {
81+
match lifetime {
82+
Some(ref mut l) => l.set_end_point(ep),
83+
None => *lifetime = Some(Lifetime::new(ep - 1)),
84+
}
85+
}
86+
}
87+
88+
#[cfg(test)]
89+
mod tests {
90+
use super::*;
91+
use bytecode::instructions::Opcode;
92+
93+
#[test]
94+
fn optimize_env_lookups_generates_correct_code() {
95+
let instrs = vec![
96+
HLInstr(Opcode::LDI, 1, 0, 0),
97+
HLInstr(Opcode::LDS, 2, 0, 0),
98+
HLInstr(Opcode::SetAttr, 0, 2, 1),
99+
HLInstr(Opcode::LDI, 3, 0, 0),
100+
];
101+
let lifetimes = vec![
102+
Lifetime::with_end_point(0, 3),
103+
Lifetime::with_end_point(0, 3),
104+
Lifetime::with_end_point(1, 3),
105+
Lifetime::with_end_point(3, 4),
106+
];
107+
let non_env_lookups = vec![1];
108+
let mut ir = LuaIR::new(instrs, ConstantsMap::new(), lifetimes, non_env_lookups);
109+
ir.optimize_env_lookups();
110+
let expected_instrs = vec![HLInstr(Opcode::LDI, 1, 0, 0), HLInstr(Opcode::LDI, 3, 0, 0)];
111+
assert_eq!(ir.instrs, expected_instrs);
112+
let expected_lifetimes = vec![
113+
Some(Lifetime::with_end_point(0, 1)),
114+
Some(Lifetime::with_end_point(0, 1)),
115+
None,
116+
Some(Lifetime::with_end_point(1, 2)),
117+
];
118+
assert_eq!(ir.lifetimes, expected_lifetimes);
119+
}
120+
121+
#[test]
122+
fn optimize_env_lookups_empty_non_env_lookups() {
123+
let instrs = vec![
124+
HLInstr(Opcode::LDI, 1, 0, 0),
125+
HLInstr(Opcode::LDS, 2, 0, 0),
126+
HLInstr(Opcode::SetAttr, 0, 2, 1),
127+
HLInstr(Opcode::LDI, 3, 0, 0),
128+
];
129+
let lifetimes = vec![
130+
Lifetime::with_end_point(0, 3),
131+
Lifetime::with_end_point(0, 3),
132+
Lifetime::with_end_point(1, 3),
133+
Lifetime::with_end_point(3, 4),
134+
];
135+
let mut ir = LuaIR::new(instrs.clone(), ConstantsMap::new(), lifetimes, vec![]);
136+
ir.optimize_env_lookups();
137+
assert_eq!(ir.instrs, instrs);
138+
let expected_lifetimes = vec![
139+
Some(Lifetime::with_end_point(0, 3)),
140+
Some(Lifetime::with_end_point(0, 3)),
141+
Some(Lifetime::with_end_point(1, 3)),
142+
Some(Lifetime::with_end_point(3, 4)),
143+
];
144+
assert_eq!(ir.lifetimes, expected_lifetimes);
145+
}
24146
}

‎luacompiler/src/lib/irgen/mod.rs

+30-28
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ struct LuaToIR<'a> {
2121
pt: &'a LuaParseTree,
2222
reg_map: RegisterMap<'a>,
2323
const_map: ConstantsMap,
24+
non_env_lookups: Vec<usize>,
2425
instrs: Vec<HLInstr>,
2526
}
2627

@@ -30,6 +31,7 @@ impl<'a> LuaToIR<'a> {
3031
pt,
3132
reg_map: RegisterMap::new(),
3233
const_map: ConstantsMap::new(),
34+
non_env_lookups: vec![],
3335
instrs: vec![],
3436
}
3537
}
@@ -56,7 +58,12 @@ impl<'a> LuaToIR<'a> {
5658
}
5759
}
5860
}
59-
LuaIR::new(self.instrs, self.const_map, self.reg_map.get_lifetimes())
61+
LuaIR::new(
62+
self.instrs,
63+
self.const_map,
64+
self.reg_map.get_lifetimes(),
65+
self.non_env_lookups,
66+
)
6067
}
6168

6269
fn compile_stat(&mut self, nodes: &Vec<Node<u8>>) {
@@ -80,6 +87,7 @@ impl<'a> LuaToIR<'a> {
8087
self.reg_map.set_reg(name, value);
8188
self.instrs
8289
.push(HLInstr(Opcode::SetAttr, env_reg, attr_reg, value));
90+
self.non_env_lookups.push(value);
8391
}
8492
_ => {}
8593
}
@@ -251,26 +259,23 @@ mod tests {
251259
}
252260
// check lifetimes
253261
let expected_lifetimes = vec![
254-
Lifetime::with_end_point(0, 15),
255-
Lifetime::with_end_point(1, 2),
256-
Lifetime::with_end_point(2, 3),
257-
Lifetime::with_end_point(3, 4),
258-
Lifetime::with_end_point(4, 5),
259-
Lifetime::with_end_point(5, 6),
260-
Lifetime::with_end_point(6, 7),
261-
Lifetime::with_end_point(7, 8),
262-
Lifetime::with_end_point(8, 9),
263-
Lifetime::with_end_point(9, 10),
264-
Lifetime::with_end_point(10, 11),
265-
Lifetime::with_end_point(11, 12),
266-
Lifetime::with_end_point(12, 13),
267-
Lifetime::with_end_point(13, 14),
268-
Lifetime::with_end_point(14, 15),
262+
Some(Lifetime::with_end_point(0, 15)),
263+
Some(Lifetime::with_end_point(1, 2)),
264+
Some(Lifetime::with_end_point(2, 3)),
265+
Some(Lifetime::with_end_point(3, 4)),
266+
Some(Lifetime::with_end_point(4, 5)),
267+
Some(Lifetime::with_end_point(5, 6)),
268+
Some(Lifetime::with_end_point(6, 7)),
269+
Some(Lifetime::with_end_point(7, 8)),
270+
Some(Lifetime::with_end_point(8, 9)),
271+
Some(Lifetime::with_end_point(9, 10)),
272+
Some(Lifetime::with_end_point(10, 11)),
273+
Some(Lifetime::with_end_point(11, 12)),
274+
Some(Lifetime::with_end_point(12, 13)),
275+
Some(Lifetime::with_end_point(13, 14)),
276+
Some(Lifetime::with_end_point(14, 15)),
269277
];
270-
assert_eq!(ir.lifetimes.len(), expected_lifetimes.len());
271-
for (lhs, rhs) in ir.lifetimes.iter().zip(expected_lifetimes.iter()) {
272-
assert_eq!(lhs, rhs);
273-
}
278+
assert_eq!(ir.lifetimes, expected_lifetimes);
274279
// check constats map
275280
let expected_ints = vec![1, 2, 3];
276281
let ints = ir.const_map.get_ints();
@@ -320,15 +325,12 @@ mod tests {
320325
}
321326
// check lifetimes
322327
let expected_lifetimes = vec![
323-
Lifetime::with_end_point(0, 4),
324-
Lifetime::with_end_point(1, 4),
325-
Lifetime::with_end_point(2, 3),
326-
Lifetime::with_end_point(3, 4),
328+
Some(Lifetime::with_end_point(0, 4)),
329+
Some(Lifetime::with_end_point(1, 4)),
330+
Some(Lifetime::with_end_point(2, 3)),
331+
Some(Lifetime::with_end_point(3, 4)),
327332
];
328-
assert_eq!(ir.lifetimes.len(), expected_lifetimes.len());
329-
for (lhs, rhs) in ir.lifetimes.iter().zip(expected_lifetimes.iter()) {
330-
assert_eq!(lhs, rhs);
331-
}
333+
assert_eq!(ir.lifetimes, expected_lifetimes);
332334
// check constats map
333335
let expected_ints = vec![1];
334336
let ints = ir.const_map.get_ints();

‎luacompiler/src/lib/irgen/register_map.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ impl Lifetime {
2828
self.1
2929
}
3030

31-
fn set_end_point(&mut self, ep: usize) {
31+
pub fn set_end_point(&mut self, ep: usize) {
3232
self.1 = ep
3333
}
3434
}

‎luacompiler/tests/integration_test.rs

+3-17
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,7 @@ fn ldi_generation() {
1414
assert_eq!(bc.reg_count(), 3);
1515
assert_eq!(bc.get_int(0), 1);
1616
assert_eq!(bc.get_string(0), "x");
17-
let expected_instrs = vec![
18-
make_instr(Opcode::LDI, 1, 0, 0),
19-
make_instr(Opcode::LDS, 2, 0, 0),
20-
make_instr(Opcode::SetAttr, 0, 2, 1),
21-
];
17+
let expected_instrs = vec![make_instr(Opcode::LDI, 1, 0, 0)];
2218
assert_eq!(bc.instrs_len(), expected_instrs.len());
2319
for i in 0..expected_instrs.len() {
2420
assert_eq!(bc.get_instr(i), expected_instrs[i]);
@@ -32,11 +28,7 @@ fn ldf_generation() {
3228
assert_eq!(bc.reg_count(), 3);
3329
assert_eq!(bc.get_float(0).to_string(), "2");
3430
assert_eq!(bc.get_string(0), "x");
35-
let expected_instrs = vec![
36-
make_instr(Opcode::LDF, 1, 0, 0),
37-
make_instr(Opcode::LDS, 2, 0, 0),
38-
make_instr(Opcode::SetAttr, 0, 2, 1),
39-
];
31+
let expected_instrs = vec![make_instr(Opcode::LDF, 1, 0, 0)];
4032
assert_eq!(bc.instrs_len(), expected_instrs.len());
4133
for i in 0..expected_instrs.len() {
4234
assert_eq!(bc.get_instr(i), expected_instrs[i]);
@@ -50,11 +42,7 @@ fn lds_generation() {
5042
assert_eq!(bc.reg_count(), 3);
5143
assert_eq!(bc.get_string(0), "1.2");
5244
assert_eq!(bc.get_string(1), "x");
53-
let expected_instrs = vec![
54-
make_instr(Opcode::LDS, 1, 0, 0),
55-
make_instr(Opcode::LDS, 2, 1, 0),
56-
make_instr(Opcode::SetAttr, 0, 2, 1),
57-
];
45+
let expected_instrs = vec![make_instr(Opcode::LDS, 1, 0, 0)];
5846
assert_eq!(bc.instrs_len(), expected_instrs.len());
5947
for i in 0..expected_instrs.len() {
6048
assert_eq!(bc.get_instr(i), expected_instrs[i]);
@@ -71,8 +59,6 @@ fn assert_bytecode(opcode: Opcode, operation: &str) {
7159
make_instr(Opcode::LDI, 1, 0, 0),
7260
make_instr(Opcode::LDI, 2, 1, 0),
7361
make_instr(opcode, 3, 1, 2),
74-
make_instr(Opcode::LDS, 4, 0, 0),
75-
make_instr(Opcode::SetAttr, 0, 4, 3),
7662
];
7763
assert_eq!(bc.instrs_len(), expected_instrs.len());
7864
for i in 0..expected_instrs.len() {

0 commit comments

Comments
 (0)