From 49bc994353185eff1b78260c971ff6fa8c2f0811 Mon Sep 17 00:00:00 2001
From: Rerumu <rerumu@protonmail.com>
Date: Thu, 7 Jul 2022 19:30:41 -0400
Subject: [PATCH] Elide block labels when not referenced

---
 codegen/luajit/src/backend/statement.rs | 50 +++++++-------
 codegen/luau/src/backend/manager.rs     | 14 ++--
 codegen/luau/src/backend/statement.rs   | 90 ++++++++++---------------
 wasm-ast/src/factory.rs                 | 51 ++++++++------
 wasm-ast/src/node.rs                    | 43 +++++-------
 wasm-ast/src/visit.rs                   | 31 ++-------
 6 files changed, 116 insertions(+), 163 deletions(-)

diff --git a/codegen/luajit/src/backend/statement.rs b/codegen/luajit/src/backend/statement.rs
index d3d86dc..1b5ea30 100644
--- a/codegen/luajit/src/backend/statement.rs
+++ b/codegen/luajit/src/backend/statement.rs
@@ -4,7 +4,7 @@ use std::{
 };
 
 use wasm_ast::node::{
-	Backward, Br, BrIf, BrTable, Call, CallIndirect, Forward, FuncData, If, MemoryGrow, SetGlobal,
+	Block, Br, BrIf, BrTable, Call, CallIndirect, FuncData, If, LabelType, MemoryGrow, SetGlobal,
 	SetLocal, SetTemporary, Statement, StoreAt, Terminator,
 };
 use wasmparser::ValType;
@@ -117,41 +117,38 @@ impl Driver for Terminator {
 	}
 }
 
