From b282bdf490c55ec3d1846da0d0d4365b560294a3 Mon Sep 17 00:00:00 2001 From: Rerumu Date: Thu, 23 Jun 2022 00:27:14 -0400 Subject: [PATCH] Fix #8, as `BrTable`s have better codegen --- codegen-luajit/src/analyzer/br_table.rs | 29 ++++++++ codegen-luajit/src/analyzer/mod.rs | 1 + codegen-luajit/src/backend/manager.rs | 20 +++++- codegen-luajit/src/backend/statement.rs | 92 ++++++++++++++++++++---- codegen-luau/src/analyzer/br_table.rs | 29 ++++++++ codegen-luau/src/analyzer/mod.rs | 1 + codegen-luau/src/backend/manager.rs | 20 +++++- codegen-luau/src/backend/statement.rs | 93 +++++++++++++++++++++---- 8 files changed, 251 insertions(+), 34 deletions(-) create mode 100644 codegen-luajit/src/analyzer/br_table.rs create mode 100644 codegen-luau/src/analyzer/br_table.rs diff --git a/codegen-luajit/src/analyzer/br_table.rs b/codegen-luajit/src/analyzer/br_table.rs new file mode 100644 index 0000000..67bcb45 --- /dev/null +++ b/codegen-luajit/src/analyzer/br_table.rs @@ -0,0 +1,29 @@ +use std::collections::HashMap; + +use wasm_ast::{ + node::{BrTable, FuncData}, + visit::{Driver, Visitor}, +}; + +struct Visit { + id_map: HashMap, +} + +impl Visitor for Visit { + fn visit_br_table(&mut self, table: &BrTable) { + let id = table as *const _ as usize; + let len = self.id_map.len() + 1; + + self.id_map.insert(id, len); + } +} + +pub fn visit(ast: &FuncData) -> HashMap { + let mut visit = Visit { + id_map: HashMap::new(), + }; + + ast.accept(&mut visit); + + visit.id_map +} diff --git a/codegen-luajit/src/analyzer/mod.rs b/codegen-luajit/src/analyzer/mod.rs index ff91cb5..394a40c 100644 --- a/codegen-luajit/src/analyzer/mod.rs +++ b/codegen-luajit/src/analyzer/mod.rs @@ -1,2 +1,3 @@ +pub mod br_table; pub mod localize; pub mod operator; diff --git a/codegen-luajit/src/backend/manager.rs b/codegen-luajit/src/backend/manager.rs index 8b54817..726f6eb 100644 --- a/codegen-luajit/src/backend/manager.rs +++ b/codegen-luajit/src/backend/manager.rs @@ -1,20 +1,36 @@ use std::{ + collections::HashMap, io::{Result, Write}, ops::Range, }; -use wasm_ast::node::{CmpOp, Expression}; +use wasm_ast::node::{BrTable, CmpOp, Expression}; use crate::analyzer::operator::cmp_symbol_of; #[derive(Default)] pub struct Manager { + table_map: HashMap, label_list: Vec, num_label: usize, - pub num_param: usize, + num_param: usize, } impl Manager { + pub fn get_table_index(&self, table: &BrTable) -> usize { + let id = table as *const _ as usize; + + self.table_map[&id] + } + + pub fn set_table_map(&mut self, map: HashMap) { + self.table_map = map; + } + + pub fn set_num_param(&mut self, num: usize) { + self.num_param = num; + } + pub fn label_list(&self) -> &[usize] { &self.label_list } diff --git a/codegen-luajit/src/backend/statement.rs b/codegen-luajit/src/backend/statement.rs index 2b9828e..caf296f 100644 --- a/codegen-luajit/src/backend/statement.rs +++ b/codegen-luajit/src/backend/statement.rs @@ -9,6 +9,8 @@ use wasm_ast::node::{ SetLocal, SetTemporary, Statement, StoreAt, Terminator, }; +use crate::analyzer::br_table; + use super::manager::{ write_ascending, write_condition, write_separated, write_variable, Driver, Manager, }; @@ -28,26 +30,80 @@ impl Driver for Br { } } +fn to_ordered_table<'a>(list: &'a [Br], default: &'a Br) -> Vec<&'a Br> { + let mut data: Vec<_> = list.iter().chain(std::iter::once(default)).collect(); + + data.sort_by_key(|v| v.target); + data.dedup_by_key(|v| v.target); + + data +} + +fn write_search_layer( + range: Range, + list: &[&Br], + mng: &mut Manager, + w: &mut dyn Write, +) -> Result<()> { + if range.len() == 1 { + return list[range.start].write(mng, w); + } + + let center = range.start + range.len() / 2; + let br = list[center]; + + if range.start != center { + write!(w, "if temp < {} then ", br.target)?; + write_search_layer(range.start..center, list, mng, w)?; + write!(w, "else")?; + } + + if range.end != center + 1 { + write!(w, "if temp > {} then ", br.target)?; + write_search_layer(center + 1..range.end, list, mng, w)?; + write!(w, "else")?; + } + + write!(w, " ")?; + br.write(mng, w)?; + write!(w, "end ") +} + +fn write_table_setup(table: &BrTable, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + let id = mng.get_table_index(table); + + write!(w, "if not br_map[{id}] then ")?; + write!(w, "br_map[{id}] = (function() return {{[0] =")?; + + table + .data + .iter() + .try_for_each(|v| write!(w, "{},", v.target))?; + + write!(w, "}} end)()")?; + write!(w, "end ")?; + + write!(w, "temp = br_map[{id}][")?; + table.cond.write(mng, w)?; + write!(w, "] or {} ", table.default.target) +} + impl Driver for BrTable { fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { - write!(w, "temp = ")?; - self.cond.write(mng, w)?; - - // Our condition should be pure so we probably don't need - // to emit it in this case. if self.data.is_empty() { + // Our condition should be pure so we probably don't need + // to emit it in this case. return self.default.write(mng, w); } - for (case, dest) in self.data.iter().enumerate() { - write!(w, "if temp == {case} then ")?; - dest.write(mng, w)?; - write!(w, "else")?; - } + // `BrTable` is optimized by first mapping all indices to targets through + // a Lua table; this reduces the size of the code generated as duplicate entries + // don't need checking. Then, for speed, a binary search is done for the target + // and the appropriate jump is performed. + let list = to_ordered_table(&self.data, &self.default); - write!(w, " ")?; - self.default.write(mng, w)?; - write!(w, "end ") + write_table_setup(self, mng, w)?; + write_search_layer(0..list.len(), &list, mng, w) } } @@ -265,11 +321,17 @@ fn write_variable_list(ast: &FuncData, w: &mut dyn Write) -> Result<()> { impl Driver for FuncData { fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + let br_map = br_table::visit(self); + write_parameter_list(self, w)?; write_variable_list(self, w)?; - write!(w, "local temp ")?; - mng.num_param = self.num_param; + if !br_map.is_empty() { + write!(w, "local br_map, temp = {{}}, nil ")?; + } + + mng.set_table_map(br_map); + mng.set_num_param(self.num_param); self.code.write(mng, w)?; if self.num_result != 0 { diff --git a/codegen-luau/src/analyzer/br_table.rs b/codegen-luau/src/analyzer/br_table.rs new file mode 100644 index 0000000..67bcb45 --- /dev/null +++ b/codegen-luau/src/analyzer/br_table.rs @@ -0,0 +1,29 @@ +use std::collections::HashMap; + +use wasm_ast::{ + node::{BrTable, FuncData}, + visit::{Driver, Visitor}, +}; + +struct Visit { + id_map: HashMap, +} + +impl Visitor for Visit { + fn visit_br_table(&mut self, table: &BrTable) { + let id = table as *const _ as usize; + let len = self.id_map.len() + 1; + + self.id_map.insert(id, len); + } +} + +pub fn visit(ast: &FuncData) -> HashMap { + let mut visit = Visit { + id_map: HashMap::new(), + }; + + ast.accept(&mut visit); + + visit.id_map +} diff --git a/codegen-luau/src/analyzer/mod.rs b/codegen-luau/src/analyzer/mod.rs index ff91cb5..394a40c 100644 --- a/codegen-luau/src/analyzer/mod.rs +++ b/codegen-luau/src/analyzer/mod.rs @@ -1,2 +1,3 @@ +pub mod br_table; pub mod localize; pub mod operator; diff --git a/codegen-luau/src/backend/manager.rs b/codegen-luau/src/backend/manager.rs index 8f76b97..2df3018 100644 --- a/codegen-luau/src/backend/manager.rs +++ b/codegen-luau/src/backend/manager.rs @@ -1,9 +1,10 @@ use std::{ + collections::HashMap, io::{Result, Write}, ops::Range, }; -use wasm_ast::node::{CmpOp, Expression}; +use wasm_ast::node::{BrTable, CmpOp, Expression}; use crate::analyzer::operator::cmp_symbol_of; @@ -15,11 +16,26 @@ pub enum Label { #[derive(Default)] pub struct Manager { + table_map: HashMap, label_list: Vec