Elide block labels when not referenced

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 }
 	}

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
 	}
 }

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),