-impl Driver for Forward {
-	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-		let label = mng.push_label();
+fn write_inner_block(block: &Block, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+	block.code().iter().try_for_each(|s| s.write(mng, w))?;
 
-		self.code().iter().try_for_each(|s| s.write(mng, w))?;
-
-		if let Some(v) = self.last() {
-			v.write(mng, w)?;
-		}
-
-		write!(w, "::continue_at_{label}::")?;
-
-		mng.pop_label();
-
-		Ok(())
+	match block.last() {
+		Some(v) => v.write(mng, w),
+		None => Ok(()),
 	}
 }
 
-impl Driver for Backward {
+impl Driver for Block {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 		let label = mng.push_label();
 
-		write!(w, "::continue_at_{label}::")?;
-		write!(w, "while true do ")?;
+		match self.label_type() {
+			Some(LabelType::Forward) => {
+				write_inner_block(self, mng, w)?;
+				write!(w, "::continue_at_{label}::")?;
+			}
+			Some(LabelType::Backward) => {
+				write!(w, "::continue_at_{label}::")?;
+				write!(w, "while true do ")?;
+				write_inner_block(self, mng, w)?;
 
-		self.code().iter().try_for_each(|s| s.write(mng, w))?;
+				if self.last().is_none() {
+					write!(w, "break ")?;
+				}
 
-		if let Some(v) = self.last() {
-			v.write(mng, w)?;
-		} else {
-			write!(w, "break ")?;
+				write!(w, "end ")?;
+			}
+			None => write_inner_block(self, mng, w)?,
 		}
 
-		write!(w, "end ")?;
-
 		mng.pop_label();
 
 		Ok(())
@@ -268,8 +265,7 @@ impl Driver for MemoryGrow {
 impl Driver for Statement {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 		match self {
-			Self::Forward(s) => s.write(mng, w),
-			Self::Backward(s) => s.write(mng, w),
+			Self::Block(s) => s.write(mng, w),
 			Self::BrIf(s) => s.write(mng, w),
 			Self::If(s) => s.write(mng, w),
 			Self::Call(s) => s.write(mng, w),
diff --git a/codegen/luau/src/backend/manager.rs b/codegen/luau/src/backend/manager.rs
index 182be4c..9c09086 100644
--- a/codegen/luau/src/backend/manager.rs
+++ b/codegen/luau/src/backend/manager.rs
@@ -4,20 +4,14 @@ use std::{
 	ops::Range,
 };
 
-use wasm_ast::node::{BrTable, CmpOp, Expression};
+use wasm_ast::node::{BrTable, CmpOp, Expression, LabelType};
 
 use crate::analyzer::as_symbol::AsSymbol;
 
-#[derive(PartialEq, Eq)]
-pub enum Label {
-	Forward,
-	Backward,
-}
-
 #[derive(Default)]
 pub struct Manager {
 	table_map: HashMap<usize, usize>,
-	label_list: Vec<Label>,
+	label_list: Vec<Option<LabelType>>,
 	num_param: usize,
 }
 
@@ -36,11 +30,11 @@ impl Manager {
 		self.num_param = num_param;
 	}
 
-	pub fn label_list(&self) -> &[Label] {
+	pub fn label_list(&self) -> &[Option<LabelType>] {
 		&self.label_list
 	}
 
-	pub fn push_label(&mut self, label: Label) -> usize {
+	pub fn push_label(&mut self, label: Option<LabelType>) -> usize {
 		self.label_list.push(label);
 
 		self.label_list.len() - 1
diff --git a/codegen/luau/src/backend/statement.rs b/codegen/luau/src/backend/statement.rs
index 726c97c..b3ea3a2 100644
--- a/codegen/luau/src/backend/statement.rs
+++ b/codegen/luau/src/backend/statement.rs
@@ -4,7 +4,7 @@ use std::{
 };
 
 use wasm_ast::node::{
-	Backward, Br, BrIf, BrTable, Call, CallIndirect, Forward, FuncData, If, MemoryGrow, SetGlobal,
+	Block, Br, BrIf, BrTable, Call, CallIndirect, FuncData, If, LabelType, MemoryGrow, SetGlobal,
 	SetLocal, SetTemporary, Statement, StoreAt, Terminator,
 };
 use wasmparser::ValType;
@@ -12,7 +12,7 @@ use wasmparser::ValType;
 use crate::analyzer::br_table;
 
 use super::manager::{
-	write_ascending, write_condition, write_separated, write_variable, Driver, Label, Manager,
+	write_ascending, write_condition, write_separated, write_variable, Driver, Manager,
 };
 
 impl Driver for Br {
@@ -25,7 +25,7 @@ impl Driver for Br {
 		}
 
 		if self.target() == 0 {
-			if let Some(&Label::Backward) = mng.label_list().last() {
+			if let Some(Some(LabelType::Backward)) = mng.label_list().last() {
 				write!(w, "continue ")
 			} else {
 				write!(w, "break ")
@@ -125,73 +125,54 @@ impl Driver for Terminator {
 	}
 }
 
-fn br_target(level: usize, in_loop: bool, w: &mut dyn Write) -> Result<()> {
-	write!(w, "if desired then ")?;
-	write!(w, "if desired == {level} then ")?;
-	write!(w, "desired = nil ")?;
-
-	if in_loop {
-		write!(w, "continue ")?;
-	}
-
-	write!(w, "else ")?;
-	write!(w, "break ")?;
-	write!(w, "end ")?;
-	write!(w, "end ")
-}
-
-fn write_br_gadget(label_list: &[Label], rem: usize, w: &mut dyn Write) -> Result<()> {
+fn write_br_gadget(
+	label_list: &[Option<LabelType>],
+	label: usize,
+	w: &mut dyn Write,
+) -> Result<()> {
 	if label_list.len() == 1 {
 		return Ok(());
 	}
 
-	match label_list.last() {
-		Some(Label::Forward) => br_target(rem, false, w),
-		Some(Label::Backward) => br_target(rem, true, w),
-		None => Ok(()),
+	write!(w, "if desired then ")?;
+
+	match label_list.last().unwrap() {
+		Some(t) => {
+			write!(w, "if desired == {label} then ")?;
+			write!(w, "desired = nil ")?;
+
+			if *t == LabelType::Backward {
+				write!(w, "continue ")?;
+			}
+
+			write!(w, "else ")?;
+			write!(w, "break ")?;
+			write!(w, "end ")?;
+		}
+		None => {
+			write!(w, "break ")?;
+		}
 	}
+
+	write!(w, "end ")
 }
 
-impl Driver for Forward {
+impl Driver for Block {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-		let rem = mng.push_label(Label::Forward);
+		let label = mng.push_label(self.label_type());
 
 		write!(w, "while true do ")?;
 
 		self.code().iter().try_for_each(|s| s.write(mng, w))?;
 
-		if let Some(v) = self.last() {
-			v.write(mng, w)?;
-		} else {
-			write!(w, "break ")?;
+		match self.last() {
+			Some(v) => v.write(mng, w)?,
+			None => write!(w, "break ")?,
 		}
 
 		write!(w, "end ")?;
 
-		write_br_gadget(mng.label_list(), rem, w)?;
-		mng.pop_label();
-
-		Ok(())
-	}
-}
-
-impl Driver for Backward {
-	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-		let rem = mng.push_label(Label::Backward);
-
-		write!(w, "while true do ")?;
-
-		self.code().iter().try_for_each(|s| s.write(mng, w))?;
-
-		if let Some(v) = self.last() {
-			v.write(mng, w)?;
-		} else {
-			write!(w, "break ")?;
-		}
-
-		write!(w, "end ")?;
-
-		write_br_gadget(mng.label_list(), rem, w)?;
+		write_br_gadget(mng.label_list(), label, w)?;
 		mng.pop_label();
 
 		Ok(())
@@ -308,8 +289,7 @@ impl Driver for MemoryGrow {
 impl Driver for Statement {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 		match self {
-			Self::Forward(s) => s.write(mng, w),
-			Self::Backward(s) => s.write(mng, w),
+			Self::Block(s) => s.write(mng, w),
 			Self::BrIf(s) => s.write(mng, w),
 			Self::If(s) => s.write(mng, w),
 			Self::Call(s) => s.write(mng, w),
diff --git a/wasm-ast/src/factory.rs b/wasm-ast/src/factory.rs
index 232cb84..3d074b4 100644
--- a/wasm-ast/src/factory.rs
+++ b/wasm-ast/src/factory.rs
@@ -3,8 +3,8 @@ use wasmparser::{BlockType, FunctionBody, MemoryImmediate, Operator};
 use crate::{
 	module::TypeInfo,
 	node::{
-		Backward, BinOp, BinOpType, Br, BrIf, BrTable, Call, CallIndirect, CmpOp, CmpOpType,
-		Expression, Forward, FuncData, GetGlobal, GetLocal, If, LoadAt, LoadType, MemoryGrow,
+		BinOp, BinOpType, Block, Br, BrIf, BrTable, Call, CallIndirect, CmpOp, CmpOpType,
+		Expression, FuncData, GetGlobal, GetLocal, If, LabelType, LoadAt, LoadType, MemoryGrow,
 		MemorySize, Select, SetGlobal, SetLocal, Statement, StoreAt, StoreType, Terminator, UnOp,
 		UnOpType, Value,
 	},
@@ -42,6 +42,17 @@ impl Default for BlockData {
 	}
 }
 
+impl From<BlockData> for LabelType {
+	fn from(data: BlockData) -> Self {
+		match data {
+			BlockData::Forward { .. } | BlockData::If { .. } | BlockData::Else { .. } => {
+				Self::Forward
+			}
+			BlockData::Backward { .. } => Self::Backward,
+		}
+	}
+}
+
 #[derive(Default)]
 struct StatList {
 	stack: Stack,
@@ -49,6 +60,7 @@ struct StatList {
 	last: Option<Terminator>,
 
 	block_data: BlockData,
+	has_reference: bool,
 }
 
 impl StatList {
@@ -191,18 +203,12 @@ impl StatList {
 	}
 }
 
-impl From<StatList> for Forward {
+impl From<StatList> for Block {
 	fn from(stat: StatList) -> Self {
-		Self {
-			code: stat.code,
-			last: stat.last,
-		}
-	}
-}
+		let label_type = stat.has_reference.then(|| stat.block_data.into());
 
-impl From<StatList> for Backward {
-	fn from(stat: StatList) -> Self {
 		Self {
+			label_type,
 			code: stat.code,
 			last: stat.last,
 		}
@@ -305,8 +311,7 @@ impl<'a> Factory<'a> {
 		self.target.stack.capacity = now.stack.capacity;
 
 		let stat = match now.block_data {
-			BlockData::Forward { .. } => Statement::Forward(now.into()),
-			BlockData::Backward { .. } => Statement::Backward(now.into()),
+			BlockData::Forward { .. } | BlockData::Backward { .. } => Statement::Block(now.into()),
 			BlockData::If { .. } => Statement::If(If {
 				condition: self.target.stack.pop(),
 				on_true: now.into(),
@@ -326,27 +331,29 @@ impl<'a> Factory<'a> {
 		self.target.code.push(stat);
 	}
 
-	fn get_relative_block(&self, index: usize) -> &StatList {
+	fn get_relative_block(&mut self, index: usize) -> &mut StatList {
 		if index == 0 {
-			&self.target
+			&mut self.target
 		} else {
-			&self.pending[self.pending.len() - index]
+			let index = self.pending.len() - index;
+
+			&mut self.pending[index]
 		}
 	}
 
-	fn get_br_terminator(&self, target: usize) -> Br {
+	fn get_br_terminator(&mut self, target: usize) -> Br {
 		let block = self.get_relative_block(target);
-		let par_result = match block.block_data {
+		let previous = block.stack.previous;
+		let result = match block.block_data {
 			BlockData::Forward { num_result }
 			| BlockData::If { num_result, .. }
 			| BlockData::Else { num_result } => num_result,
 			BlockData::Backward { num_param } => num_param,
 		};
 
-		let align = self
-			.target
-			.stack
-			.get_br_alignment(block.stack.previous, par_result);
+		block.has_reference = true;
+
+		let align = self.target.stack.get_br_alignment(previous, result);
 
 		Br { target, align }
 	}
diff --git a/wasm-ast/src/node.rs b/wasm-ast/src/node.rs
index d39b1f9..bc56778 100644
--- a/wasm-ast/src/node.rs
+++ b/wasm-ast/src/node.rs
@@ -847,6 +847,12 @@ impl BrTable {
 	}
 }
 
+#[derive(PartialEq, Eq, Clone, Copy)]
+pub enum LabelType {
+	Forward,
+	Backward,
+}
+
 pub enum Terminator {
 	Unreachable,
 	Br(Br),
@@ -854,30 +860,18 @@ pub enum Terminator {
 }
 
 #[derive(Default)]
-pub struct Forward {
+pub struct Block {
+	pub(crate) label_type: Option<LabelType>,
 	pub(crate) code: Vec<Statement>,
 	pub(crate) last: Option<Terminator>,
 }
 
-impl Forward {
+impl Block {
 	#[must_use]
-	pub fn code(&self) -> &[Statement] {
-		&self.code
+	pub fn label_type(&self) -> Option<LabelType> {
+		self.label_type
 	}
 
-	#[must_use]
-	pub fn last(&self) -> Option<&Terminator> {
-		self.last.as_ref()
-	}
-}
-
-#[derive(Default)]
-pub struct Backward {
-	pub(crate) code: Vec<Statement>,
-	pub(crate) last: Option<Terminator>,
-}
-
-impl Backward {
 	#[must_use]
 	pub fn code(&self) -> &[Statement] {
 		&self.code
@@ -908,8 +902,8 @@ impl BrIf {
 
 pub struct If {
 	pub(crate) condition: Expression,
-	pub(crate) on_true: Forward,
-	pub(crate) on_false: Option<Forward>,
+	pub(crate) on_true: Block,
+	pub(crate) on_false: Option<Block>,
 }
 
 impl If {
@@ -919,12 +913,12 @@ impl If {
 	}
 
 	#[must_use]
-	pub fn on_true(&self) -> &Forward {
+	pub fn on_true(&self) -> &Block {
 		&self.on_true
 	}
 
 	#[must_use]
-	pub fn on_false(&self) -> Option<&Forward> {
+	pub fn on_false(&self) -> Option<&Block> {
 		self.on_false.as_ref()
 	}
 }
@@ -1091,8 +1085,7 @@ impl MemoryGrow {
 }
 
 pub enum Statement {
-	Forward(Forward),
-	Backward(Backward),
+	Block(Block),
 	BrIf(BrIf),
 	If(If),
 	Call(Call),
@@ -1109,7 +1102,7 @@ pub struct FuncData {
 	pub(crate) num_result: usize,
 	pub(crate) num_param: usize,
 	pub(crate) num_stack: usize,
-	pub(crate) code: Forward,
+	pub(crate) code: Block,
 }
 
 impl FuncData {
@@ -1134,7 +1127,7 @@ impl FuncData {
 	}
 
 	#[must_use]
-	pub fn code(&self) -> &Forward {
+	pub fn code(&self) -> &Block {
 		&self.code
 	}
 }
diff --git a/wasm-ast/src/visit.rs b/wasm-ast/src/visit.rs
index d771c86..0961e76 100644
--- a/wasm-ast/src/visit.rs
+++ b/wasm-ast/src/visit.rs
@@ -1,7 +1,7 @@
 use crate::node::{
-	Backward, BinOp, Br, BrIf, BrTable, Call, CallIndirect, CmpOp, Expression, Forward, FuncData,
-	GetGlobal, GetLocal, GetTemporary, If, LoadAt, MemoryGrow, MemorySize, Select, SetGlobal,
-	SetLocal, SetTemporary, Statement, StoreAt, Terminator, UnOp, Value,
+	BinOp, Block, Br, BrIf, BrTable, Call, CallIndirect, CmpOp, Expression, FuncData, GetGlobal,
+	GetLocal, GetTemporary, If, LoadAt, MemoryGrow, MemorySize, Select, SetGlobal, SetLocal,
+	SetTemporary, Statement, StoreAt, Terminator, UnOp, Value,
 };
 
 pub trait Visitor {
@@ -35,9 +35,7 @@ pub trait Visitor {
 
 	fn visit_terminator(&mut self, _: &Terminator) {}
 
-	fn visit_forward(&mut self, _: &Forward) {}
-
-	fn visit_backward(&mut self, _: &Backward) {}
+	fn visit_block(&mut self, _: &Block) {}
 
 	fn visit_br_if(&mut self, _: &BrIf) {}
 
@@ -183,7 +181,7 @@ impl<T: Visitor> Driver<T> for Terminator {
 	}
 }
 
-impl<T: Visitor> Driver<T> for Forward {
+impl<T: Visitor> Driver<T> for Block {
 	fn accept(&self, visitor: &mut T) {
 		for v in self.code() {
 			v.accept(visitor);
@@ -193,21 +191,7 @@ impl<T: Visitor> Driver<T> for Forward {
 			v.accept(visitor);
 		}
 
-		visitor.visit_forward(self);
-	}
-}
-
-impl<T: Visitor> Driver<T> for Backward {
-	fn accept(&self, visitor: &mut T) {
-		for v in self.code() {
-			v.accept(visitor);
-		}
-
-		if let Some(v) = self.last() {
-			v.accept(visitor);
-		}
-
-		visitor.visit_backward(self);
+		visitor.visit_block(self);
 	}
 }
 
@@ -298,8 +282,7 @@ impl<T: Visitor> Driver<T> for MemoryGrow {
 impl<T: Visitor> Driver<T> for Statement {
 	fn accept(&self, visitor: &mut T) {
 		match self {
-			Self::Forward(v) => v.accept(visitor),
-			Self::Backward(v) => v.accept(visitor),
+			Self::Block(v) => v.accept(visitor),
 			Self::BrIf(v) => v.accept(visitor),
 			Self::If(v) => v.accept(visitor),
 			Self::Call(v) => v.accept(visitor),