From d33e4a6b3edbffacafd2a8b8f58b21121940ed88 Mon Sep 17 00:00:00 2001 From: Rerumu <25379555+Rerumu@users.noreply.github.com> Date: Sun, 25 Jun 2023 22:16:37 -0400 Subject: [PATCH] Add naive local and temporary spills --- codegen/luajit/src/backend/expression.rs | 129 +++++++------- codegen/luajit/src/backend/manager.rs | 79 +++++++-- codegen/luajit/src/backend/statement.rs | 183 +++++++++++--------- codegen/luajit/src/translator.rs | 13 +- codegen/luau/src/analyzer/localize.rs | 2 +- codegen/luau/src/backend/expression.rs | 129 +++++++------- codegen/luau/src/backend/manager.rs | 78 +++++++-- codegen/luau/src/backend/statement.rs | 205 ++++++++++++----------- codegen/luau/src/translator.rs | 12 +- wasm-ast/src/factory.rs | 30 ++-- wasm-ast/src/module.rs | 13 +- wasm-ast/src/node.rs | 86 ++++++---- wasm-ast/src/stack.rs | 20 ++- wasm-ast/src/visit.rs | 12 +- 14 files changed, 589 insertions(+), 402 deletions(-) diff --git a/codegen/luajit/src/backend/expression.rs b/codegen/luajit/src/backend/expression.rs index 6a6a65f..c960afe 100644 --- a/codegen/luajit/src/backend/expression.rs +++ b/codegen/luajit/src/backend/expression.rs @@ -4,13 +4,12 @@ use std::{ }; use wasm_ast::node::{ - BinOp, CmpOp, Expression, GetGlobal, GetLocal, GetTemporary, LoadAt, MemorySize, Select, UnOp, - Value, + BinOp, CmpOp, Expression, GetGlobal, LoadAt, Local, MemorySize, Select, Temporary, UnOp, Value, }; use crate::analyzer::as_symbol::AsSymbol; -use super::manager::{write_separated, DriverNoContext}; +use super::manager::{write_separated, Driver, Manager}; macro_rules! impl_write_number { ($name:tt, $numeric:ty) => { @@ -26,43 +25,55 @@ macro_rules! impl_write_number { }; } -impl DriverNoContext for Select { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for Select { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { write!(w, "(")?; - Condition(self.condition()).write(w)?; + Condition(self.condition()).write(mng, w)?; write!(w, " and ")?; - self.on_true().write(w)?; + self.on_true().write(mng, w)?; write!(w, " or ")?; - self.on_false().write(w)?; + self.on_false().write(mng, w)?; write!(w, ")") } } -impl DriverNoContext for GetTemporary { - fn write(&self, w: &mut dyn Write) -> Result<()> { - write!(w, "reg_{}", self.var()) +impl Driver for Temporary { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + let var = self.var(); + + if let Some(var) = var.checked_sub(mng.num_temp()) { + write!(w, "reg_spill[{}]", var + 1) + } else { + write!(w, "reg_{var}") + } } } -impl DriverNoContext for GetLocal { - fn write(&self, w: &mut dyn Write) -> Result<()> { - write!(w, "loc_{}", self.var()) +impl Driver for Local { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + let var = self.var(); + + if let Some(var) = var.checked_sub(mng.num_local()) { + write!(w, "loc_spill[{}]", var + 1) + } else { + write!(w, "loc_{var}") + } } } -impl DriverNoContext for GetGlobal { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for GetGlobal { + fn write(&self, _mng: &mut Manager, w: &mut dyn Write) -> Result<()> { write!(w, "GLOBAL_LIST[{}].value", self.var()) } } -impl DriverNoContext for LoadAt { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for LoadAt { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { let name = self.load_type().as_name(); let memory = self.memory(); write!(w, "load_{name}(memory_at_{memory}, ")?; - self.pointer().write(w)?; + self.pointer().write(mng, w)?; if self.offset() != 0 { write!(w, " + {}", self.offset())?; @@ -72,8 +83,8 @@ impl DriverNoContext for LoadAt { } } -impl DriverNoContext for MemorySize { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for MemorySize { + fn write(&self, _mng: &mut Manager, w: &mut dyn Write) -> Result<()> { write!(w, "memory_at_{}.min", self.memory()) } } @@ -81,8 +92,8 @@ impl DriverNoContext for MemorySize { impl_write_number!(write_f32, f32); impl_write_number!(write_f64, f64); -impl DriverNoContext for Value { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for Value { + fn write(&self, _mng: &mut Manager, w: &mut dyn Write) -> Result<()> { match self { Self::I32(i) => write!(w, "{i}"), Self::I64(i) => write!(w, "{i}LL"), @@ -92,97 +103,97 @@ impl DriverNoContext for Value { } } -impl DriverNoContext for UnOp { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for UnOp { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { let (a, b) = self.op_type().as_name(); write!(w, "{a}_{b}(")?; - self.rhs().write(w)?; + self.rhs().write(mng, w)?; write!(w, ")") } } -impl DriverNoContext for BinOp { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for BinOp { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { if let Some(symbol) = self.op_type().as_symbol() { write!(w, "(")?; - self.lhs().write(w)?; + self.lhs().write(mng, w)?; write!(w, " {symbol} ")?; } else { let (head, tail) = self.op_type().as_name(); write!(w, "{head}_{tail}(")?; - self.lhs().write(w)?; + self.lhs().write(mng, w)?; write!(w, ", ")?; } - self.rhs().write(w)?; + self.rhs().write(mng, w)?; write!(w, ")") } } struct CmpOpBoolean<'a>(&'a CmpOp); -impl DriverNoContext for CmpOpBoolean<'_> { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for CmpOpBoolean<'_> { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { let cmp = self.0; if let Some(symbol) = cmp.op_type().as_symbol() { - cmp.lhs().write(w)?; + cmp.lhs().write(mng, w)?; write!(w, " {symbol} ")?; - cmp.rhs().write(w) + cmp.rhs().write(mng, w) } else { let (head, tail) = cmp.op_type().as_name(); write!(w, "{head}_{tail}(")?; - cmp.lhs().write(w)?; + cmp.lhs().write(mng, w)?; write!(w, ", ")?; - cmp.rhs().write(w)?; + cmp.rhs().write(mng, w)?; write!(w, ")") } } } -impl DriverNoContext for CmpOp { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for CmpOp { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { write!(w, "(")?; - CmpOpBoolean(self).write(w)?; + CmpOpBoolean(self).write(mng, w)?; write!(w, " and 1 or 0)") } } pub struct Condition<'a>(pub &'a Expression); -impl DriverNoContext for Condition<'_> { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for Condition<'_> { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { if let Expression::CmpOp(node) = self.0 { - CmpOpBoolean(node).write(w) + CmpOpBoolean(node).write(mng, w) } else { - self.0.write(w)?; + self.0.write(mng, w)?; write!(w, " ~= 0") } } } -impl DriverNoContext for Expression { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for Expression { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { match self { - Self::Select(e) => e.write(w), - Self::GetTemporary(e) => e.write(w), - Self::GetLocal(e) => e.write(w), - Self::GetGlobal(e) => e.write(w), - Self::LoadAt(e) => e.write(w), - Self::MemorySize(e) => e.write(w), - Self::Value(e) => e.write(w), - Self::UnOp(e) => e.write(w), - Self::BinOp(e) => e.write(w), - Self::CmpOp(e) => e.write(w), + Self::Select(e) => e.write(mng, w), + Self::GetTemporary(e) => e.write(mng, w), + Self::GetLocal(e) => e.write(mng, w), + Self::GetGlobal(e) => e.write(mng, w), + Self::LoadAt(e) => e.write(mng, w), + Self::MemorySize(e) => e.write(mng, w), + Self::Value(e) => e.write(mng, w), + Self::UnOp(e) => e.write(mng, w), + Self::BinOp(e) => e.write(mng, w), + Self::CmpOp(e) => e.write(mng, w), } } } -impl DriverNoContext for &[Expression] { - fn write(&self, w: &mut dyn Write) -> Result<()> { - write_separated(self.iter(), |e, w| e.write(w), w) +impl Driver for &[Expression] { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + write_separated(self.iter(), |e, w| e.write(mng, w), w) } } diff --git a/codegen/luajit/src/backend/manager.rs b/codegen/luajit/src/backend/manager.rs index ee31d1e..04af04e 100644 --- a/codegen/luajit/src/backend/manager.rs +++ b/codegen/luajit/src/backend/manager.rs @@ -1,10 +1,11 @@ use std::{ collections::HashMap, io::{Result, Write}, - ops::Range, }; -use wasm_ast::node::BrTable; +use wasm_ast::node::{BrTable, FuncData}; + +use crate::analyzer::{br_table, localize}; #[macro_export] macro_rules! indentation { @@ -31,23 +32,81 @@ macro_rules! line { }}; } -#[derive(Default)] +fn get_pinned_registers( + upvalues: usize, + params: usize, + locals: usize, + temporaries: usize, +) -> (usize, usize) { + const MAX_LOCAL_COUNT: usize = 170; + + let available = MAX_LOCAL_COUNT + .saturating_sub(upvalues) + .saturating_sub(params); + + let locals = available.min(locals); + let temporaries = available.saturating_sub(locals).min(temporaries); + + (params + locals, temporaries) +} + pub struct Manager { table_map: HashMap, - label_list: Vec, + num_local: usize, + num_temp: usize, num_label: usize, + label_list: Vec, indentation: usize, } impl Manager { + pub fn empty() -> Self { + Self { + table_map: HashMap::new(), + num_local: 0, + num_temp: usize::MAX, + num_label: 0, + label_list: Vec::new(), + indentation: 0, + } + } + + pub fn function(ast: &FuncData) -> Self { + let (upvalues, memories) = localize::visit(ast); + let table_map = br_table::visit(ast); + let (num_local, num_temp) = get_pinned_registers( + upvalues.len() + memories.len(), + ast.num_param(), + ast.local_data().len(), + ast.num_stack(), + ); + + Self { + table_map, + num_local, + num_temp, + num_label: 0, + label_list: Vec::new(), + indentation: 0, + } + } + 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 has_table(&self) -> bool { + !self.table_map.is_empty() + } + + pub const fn num_local(&self) -> usize { + self.num_local + } + + pub const fn num_temp(&self) -> usize { + self.num_temp } pub fn label_list(&self) -> &[usize] { @@ -82,10 +141,6 @@ pub trait Driver { fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()>; } -pub trait DriverNoContext { - fn write(&self, w: &mut dyn Write) -> Result<()>; -} - pub fn write_separated(mut iter: I, mut func: M, w: &mut dyn Write) -> Result<()> where M: FnMut(T, &mut dyn Write) -> Result<()>, @@ -101,7 +156,3 @@ where func(v, w) }) } - -pub fn write_ascending(prefix: &str, range: Range, w: &mut dyn Write) -> Result<()> { - write_separated(range, |i, w| write!(w, "{prefix}_{i}"), w) -} diff --git a/codegen/luajit/src/backend/statement.rs b/codegen/luajit/src/backend/statement.rs index 86e3328..5256f02 100644 --- a/codegen/luajit/src/backend/statement.rs +++ b/codegen/luajit/src/backend/statement.rs @@ -5,26 +5,32 @@ use std::{ use wasm_ast::node::{ Block, Br, BrIf, BrTable, Call, CallIndirect, FuncData, If, LabelType, MemoryCopy, MemoryFill, - MemoryGrow, SetGlobal, SetLocal, SetTemporary, Statement, StoreAt, Terminator, + MemoryGrow, ResultList, SetGlobal, SetLocal, SetTemporary, Statement, StoreAt, Terminator, }; use wasmparser::ValType; -use crate::{analyzer::br_table, indentation, indented, line}; +use crate::{backend::manager::write_separated, indentation, indented, line}; use super::{ expression::Condition, - manager::{write_ascending, write_separated, Driver, DriverNoContext, Manager}, + manager::{Driver, Manager}, }; +impl Driver for ResultList { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + write_separated(self.iter(), |t, w| t.write(mng, w), w) + } +} + impl Driver for Br { fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { let level = *mng.label_list().iter().nth_back(self.target()).unwrap(); if !self.align().is_aligned() { indentation!(mng, w)?; - write_ascending("reg", self.align().new_range(), w)?; + self.align().new_range().write(mng, w)?; write!(w, " = ")?; - write_ascending("reg", self.align().old_range(), w)?; + self.align().old_range().write(mng, w)?; writeln!(w)?; } @@ -97,7 +103,7 @@ fn write_table_setup(table: &BrTable, mng: &mut Manager, w: &mut dyn Write) -> R line!(mng, w, "end")?; indented!(mng, w, "temp = br_map[{id}][")?; - table.condition().write(w)?; + table.condition().write(mng, w)?; writeln!(w, "] or {}", table.default().target()) } @@ -174,7 +180,7 @@ impl Driver for Block { impl Driver for BrIf { fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { indented!(mng, w, "if ")?; - Condition(self.condition()).write(w)?; + Condition(self.condition()).write(mng, w)?; writeln!(w, " then")?; mng.indent(); self.target().write(mng, w)?; @@ -186,7 +192,7 @@ impl Driver for BrIf { impl Driver for If { fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { indented!(mng, w, "if ")?; - Condition(self.condition()).write(w)?; + Condition(self.condition()).write(mng, w)?; writeln!(w, " then")?; mng.indent(); @@ -204,120 +210,119 @@ impl Driver for If { } } -fn write_call_store(result: Range, w: &mut dyn Write) -> Result<()> { - if result.is_empty() { - return Ok(()); - } - - write_ascending("reg", result, w)?; - write!(w, " = ") -} - -impl DriverNoContext for Call { - fn write(&self, w: &mut dyn Write) -> Result<()> { - write_call_store(self.result(), w)?; +impl Driver for Call { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + if !self.result_list().is_empty() { + self.result_list().write(mng, w)?; + write!(w, " = ")?; + } write!(w, "FUNC_LIST[{}](", self.function())?; - self.param_list().write(w)?; + self.param_list().write(mng, w)?; write!(w, ")") } } -impl DriverNoContext for CallIndirect { - fn write(&self, w: &mut dyn Write) -> Result<()> { - write_call_store(self.result(), w)?; +impl Driver for CallIndirect { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + if !self.result_list().is_empty() { + self.result_list().write(mng, w)?; + write!(w, " = ")?; + } write!(w, "TABLE_LIST[{}].data[", self.table())?; - self.index().write(w)?; + self.index().write(mng, w)?; write!(w, "](")?; - self.param_list().write(w)?; + self.param_list().write(mng, w)?; write!(w, ")") } } -impl DriverNoContext for SetTemporary { - fn write(&self, w: &mut dyn Write) -> Result<()> { - write!(w, "reg_{} = ", self.var())?; - self.value().write(w) +impl Driver for SetTemporary { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + self.var().write(mng, w)?; + write!(w, " = ")?; + self.value().write(mng, w) } } -impl DriverNoContext for SetLocal { - fn write(&self, w: &mut dyn Write) -> Result<()> { - write!(w, "loc_{} = ", self.var())?; - self.value().write(w) +impl Driver for SetLocal { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + self.var().write(mng, w)?; + write!(w, " = ")?; + self.value().write(mng, w) } } -impl DriverNoContext for SetGlobal { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for SetGlobal { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { write!(w, "GLOBAL_LIST[{}].value = ", self.var())?; - self.value().write(w) + self.value().write(mng, w) } } -impl DriverNoContext for StoreAt { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for StoreAt { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { let name = self.store_type().as_name(); let memory = self.memory(); write!(w, "store_{name}(memory_at_{memory}, ")?; - self.pointer().write(w)?; + self.pointer().write(mng, w)?; if self.offset() != 0 { write!(w, " + {}", self.offset())?; } write!(w, ", ")?; - self.value().write(w)?; + self.value().write(mng, w)?; write!(w, ")") } } -impl DriverNoContext for MemoryGrow { - fn write(&self, w: &mut dyn Write) -> Result<()> { - let result = self.result(); +impl Driver for MemoryGrow { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { let memory = self.memory(); - write!(w, "reg_{result} = rt.allocator.grow(memory_at_{memory}, ")?; - self.size().write(w)?; + self.result().write(mng, w)?; + write!(w, " = rt.allocator.grow(memory_at_{memory}, ")?; + self.size().write(mng, w)?; write!(w, ")") } } -impl DriverNoContext for MemoryCopy { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for MemoryCopy { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { let memory_1 = self.destination().memory(); let memory_2 = self.source().memory(); write!(w, "rt.store.copy(memory_at_{memory_1}, ")?; - self.destination().pointer().write(w)?; + self.destination().pointer().write(mng, w)?; write!(w, ", memory_at_{memory_2}, ")?; - self.source().pointer().write(w)?; + self.source().pointer().write(mng, w)?; write!(w, ", ")?; - self.size().write(w)?; + self.size().write(mng, w)?; write!(w, ")") } } -impl DriverNoContext for MemoryFill { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for MemoryFill { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { let memory = self.destination().memory(); write!(w, "rt.store.fill(memory_at_{memory}, ")?; - self.destination().pointer().write(w)?; + self.destination().pointer().write(mng, w)?; write!(w, ", ")?; - self.size().write(w)?; + self.size().write(mng, w)?; write!(w, ", ")?; - self.value().write(w)?; + self.value().write(mng, w)?; write!(w, ")") } } -fn write_stat(stat: &dyn DriverNoContext, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { +fn write_stat(stat: &dyn Driver, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { indentation!(mng, w)?; - stat.write(w)?; + stat.write(mng, w)?; writeln!(w) } @@ -342,30 +347,51 @@ impl Driver for Statement { fn write_parameter_list(ast: &FuncData, w: &mut dyn Write) -> Result<()> { write!(w, "function(")?; - write_ascending("loc", 0..ast.num_param(), w)?; + write_separated(0..ast.num_param(), |i, w| write!(w, "loc_{i}"), w)?; writeln!(w, ")") } +const fn type_to_zero(typ: ValType) -> &'static str { + match typ { + ValType::F32 | ValType::F64 => "0.0", + ValType::I64 => "0LL", + _ => "0", + } +} + fn write_variable_list(ast: &FuncData, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { - let mut total = ast.num_param(); + let mut locals = ast.local_data().iter().copied(); + let num_local = mng.num_local() - ast.num_param(); - for data in ast.local_data().iter().filter(|v| v.0 != 0) { - let range = total..total + usize::try_from(data.0).unwrap(); - let typed = if data.1 == ValType::I64 { "0LL" } else { "0" }.as_bytes(); + for (i, typ) in locals.by_ref().enumerate().take(num_local) { + let index = ast.num_param() + i; + let zero = type_to_zero(typ); - total = range.end; - - indented!(mng, w, "local ")?; - write_ascending("loc", range.clone(), w)?; - write!(w, " = ")?; - write_separated(range, |_, w| w.write_all(typed), w)?; - writeln!(w)?; + line!(mng, w, "local loc_{index} = {zero}")?; } - if ast.num_stack() != 0 { - indented!(mng, w, "local ")?; - write_ascending("reg", 0..ast.num_stack(), w)?; - writeln!(w)?; + if locals.len() != 0 { + indented!(mng, w, "local loc_spill = {{ ")?; + + for typ in locals { + let zero = type_to_zero(typ); + + write!(w, "{zero}, ")?; + } + + writeln!(w, "}}")?; + } + + let mut temporaries = 0..ast.num_stack(); + + for i in temporaries.by_ref().take(mng.num_temp()) { + line!(mng, w, "local reg_{i}")?; + } + + if !temporaries.is_empty() { + let len = temporaries.len(); + + line!(mng, w, "local reg_spill = table.create({len})")?; } Ok(()) @@ -373,23 +399,22 @@ fn write_variable_list(ast: &FuncData, mng: &mut Manager, w: &mut dyn Write) -> impl Driver for FuncData { fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { - let br_map = br_table::visit(self); - mng.indent(); write_parameter_list(self, w)?; write_variable_list(self, mng, w)?; - if !br_map.is_empty() { + if mng.has_table() { line!(mng, w, "local br_map, temp = {{}}, nil")?; } - mng.set_table_map(br_map); self.code().write(mng, w)?; if self.num_result() != 0 { indented!(mng, w, "return ")?; - write_ascending("reg", 0..self.num_result(), w)?; + + ResultList::new(0, self.num_result()).write(mng, w)?; + writeln!(w)?; } diff --git a/codegen/luajit/src/translator.rs b/codegen/luajit/src/translator.rs index 3422835..496f42a 100644 --- a/codegen/luajit/src/translator.rs +++ b/codegen/luajit/src/translator.rs @@ -15,7 +15,7 @@ use wasmparser::{ use crate::{ analyzer::localize, - backend::manager::{Driver, DriverNoContext, Manager}, + backend::manager::{Driver, Manager}, }; trait AsIEName { @@ -51,7 +51,7 @@ fn write_constant(init: &ConstExpr, type_info: &TypeInfo, w: &mut dyn Write) -> let func = Factory::from_type_info(type_info).create_anonymous(&code); if let Some(Statement::SetTemporary(stat)) = func.code().code().last() { - stat.value().write(w) + stat.value().write(&mut Manager::empty(), w) } else { writeln!(w, r#"error("Valueless constant")"#) } @@ -178,6 +178,7 @@ fn write_element_list(list: &[Element], type_info: &TypeInfo, w: &mut dyn Write) } } } + writeln!(w, " }}")?; writeln!(w, "\t\ttable.move(data, 1, #data, offset, target)")?; writeln!(w, "\tend")?; @@ -270,7 +271,7 @@ fn write_func_list(wasm: &Module, func_list: &[FuncData], w: &mut dyn Write) -> write_func_start(wasm, index, w)?; - v.write(&mut Manager::default(), w) + v.write(&mut Manager::function(v), w) }) } @@ -309,9 +310,9 @@ fn write_module_start( /// # Errors /// Returns `Err` if writing to `Write` failed. pub fn from_inst_list(code: &[Operator], type_info: &TypeInfo, w: &mut dyn Write) -> Result<()> { - Factory::from_type_info(type_info) - .create_anonymous(code) - .write(&mut Manager::default(), w) + let ast = Factory::from_type_info(type_info).create_anonymous(code); + + ast.write(&mut Manager::function(&ast), w) } /// # Errors diff --git a/codegen/luau/src/analyzer/localize.rs b/codegen/luau/src/analyzer/localize.rs index 6755072..4f873de 100644 --- a/codegen/luau/src/analyzer/localize.rs +++ b/codegen/luau/src/analyzer/localize.rs @@ -92,7 +92,7 @@ pub fn visit(ast: &FuncData) -> (BTreeSet<(&'static str, &'static str)>, BTreeSe memory_set: BTreeSet::new(), }; - if ast.local_data().iter().any(|v| v.1 == ValType::I64) { + if ast.local_data().iter().any(|&v| v == ValType::I64) { visit.local_set.insert(("i64", "ZERO")); } diff --git a/codegen/luau/src/backend/expression.rs b/codegen/luau/src/backend/expression.rs index 14e4513..956ab45 100644 --- a/codegen/luau/src/backend/expression.rs +++ b/codegen/luau/src/backend/expression.rs @@ -4,13 +4,12 @@ use std::{ }; use wasm_ast::node::{ - BinOp, CmpOp, Expression, GetGlobal, GetLocal, GetTemporary, LoadAt, MemorySize, Select, UnOp, - Value, + BinOp, CmpOp, Expression, GetGlobal, LoadAt, Local, MemorySize, Select, Temporary, UnOp, Value, }; use crate::analyzer::as_symbol::AsSymbol; -use super::manager::{write_separated, DriverNoContext}; +use super::manager::{write_separated, Driver, Manager}; macro_rules! impl_write_number { ($name:tt, $numeric:ty) => { @@ -26,43 +25,55 @@ macro_rules! impl_write_number { }; } -impl DriverNoContext for Select { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for Select { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { write!(w, "(if ")?; - Condition(self.condition()).write(w)?; + Condition(self.condition()).write(mng, w)?; write!(w, " then ")?; - self.on_true().write(w)?; + self.on_true().write(mng, w)?; write!(w, " else ")?; - self.on_false().write(w)?; + self.on_false().write(mng, w)?; write!(w, ")") } } -impl DriverNoContext for GetTemporary { - fn write(&self, w: &mut dyn Write) -> Result<()> { - write!(w, "reg_{}", self.var()) +impl Driver for Temporary { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + let var = self.var(); + + if let Some(var) = var.checked_sub(mng.num_temp()) { + write!(w, "reg_spill[{}]", var + 1) + } else { + write!(w, "reg_{var}") + } } } -impl DriverNoContext for GetLocal { - fn write(&self, w: &mut dyn Write) -> Result<()> { - write!(w, "loc_{}", self.var()) +impl Driver for Local { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + let var = self.var(); + + if let Some(var) = var.checked_sub(mng.num_local()) { + write!(w, "loc_spill[{}]", var + 1) + } else { + write!(w, "loc_{var}") + } } } -impl DriverNoContext for GetGlobal { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for GetGlobal { + fn write(&self, _mng: &mut Manager, w: &mut dyn Write) -> Result<()> { write!(w, "GLOBAL_LIST[{}].value", self.var()) } } -impl DriverNoContext for LoadAt { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for LoadAt { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { let name = self.load_type().as_name(); let memory = self.memory(); write!(w, "load_{name}(memory_at_{memory}, ")?; - self.pointer().write(w)?; + self.pointer().write(mng, w)?; if self.offset() != 0 { write!(w, " + {}", self.offset())?; @@ -72,8 +83,8 @@ impl DriverNoContext for LoadAt { } } -impl DriverNoContext for MemorySize { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for MemorySize { + fn write(&self, _mng: &mut Manager, w: &mut dyn Write) -> Result<()> { write!(w, "memory_at_{}.min", self.memory()) } } @@ -101,8 +112,8 @@ fn write_i64(number: i64, w: &mut dyn Write) -> Result<()> { impl_write_number!(write_f32, f32); impl_write_number!(write_f64, f64); -impl DriverNoContext for Value { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for Value { + fn write(&self, _mng: &mut Manager, w: &mut dyn Write) -> Result<()> { match self { Self::I32(i) => write_i32(*i, w), Self::I64(i) => write_i64(*i, w), @@ -112,97 +123,97 @@ impl DriverNoContext for Value { } } -impl DriverNoContext for UnOp { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for UnOp { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { let (a, b) = self.op_type().as_name(); write!(w, "{a}_{b}(")?; - self.rhs().write(w)?; + self.rhs().write(mng, w)?; write!(w, ")") } } -impl DriverNoContext for BinOp { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for BinOp { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { if let Some(symbol) = self.op_type().as_symbol() { write!(w, "(")?; - self.lhs().write(w)?; + self.lhs().write(mng, w)?; write!(w, " {symbol} ")?; } else { let (head, tail) = self.op_type().as_name(); write!(w, "{head}_{tail}(")?; - self.lhs().write(w)?; + self.lhs().write(mng, w)?; write!(w, ", ")?; } - self.rhs().write(w)?; + self.rhs().write(mng, w)?; write!(w, ")") } } struct CmpOpBoolean<'a>(&'a CmpOp); -impl DriverNoContext for CmpOpBoolean<'_> { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for CmpOpBoolean<'_> { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { let cmp = self.0; if let Some(symbol) = cmp.op_type().as_symbol() { - cmp.lhs().write(w)?; + cmp.lhs().write(mng, w)?; write!(w, " {symbol} ")?; - cmp.rhs().write(w) + cmp.rhs().write(mng, w) } else { let (head, tail) = cmp.op_type().as_name(); write!(w, "{head}_{tail}(")?; - cmp.lhs().write(w)?; + cmp.lhs().write(mng, w)?; write!(w, ", ")?; - cmp.rhs().write(w)?; + cmp.rhs().write(mng, w)?; write!(w, ")") } } } -impl DriverNoContext for CmpOp { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for CmpOp { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { write!(w, "(if ")?; - CmpOpBoolean(self).write(w)?; + CmpOpBoolean(self).write(mng, w)?; write!(w, " then 1 else 0)") } } pub struct Condition<'a>(pub &'a Expression); -impl DriverNoContext for Condition<'_> { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for Condition<'_> { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { if let Expression::CmpOp(node) = self.0 { - CmpOpBoolean(node).write(w) + CmpOpBoolean(node).write(mng, w) } else { - self.0.write(w)?; + self.0.write(mng, w)?; write!(w, " ~= 0") } } } -impl DriverNoContext for Expression { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for Expression { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { match self { - Self::Select(e) => e.write(w), - Self::GetTemporary(e) => e.write(w), - Self::GetLocal(e) => e.write(w), - Self::GetGlobal(e) => e.write(w), - Self::LoadAt(e) => e.write(w), - Self::MemorySize(e) => e.write(w), - Self::Value(e) => e.write(w), - Self::UnOp(e) => e.write(w), - Self::BinOp(e) => e.write(w), - Self::CmpOp(e) => e.write(w), + Self::Select(e) => e.write(mng, w), + Self::GetTemporary(e) => e.write(mng, w), + Self::GetLocal(e) => e.write(mng, w), + Self::GetGlobal(e) => e.write(mng, w), + Self::LoadAt(e) => e.write(mng, w), + Self::MemorySize(e) => e.write(mng, w), + Self::Value(e) => e.write(mng, w), + Self::UnOp(e) => e.write(mng, w), + Self::BinOp(e) => e.write(mng, w), + Self::CmpOp(e) => e.write(mng, w), } } } -impl DriverNoContext for &[Expression] { - fn write(&self, w: &mut dyn Write) -> Result<()> { - write_separated(self.iter(), |e, w| e.write(w), w) +impl Driver for &[Expression] { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + write_separated(self.iter(), |e, w| e.write(mng, w), w) } } diff --git a/codegen/luau/src/backend/manager.rs b/codegen/luau/src/backend/manager.rs index ed57b5c..b237cf2 100644 --- a/codegen/luau/src/backend/manager.rs +++ b/codegen/luau/src/backend/manager.rs @@ -1,10 +1,11 @@ use std::{ collections::HashMap, io::{Result, Write}, - ops::Range, }; -use wasm_ast::node::{BrTable, LabelType}; +use wasm_ast::node::{BrTable, FuncData, LabelType}; + +use crate::analyzer::{br_target, localize}; #[macro_export] macro_rules! indentation { @@ -31,30 +32,87 @@ macro_rules! line { }}; } -#[derive(Default)] +fn get_pinned_registers( + upvalues: usize, + params: usize, + locals: usize, + temporaries: usize, +) -> (usize, usize) { + const MAX_LOCAL_COUNT: usize = 170; + + let available = MAX_LOCAL_COUNT + .saturating_sub(upvalues) + .saturating_sub(params); + + let locals = available.min(locals); + let temporaries = available.saturating_sub(locals).min(temporaries); + + (params + locals, temporaries) +} + pub struct Manager { table_map: HashMap, has_branch: bool, + num_local: usize, + num_temp: usize, label_list: Vec>, indentation: usize, } impl Manager { + pub fn empty() -> Self { + Self { + table_map: HashMap::new(), + has_branch: false, + num_local: 0, + num_temp: usize::MAX, + label_list: Vec::new(), + indentation: 0, + } + } + + pub fn function(ast: &FuncData) -> Self { + let (upvalues, memories) = localize::visit(ast); + let (table_map, has_branch) = br_target::visit(ast); + let (num_local, num_temp) = get_pinned_registers( + upvalues.len() + memories.len(), + ast.num_param(), + ast.local_data().len(), + ast.num_stack(), + ); + + Self { + table_map, + has_branch, + num_local, + num_temp, + label_list: Vec::new(), + indentation: 0, + } + } + pub fn get_table_index(&self, table: &BrTable) -> usize { let id = table as *const _ as usize; self.table_map[&id] } - pub fn set_branch_information(&mut self, table_map: HashMap, has_branch: bool) { - self.table_map = table_map; - self.has_branch = has_branch; + pub fn has_table(&self) -> bool { + !self.table_map.is_empty() } pub const fn has_branch(&self) -> bool { self.has_branch } + pub const fn num_local(&self) -> usize { + self.num_local + } + + pub const fn num_temp(&self) -> usize { + self.num_temp + } + pub fn label_list(&self) -> &[Option] { &self.label_list } @@ -84,10 +142,6 @@ pub trait Driver { fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()>; } -pub trait DriverNoContext { - fn write(&self, w: &mut dyn Write) -> Result<()>; -} - pub fn write_separated(mut iter: I, mut func: M, w: &mut dyn Write) -> Result<()> where M: FnMut(T, &mut dyn Write) -> Result<()>, @@ -103,7 +157,3 @@ where func(v, w) }) } - -pub fn write_ascending(prefix: &str, range: Range, w: &mut dyn Write) -> Result<()> { - write_separated(range, |i, w| write!(w, "{prefix}_{i}"), w) -} diff --git a/codegen/luau/src/backend/statement.rs b/codegen/luau/src/backend/statement.rs index b25979c..02a30d1 100644 --- a/codegen/luau/src/backend/statement.rs +++ b/codegen/luau/src/backend/statement.rs @@ -5,24 +5,30 @@ use std::{ use wasm_ast::node::{ Block, Br, BrIf, BrTable, Call, CallIndirect, FuncData, If, LabelType, MemoryCopy, MemoryFill, - MemoryGrow, SetGlobal, SetLocal, SetTemporary, Statement, StoreAt, Terminator, + MemoryGrow, ResultList, SetGlobal, SetLocal, SetTemporary, Statement, StoreAt, Terminator, }; use wasmparser::ValType; -use crate::{analyzer::br_target, indentation, indented, line}; +use crate::{backend::manager::write_separated, indentation, indented, line}; use super::{ expression::Condition, - manager::{write_ascending, write_separated, Driver, DriverNoContext, Manager}, + manager::{Driver, Manager}, }; +impl Driver for ResultList { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + write_separated(self.iter(), |t, w| t.write(mng, w), w) + } +} + impl Driver for Br { fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { if !self.align().is_aligned() { indentation!(mng, w)?; - write_ascending("reg", self.align().new_range(), w)?; + self.align().new_range().write(mng, w)?; write!(w, " = ")?; - write_ascending("reg", self.align().old_range(), w)?; + self.align().old_range().write(mng, w)?; writeln!(w)?; } @@ -106,7 +112,7 @@ fn write_table_setup(table: &BrTable, mng: &mut Manager, w: &mut dyn Write) -> R line!(mng, w, "end")?; indented!(mng, w, "temp = br_map[{id}][")?; - table.condition().write(w)?; + table.condition().write(mng, w)?; writeln!(w, "] or {}", table.default().target()) } @@ -193,7 +199,7 @@ impl Driver for Block { impl Driver for BrIf { fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { indented!(mng, w, "if ")?; - Condition(self.condition()).write(w)?; + Condition(self.condition()).write(mng, w)?; writeln!(w, " then")?; mng.indent(); self.target().write(mng, w)?; @@ -205,7 +211,7 @@ impl Driver for BrIf { impl Driver for If { fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { indented!(mng, w, "if ")?; - Condition(self.condition()).write(w)?; + Condition(self.condition()).write(mng, w)?; writeln!(w, " then")?; mng.indent(); @@ -223,120 +229,119 @@ impl Driver for If { } } -fn write_call_store(result: Range, w: &mut dyn Write) -> Result<()> { - if result.is_empty() { - return Ok(()); - } - - write_ascending("reg", result, w)?; - write!(w, " = ") -} - -impl DriverNoContext for Call { - fn write(&self, w: &mut dyn Write) -> Result<()> { - write_call_store(self.result(), w)?; +impl Driver for Call { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + if !self.result_list().is_empty() { + self.result_list().write(mng, w)?; + write!(w, " = ")?; + } write!(w, "FUNC_LIST[{}](", self.function())?; - self.param_list().write(w)?; + self.param_list().write(mng, w)?; write!(w, ")") } } -impl DriverNoContext for CallIndirect { - fn write(&self, w: &mut dyn Write) -> Result<()> { - write_call_store(self.result(), w)?; +impl Driver for CallIndirect { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + if !self.result_list().is_empty() { + self.result_list().write(mng, w)?; + write!(w, " = ")?; + } write!(w, "TABLE_LIST[{}].data[", self.table())?; - self.index().write(w)?; + self.index().write(mng, w)?; write!(w, "](")?; - self.param_list().write(w)?; + self.param_list().write(mng, w)?; write!(w, ")") } } -impl DriverNoContext for SetTemporary { - fn write(&self, w: &mut dyn Write) -> Result<()> { - write!(w, "reg_{} = ", self.var())?; - self.value().write(w) +impl Driver for SetTemporary { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + self.var().write(mng, w)?; + write!(w, " = ")?; + self.value().write(mng, w) } } -impl DriverNoContext for SetLocal { - fn write(&self, w: &mut dyn Write) -> Result<()> { - write!(w, "loc_{} = ", self.var())?; - self.value().write(w) +impl Driver for SetLocal { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { + self.var().write(mng, w)?; + write!(w, " = ")?; + self.value().write(mng, w) } } -impl DriverNoContext for SetGlobal { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for SetGlobal { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { write!(w, "GLOBAL_LIST[{}].value = ", self.var())?; - self.value().write(w) + self.value().write(mng, w) } } -impl DriverNoContext for StoreAt { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for StoreAt { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { let name = self.store_type().as_name(); let memory = self.memory(); write!(w, "store_{name}(memory_at_{memory}, ")?; - self.pointer().write(w)?; + self.pointer().write(mng, w)?; if self.offset() != 0 { write!(w, " + {}", self.offset())?; } write!(w, ", ")?; - self.value().write(w)?; + self.value().write(mng, w)?; write!(w, ")") } } -impl DriverNoContext for MemoryGrow { - fn write(&self, w: &mut dyn Write) -> Result<()> { - let result = self.result(); +impl Driver for MemoryGrow { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { let memory = self.memory(); - write!(w, "reg_{result} = rt.allocator.grow(memory_at_{memory}, ")?; - self.size().write(w)?; + self.result().write(mng, w)?; + write!(w, " = rt.allocator.grow(memory_at_{memory}, ")?; + self.size().write(mng, w)?; write!(w, ")") } } -impl DriverNoContext for MemoryCopy { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for MemoryCopy { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { let memory_1 = self.destination().memory(); let memory_2 = self.source().memory(); write!(w, "rt.store.copy(memory_at_{memory_1}, ")?; - self.destination().pointer().write(w)?; + self.destination().pointer().write(mng, w)?; write!(w, ", memory_at_{memory_2}, ")?; - self.source().pointer().write(w)?; + self.source().pointer().write(mng, w)?; write!(w, ", ")?; - self.size().write(w)?; + self.size().write(mng, w)?; write!(w, ")") } } -impl DriverNoContext for MemoryFill { - fn write(&self, w: &mut dyn Write) -> Result<()> { +impl Driver for MemoryFill { + fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { let memory = self.destination().memory(); write!(w, "rt.store.fill(memory_at_{memory}, ")?; - self.destination().pointer().write(w)?; + self.destination().pointer().write(mng, w)?; write!(w, ", ")?; - self.size().write(w)?; + self.size().write(mng, w)?; write!(w, ", ")?; - self.value().write(w)?; + self.value().write(mng, w)?; write!(w, ")") } } -fn write_stat(stat: &dyn DriverNoContext, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { +fn write_stat(stat: &dyn Driver, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { indentation!(mng, w)?; - stat.write(w)?; + stat.write(mng, w)?; writeln!(w) } @@ -359,45 +364,53 @@ impl Driver for Statement { } } -fn has_sane_variables(ast: &FuncData) -> bool { - let local_count = ast.local_data().iter().map(|v| v.0).sum::(); - let local_count = usize::try_from(local_count).unwrap(); - let param_count = ast.num_param(); - let temp_count = ast.num_stack(); - - temp_count + param_count + local_count < 170 -} - fn write_parameter_list(ast: &FuncData, w: &mut dyn Write) -> Result<()> { write!(w, "function(")?; - write_ascending("loc", 0..ast.num_param(), w)?; + write_separated(0..ast.num_param(), |i, w| write!(w, "loc_{i}"), w)?; writeln!(w, ")") } +const fn type_to_zero(typ: ValType) -> &'static str { + match typ { + ValType::F32 | ValType::F64 => "0.0", + ValType::I64 => "i64_ZERO", + _ => "0", + } +} + fn write_variable_list(ast: &FuncData, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { - let mut total = ast.num_param(); + let mut locals = ast.local_data().iter().copied(); + let num_local = mng.num_local() - ast.num_param(); - for data in ast.local_data().iter().filter(|v| v.0 != 0) { - let range = total..total + usize::try_from(data.0).unwrap(); - let zero = if data.1 == ValType::I64 { - "i64_ZERO" - } else { - "0" - }; + for (i, typ) in locals.by_ref().enumerate().take(num_local) { + let index = ast.num_param() + i; + let zero = type_to_zero(typ); - total = range.end; - - indented!(mng, w, "local ")?; - write_ascending("loc", range.clone(), w)?; - write!(w, " = ")?; - write_separated(range, |_, w| w.write_all(zero.as_bytes()), w)?; - writeln!(w)?; + line!(mng, w, "local loc_{index} = {zero}")?; } - if ast.num_stack() != 0 { - indented!(mng, w, "local ")?; - write_ascending("reg", 0..ast.num_stack(), w)?; - writeln!(w)?; + if locals.len() != 0 { + indented!(mng, w, "local loc_spill = {{ ")?; + + for typ in locals { + let zero = type_to_zero(typ); + + write!(w, "{zero}, ")?; + } + + writeln!(w, "}}")?; + } + + let mut temporaries = 0..ast.num_stack(); + + for i in temporaries.by_ref().take(mng.num_temp()) { + line!(mng, w, "local reg_{i}")?; + } + + if !temporaries.is_empty() { + let len = temporaries.len(); + + line!(mng, w, "local reg_spill = table.create({len})")?; } Ok(()) @@ -405,34 +418,26 @@ fn write_variable_list(ast: &FuncData, mng: &mut Manager, w: &mut dyn Write) -> impl Driver for FuncData { fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> { - if !has_sane_variables(self) { - return Err(std::io::Error::new( - std::io::ErrorKind::Other, - "too many variables in function", - )); - } - - let (table_map, has_branch) = br_target::visit(self); - mng.indent(); write_parameter_list(self, w)?; write_variable_list(self, mng, w)?; - if has_branch { + if mng.has_branch() { line!(mng, w, "local desired")?; } - if !table_map.is_empty() { + if mng.has_table() { line!(mng, w, "local br_map = {{}}")?; } - mng.set_branch_information(table_map, has_branch); self.code().write(mng, w)?; if self.num_result() != 0 { indented!(mng, w, "return ")?; - write_ascending("reg", 0..self.num_result(), w)?; + + ResultList::new(0, self.num_result()).write(mng, w)?; + writeln!(w)?; } diff --git a/codegen/luau/src/translator.rs b/codegen/luau/src/translator.rs index 54b690c..7dbb21c 100644 --- a/codegen/luau/src/translator.rs +++ b/codegen/luau/src/translator.rs @@ -15,7 +15,7 @@ use wasmparser::{ use crate::{ analyzer::localize, - backend::manager::{Driver, DriverNoContext, Manager}, + backend::manager::{Driver, Manager}, }; trait AsIEName { @@ -51,7 +51,7 @@ fn write_constant(init: &ConstExpr, type_info: &TypeInfo, w: &mut dyn Write) -> let func = Factory::from_type_info(type_info).create_anonymous(&code); if let Some(Statement::SetTemporary(stat)) = func.code().code().last() { - stat.value().write(w) + stat.value().write(&mut Manager::empty(), w) } else { writeln!(w, r#"error("Valueless constant")"#) } @@ -283,7 +283,7 @@ fn write_func_list(wasm: &Module, func_list: &[FuncData], w: &mut dyn Write) -> write_func_start(wasm, index, w)?; - v.write(&mut Manager::default(), w) + v.write(&mut Manager::function(v), w) }) } @@ -322,9 +322,9 @@ fn write_module_start( /// # Errors /// Returns `Err` if writing to `Write` failed. pub fn from_inst_list(code: &[Operator], type_info: &TypeInfo, w: &mut dyn Write) -> Result<()> { - Factory::from_type_info(type_info) - .create_anonymous(code) - .write(&mut Manager::default(), w) + let ast = Factory::from_type_info(type_info).create_anonymous(code); + + ast.write(&mut Manager::function(&ast), w) } /// # Errors diff --git a/wasm-ast/src/factory.rs b/wasm-ast/src/factory.rs index 15f6e51..8469ec2 100644 --- a/wasm-ast/src/factory.rs +++ b/wasm-ast/src/factory.rs @@ -1,10 +1,10 @@ use wasmparser::{BlockType, FunctionBody, MemArg, Operator, Result}; use crate::{ - module::{read_checked, TypeInfo}, + module::{read_checked, read_checked_locals, TypeInfo}, node::{ BinOp, BinOpType, Block, Br, BrIf, BrTable, Call, CallIndirect, CmpOp, CmpOpType, - Expression, FuncData, GetGlobal, GetLocal, If, LabelType, LoadAt, LoadType, MemoryArgument, + Expression, FuncData, GetGlobal, If, LabelType, LoadAt, LoadType, Local, MemoryArgument, MemoryCopy, MemoryFill, MemoryGrow, MemorySize, Select, SetGlobal, SetLocal, Statement, StoreAt, StoreType, Terminator, UnOp, UnOpType, Value, }, @@ -253,13 +253,13 @@ impl<'a> Factory<'a> { /// Returns an error if the function is malformed. pub fn create_indexed(&mut self, index: usize, func: &FunctionBody) -> Result { let code = read_checked(func.get_operators_reader()?)?; - let local = read_checked(func.get_locals_reader()?)?; + let local_data = read_checked_locals(func.get_locals_reader()?)?; let (num_param, num_result) = self.type_info.by_func_index(index); let data = self.build_stat_list(&code, num_result); Ok(FuncData { - local_data: local, + local_data, num_result, num_param, num_stack: data.stack.capacity, @@ -279,7 +279,7 @@ impl<'a> Factory<'a> { BlockVariant::If => BlockData::If { num_result, ty }, BlockVariant::Else => { old.stack.pop_len(num_result).for_each(drop); - old.stack.push_temporary(num_param); + old.stack.push_temporaries(num_param); BlockData::Else { num_result } } @@ -287,7 +287,7 @@ impl<'a> Factory<'a> { self.target.stack = old.stack.split_last(num_param, num_result); - old.stack.push_temporary(num_result); + old.stack.push_temporaries(num_result); self.pending.push(old); } @@ -358,12 +358,12 @@ impl<'a> Factory<'a> { self.target.leak_pre_call(); - let result = self.target.stack.push_temporary(num_result); + let result_list = self.target.stack.push_temporaries(num_result); let data = Statement::Call(Call { function, - result, param_list, + result_list, }); self.target.code.push(data); @@ -376,13 +376,13 @@ impl<'a> Factory<'a> { self.target.leak_pre_call(); - let result = self.target.stack.push_temporary(num_result); + let result_list = self.target.stack.push_temporaries(num_result); let data = Statement::CallIndirect(CallIndirect { table, index, - result, param_list, + result_list, }); self.target.code.push(data); @@ -521,14 +521,14 @@ impl<'a> Factory<'a> { } Operator::LocalGet { local_index } => { let var = local_index.try_into().unwrap(); - let data = Expression::GetLocal(GetLocal { var }); + let data = Expression::GetLocal(Local { var }); self.target.stack.push_with_single(data); } Operator::LocalSet { local_index } => { let var = local_index.try_into().unwrap(); let data = Statement::SetLocal(SetLocal { - var, + var: Local { var }, value: self.target.stack.pop().into(), }); @@ -537,9 +537,9 @@ impl<'a> Factory<'a> { } Operator::LocalTee { local_index } => { let var = local_index.try_into().unwrap(); - let get = Expression::GetLocal(GetLocal { var }); + let get = Expression::GetLocal(Local { var }); let set = Statement::SetLocal(SetLocal { - var, + var: Local { var }, value: self.target.stack.pop().into(), }); @@ -594,7 +594,7 @@ impl<'a> Factory<'a> { } Operator::MemoryGrow { mem, .. } => { let size = self.target.stack.pop().into(); - let result = self.target.stack.push_temporary(1).start; + let result = self.target.stack.push_temporary(); let memory = mem.try_into().unwrap(); let data = Statement::MemoryGrow(MemoryGrow { diff --git a/wasm-ast/src/module.rs b/wasm-ast/src/module.rs index 5b890e3..48f4933 100644 --- a/wasm-ast/src/module.rs +++ b/wasm-ast/src/module.rs @@ -1,8 +1,9 @@ use std::collections::HashMap; use wasmparser::{ - BlockType, Data, Element, Export, ExternalKind, FunctionBody, Global, Import, MemoryType, Name, - NameSectionReader, Parser, Payload, Result, TableType, Type, TypeRef, + BlockType, Data, Element, Export, ExternalKind, FunctionBody, Global, Import, LocalsReader, + MemoryType, Name, NameSectionReader, Parser, Payload, Result, TableType, Type, TypeRef, + ValType, }; #[derive(PartialEq, Eq, Clone, Copy)] @@ -45,6 +46,14 @@ where reader.into_iter().collect() } +pub(crate) fn read_checked_locals(reader: LocalsReader) -> Result> { + read_checked(reader).map(|locals| { + let convert = |(a, b)| std::iter::repeat(b).take(usize::try_from(a).unwrap()); + + locals.into_iter().flat_map(convert).collect() + }) +} + pub struct Module<'a> { type_section: Vec, import_section: Vec>, diff --git a/wasm-ast/src/node.rs b/wasm-ast/src/node.rs index 7e17e79..8a094d1 100644 --- a/wasm-ast/src/node.rs +++ b/wasm-ast/src/node.rs @@ -1,5 +1,3 @@ -use std::ops::Range; - use wasmparser::{Operator, ValType}; #[allow(non_camel_case_types)] @@ -627,22 +625,22 @@ impl Select { } } -pub struct GetTemporary { +pub struct Temporary { pub(crate) var: usize, } -impl GetTemporary { +impl Temporary { #[must_use] pub const fn var(&self) -> usize { self.var } } -pub struct GetLocal { +pub struct Local { pub(crate) var: usize, } -impl GetLocal { +impl Local { #[must_use] pub const fn var(&self) -> usize { self.var @@ -797,8 +795,8 @@ impl CmpOp { pub enum Expression { Select(Select), - GetTemporary(GetTemporary), - GetLocal(GetLocal), + GetTemporary(Temporary), + GetLocal(Local), GetGlobal(GetGlobal), LoadAt(LoadAt), MemorySize(MemorySize), @@ -808,6 +806,28 @@ pub enum Expression { CmpOp(CmpOp), } +#[derive(Clone, Copy)] +pub struct ResultList { + start: usize, + end: usize, +} + +impl ResultList { + #[must_use] + pub const fn new(start: usize, end: usize) -> Self { + Self { start, end } + } + + #[must_use] + pub const fn is_empty(self) -> bool { + self.start == self.end + } + + pub fn iter(self) -> impl Iterator { + (self.start..self.end).map(|var| Temporary { var }) + } +} + pub struct Align { pub(crate) new: usize, pub(crate) old: usize, @@ -821,13 +841,13 @@ impl Align { } #[must_use] - pub const fn new_range(&self) -> Range { - self.new..self.new + self.length + pub const fn new_range(&self) -> ResultList { + ResultList::new(self.new, self.new + self.length) } #[must_use] - pub const fn old_range(&self) -> Range { - self.old..self.old + self.length + pub const fn old_range(&self) -> ResultList { + ResultList::new(self.old, self.old + self.length) } } @@ -949,8 +969,8 @@ impl If { pub struct Call { pub(crate) function: usize, - pub(crate) result: Range, pub(crate) param_list: Vec, + pub(crate) result_list: ResultList, } impl Call { @@ -960,21 +980,21 @@ impl Call { } #[must_use] - pub fn result(&self) -> Range { - self.result.clone() + pub fn param_list(&self) -> &[Expression] { + &self.param_list } #[must_use] - pub fn param_list(&self) -> &[Expression] { - &self.param_list + pub const fn result_list(&self) -> ResultList { + self.result_list } } pub struct CallIndirect { pub(crate) table: usize, pub(crate) index: Box, - pub(crate) result: Range, pub(crate) param_list: Vec, + pub(crate) result_list: ResultList, } impl CallIndirect { @@ -989,25 +1009,25 @@ impl CallIndirect { } #[must_use] - pub fn result(&self) -> Range { - self.result.clone() + pub fn param_list(&self) -> &[Expression] { + &self.param_list } #[must_use] - pub fn param_list(&self) -> &[Expression] { - &self.param_list + pub const fn result_list(&self) -> ResultList { + self.result_list } } pub struct SetTemporary { - pub(crate) var: usize, + pub(crate) var: Temporary, pub(crate) value: Box, } impl SetTemporary { #[must_use] - pub const fn var(&self) -> usize { - self.var + pub const fn var(&self) -> &Temporary { + &self.var } #[must_use] @@ -1017,14 +1037,14 @@ impl SetTemporary { } pub struct SetLocal { - pub(crate) var: usize, + pub(crate) var: Local, pub(crate) value: Box, } impl SetLocal { #[must_use] - pub const fn var(&self) -> usize { - self.var + pub const fn var(&self) -> &Local { + &self.var } #[must_use] @@ -1087,7 +1107,7 @@ impl StoreAt { pub struct MemoryGrow { pub(crate) memory: usize, - pub(crate) result: usize, + pub(crate) result: Temporary, pub(crate) size: Box, } @@ -1098,8 +1118,8 @@ impl MemoryGrow { } #[must_use] - pub const fn result(&self) -> usize { - self.result + pub const fn result(&self) -> &Temporary { + &self.result } #[must_use] @@ -1187,7 +1207,7 @@ pub enum Statement { } pub struct FuncData { - pub(crate) local_data: Vec<(u32, ValType)>, + pub(crate) local_data: Vec, pub(crate) num_result: usize, pub(crate) num_param: usize, pub(crate) num_stack: usize, @@ -1196,7 +1216,7 @@ pub struct FuncData { impl FuncData { #[must_use] - pub fn local_data(&self) -> &[(u32, ValType)] { + pub fn local_data(&self) -> &[ValType] { &self.local_data } diff --git a/wasm-ast/src/stack.rs b/wasm-ast/src/stack.rs index a8290f2..9a73430 100644 --- a/wasm-ast/src/stack.rs +++ b/wasm-ast/src/stack.rs @@ -1,7 +1,7 @@ -use std::{collections::HashSet, ops::Range}; +use std::collections::HashSet; use crate::node::{ - Align, Expression, GetGlobal, GetLocal, GetTemporary, LoadAt, SetTemporary, Statement, + Align, Expression, GetGlobal, LoadAt, Local, ResultList, SetTemporary, Statement, Temporary, }; #[derive(Clone, Copy, PartialEq, Eq, Hash)] @@ -68,7 +68,7 @@ impl Stack { pub fn push_with_single(&mut self, data: Expression) { let mut read = HashSet::new(); let elem = match data { - Expression::GetLocal(GetLocal { var }) => ReadType::Local(var), + Expression::GetLocal(Local { var }) => ReadType::Local(var), Expression::GetGlobal(GetGlobal { var }) => ReadType::Global(var), Expression::LoadAt(LoadAt { memory, .. }) => ReadType::Memory(memory), _ => unreachable!(), @@ -94,19 +94,23 @@ impl Stack { self.var_list.drain(desired..).map(|v| v.data) } - pub fn push_temporary(&mut self, num: usize) -> Range { + pub fn push_temporaries(&mut self, num: usize) -> ResultList { let start = self.previous + self.len(); let range = start..start + num; self.capacity = self.capacity.max(range.end); for var in range.clone() { - let data = Expression::GetTemporary(GetTemporary { var }); + let data = Expression::GetTemporary(Temporary { var }); self.push(data); } - range + ResultList::new(range.start, range.end) + } + + pub fn push_temporary(&mut self) -> Temporary { + self.push_temporaries(1).iter().next().unwrap() } // Return the alignment necessary for this block to branch out to a @@ -136,9 +140,9 @@ impl Stack { old.read.clear(); - let get = Expression::GetTemporary(GetTemporary { var }); + let get = Expression::GetTemporary(Temporary { var }); let set = Statement::SetTemporary(SetTemporary { - var, + var: Temporary { var }, value: std::mem::replace(&mut old.data, get).into(), }); diff --git a/wasm-ast/src/visit.rs b/wasm-ast/src/visit.rs index b8b9605..5a13e79 100644 --- a/wasm-ast/src/visit.rs +++ b/wasm-ast/src/visit.rs @@ -1,15 +1,15 @@ use crate::node::{ BinOp, Block, Br, BrIf, BrTable, Call, CallIndirect, CmpOp, Expression, FuncData, GetGlobal, - GetLocal, GetTemporary, If, LoadAt, MemoryCopy, MemoryFill, MemoryGrow, MemorySize, Select, - SetGlobal, SetLocal, SetTemporary, Statement, StoreAt, Terminator, UnOp, Value, + If, LoadAt, Local, MemoryCopy, MemoryFill, MemoryGrow, MemorySize, Select, SetGlobal, SetLocal, + SetTemporary, Statement, StoreAt, Temporary, Terminator, UnOp, Value, }; pub trait Visitor { fn visit_select(&mut self, _: &Select) {} - fn visit_get_temporary(&mut self, _: &GetTemporary) {} + fn visit_get_temporary(&mut self, _: &Temporary) {} - fn visit_get_local(&mut self, _: &GetLocal) {} + fn visit_get_local(&mut self, _: &Local) {} fn visit_get_global(&mut self, _: &GetGlobal) {} @@ -76,13 +76,13 @@ impl Driver for Select { } } -impl Driver for GetTemporary { +impl Driver for Temporary { fn accept(&self, visitor: &mut T) { visitor.visit_get_temporary(self); } } -impl Driver for GetLocal { +impl Driver for Local { fn accept(&self, visitor: &mut T) { visitor.visit_get_local(self); }