Refactor Terminator from Statement

codegen-luajit/src/backend/statement.rs

@@ -6,70 +6,13 @@ use std::{
 use parity_wasm::elements::ValueType;
 use wasm_ast::node::{
 	Backward, Br, BrIf, BrTable, Call, CallIndirect, Forward, If, Intermediate, Return, SetGlobal,
-	SetLocal, SetTemporary, Statement, StoreAt,
+	SetLocal, SetTemporary, Statement, StoreAt, Terminator,
 };
 
 use super::manager::{
 	write_ascending, write_condition, write_separated, write_variable, Driver, Manager,
 };
 
-impl Driver for Forward {
-	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-		let label = mng.push_label();
-
-		self.body.iter().try_for_each(|s| s.write(mng, w))?;
-
-		write!(w, "::continue_at_{label}::")?;
-
-		mng.pop_label();
-
-		Ok(())
-	}
-}
-
-impl Driver for Backward {
-	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 ")?;
-
-		self.body.iter().try_for_each(|s| s.write(mng, w))?;
-
-		write!(w, "break ")?;
-		write!(w, "end ")?;
-
-		mng.pop_label();
-
-		Ok(())
-	}
-}
-
-impl Driver for If {
-	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-		let label = mng.push_label();
-
-		write!(w, "if ")?;
-		write_condition(&self.cond, mng, w)?;
-		write!(w, "then ")?;
-
-		self.truthy.iter().try_for_each(|s| s.write(mng, w))?;
-
-		if !self.falsey.is_empty() {
-			write!(w, "else ")?;
-
-			self.falsey.iter().try_for_each(|s| s.write(mng, w))?;
-		}
-
-		write!(w, "end ")?;
-		write!(w, "::continue_at_{label}::")?;
-
-		mng.pop_label();
-
-		Ok(())
-	}
-}
-
 fn write_br_at(up: usize, mng: &Manager, w: &mut dyn Write) -> Result<()> {
 	let level = mng.label_list().iter().nth_back(up).unwrap();
 
@@ -82,16 +25,6 @@ impl Driver for Br {
 	}
 }
 
-impl Driver for BrIf {
-	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-		write!(w, "if ")?;
-		write_condition(&self.cond, mng, w)?;
-		write!(w, "then ")?;
-		write_br_at(self.target, mng, w)?;
-		write!(w, "end ")
-	}
-}
-
 fn condense_jump_table(list: &[u32]) -> Vec<(usize, usize, u32)> {
 	let mut result = Vec::new();
 	let mut index = 0;
@@ -152,6 +85,92 @@ impl Driver for Return {
 	}
 }
 
+impl Driver for Terminator {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		match self {
+			Self::Unreachable => write!(w, "error(\"out of code bounds\")"),
+			Self::Br(s) => s.write(mng, w),
+			Self::BrTable(s) => s.write(mng, w),
+			Self::Return(s) => s.write(mng, w),
+		}
+	}
+}
+
+impl Driver for Forward {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		let label = mng.push_label();
+
+		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(())
+	}
+}
+
+impl Driver for Backward {
+	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 ")?;
+
+		self.code.iter().try_for_each(|s| s.write(mng, w))?;
+
+		if let Some(v) = &self.last {
+			v.write(mng, w)?;
+		}
+
+		write!(w, "break ")?;
+		write!(w, "end ")?;
+
+		mng.pop_label();
+
+		Ok(())
+	}
+}
+
+impl Driver for If {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		let label = mng.push_label();
+
+		write!(w, "if ")?;
+		write_condition(&self.cond, mng, w)?;
+		write!(w, "then ")?;
+
+		self.truthy.write(mng, w)?;
+
+		if let Some(falsey) = &self.falsey {
+			write!(w, "else ")?;
+
+			falsey.write(mng, w)?;
+		}
+
+		write!(w, "end ")?;
+		write!(w, "::continue_at_{label}::")?;
+
+		mng.pop_label();
+
+		Ok(())
+	}
+}
+
+impl Driver for BrIf {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write!(w, "if ")?;
+		write_condition(&self.cond, mng, w)?;
+		write!(w, "then ")?;
+		write_br_at(self.target, mng, w)?;
+		write!(w, "end ")
+	}
+}
+
 fn write_call_store(result: Range<usize>, w: &mut dyn Write) -> Result<()> {
 	if result.is_empty() {
 		return Ok(());
@@ -223,14 +242,10 @@ impl Driver for StoreAt {
 impl Driver for Statement {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 		match self {
-			Self::Unreachable => write!(w, "error(\"out of code bounds\")"),
 			Self::Forward(s) => s.write(mng, w),
 			Self::Backward(s) => s.write(mng, w),
 			Self::If(s) => s.write(mng, w),
-			Self::Br(s) => s.write(mng, w),
 			Self::BrIf(s) => s.write(mng, w),
-			Self::BrTable(s) => s.write(mng, w),
-			Self::Return(s) => s.write(mng, w),
 			Self::Call(s) => s.write(mng, w),
 			Self::CallIndirect(s) => s.write(mng, w),
 			Self::SetTemporary(s) => s.write(mng, w),

codegen-luau/src/backend/statement.rs

@@ -6,93 +6,13 @@ use std::{
 use parity_wasm::elements::ValueType;
 use wasm_ast::node::{
 	Backward, Br, BrIf, BrTable, Call, CallIndirect, Forward, If, Intermediate, Return, SetGlobal,
-	SetLocal, SetTemporary, Statement, StoreAt,
+	SetLocal, SetTemporary, Statement, StoreAt, Terminator,
 };
 
 use super::manager::{
 	write_ascending, write_condition, write_separated, write_variable, Driver, Label, Manager,
 };
 
-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, "end ")?;
-	write!(w, "break ")?;
-	write!(w, "end ")
-}
-
-fn write_br_gadget(label_list: &[Label], rem: usize, w: &mut dyn Write) -> Result<()> {
-	match label_list.last() {
-		Some(Label::Forward | Label::If) => br_target(rem, false, w),
-		Some(Label::Backward) => br_target(rem, true, w),
-		None => Ok(()),
-	}
-}
-
-impl Driver for Forward {
-	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-		let rem = mng.push_label(Label::Forward);
-
-		write!(w, "while true do ")?;
-
-		self.body.iter().try_for_each(|s| s.write(mng, w))?;
-
-		write!(w, "break ")?;
-		write!(w, "end ")?;
-
-		mng.pop_label();
-		write_br_gadget(mng.label_list(), rem, w)
-	}
-}
-
-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.body.iter().try_for_each(|s| s.write(mng, w))?;
-
-		write!(w, "break ")?;
-		write!(w, "end ")?;
-
-		mng.pop_label();
-		write_br_gadget(mng.label_list(), rem, w)
-	}
-}
-
-impl Driver for If {
-	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-		let rem = mng.push_label(Label::If);
-
-		write!(w, "while true do ")?;
-		write!(w, "if ")?;
-		write_condition(&self.cond, mng, w)?;
-		write!(w, "then ")?;
-
-		self.truthy.iter().try_for_each(|s| s.write(mng, w))?;
-
-		if !self.falsey.is_empty() {
-			write!(w, "else ")?;
-
-			self.falsey.iter().try_for_each(|s| s.write(mng, w))?;
-		}
-
-		write!(w, "end ")?;
-		write!(w, "break ")?;
-		write!(w, "end ")?;
-
-		mng.pop_label();
-		write_br_gadget(mng.label_list(), rem, w)
-	}
-}
-
 fn write_br_at(up: usize, mng: &Manager, w: &mut dyn Write) -> Result<()> {
 	write!(w, "do ")?;
 
@@ -118,16 +38,6 @@ impl Driver for Br {
 	}
 }
 
-impl Driver for BrIf {
-	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-		write!(w, "if ")?;
-		write_condition(&self.cond, mng, w)?;
-		write!(w, "then ")?;
-		write_br_at(self.target, mng, w)?;
-		write!(w, "end ")
-	}
-}
-
 impl Driver for BrTable {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 		write!(w, "do ")?;
@@ -159,6 +69,115 @@ impl Driver for Return {
 	}
 }
 
+impl Driver for Terminator {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		match self {
+			Self::Unreachable => write!(w, "error(\"out of code bounds\")"),
+			Self::Br(s) => s.write(mng, w),
+			Self::BrTable(s) => s.write(mng, w),
+			Self::Return(s) => s.write(mng, w),
+		}
+	}
+}
+
+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, "end ")?;
+	write!(w, "break ")?;
+	write!(w, "end ")
+}
+
+fn write_br_gadget(label_list: &[Label], rem: usize, w: &mut dyn Write) -> Result<()> {
+	match label_list.last() {
+		Some(Label::Forward | Label::If) => br_target(rem, false, w),
+		Some(Label::Backward) => br_target(rem, true, w),
+		None => Ok(()),
+	}
+}
+
+impl Driver for Forward {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		let rem = mng.push_label(Label::Forward);
+
+		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)?;
+		}
+
+		write!(w, "break ")?;
+		write!(w, "end ")?;
+
+		mng.pop_label();
+		write_br_gadget(mng.label_list(), rem, w)
+	}
+}
+
+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)?;
+		}
+
+		write!(w, "break ")?;
+		write!(w, "end ")?;
+
+		mng.pop_label();
+		write_br_gadget(mng.label_list(), rem, w)
+	}
+}
+
+impl Driver for If {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		let rem = mng.push_label(Label::If);
+
+		write!(w, "while true do ")?;
+		write!(w, "if ")?;
+		write_condition(&self.cond, mng, w)?;
+		write!(w, "then ")?;
+
+		self.truthy.write(mng, w)?;
+
+		if let Some(falsey) = &self.falsey {
+			write!(w, "else ")?;
+
+			falsey.write(mng, w)?;
+		}
+
+		write!(w, "end ")?;
+		write!(w, "break ")?;
+		write!(w, "end ")?;
+
+		mng.pop_label();
+		write_br_gadget(mng.label_list(), rem, w)
+	}
+}
+
+impl Driver for BrIf {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write!(w, "if ")?;
+		write_condition(&self.cond, mng, w)?;
+		write!(w, "then ")?;
+		write_br_at(self.target, mng, w)?;
+		write!(w, "end ")
+	}
+}
+
 fn write_call_store(result: Range<usize>, w: &mut dyn Write) -> Result<()> {
 	if result.is_empty() {
 		return Ok(());
@@ -230,14 +249,10 @@ impl Driver for StoreAt {
 impl Driver for Statement {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 		match self {
-			Self::Unreachable => write!(w, "error(\"out of code bounds\")"),
 			Self::Forward(s) => s.write(mng, w),
 			Self::Backward(s) => s.write(mng, w),
 			Self::If(s) => s.write(mng, w),
-			Self::Br(s) => s.write(mng, w),
 			Self::BrIf(s) => s.write(mng, w),
-			Self::BrTable(s) => s.write(mng, w),
-			Self::Return(s) => s.write(mng, w),
 			Self::Call(s) => s.write(mng, w),
 			Self::CallIndirect(s) => s.write(mng, w),
 			Self::SetTemporary(s) => s.write(mng, w),

wasm-ast/src/builder.rs

@@ -7,7 +7,7 @@ use crate::node::{
 	Backward, BinOp, BinOpType, Br, BrIf, BrTable, Call, CallIndirect, CmpOp, CmpOpType,
 	Expression, Forward, GetGlobal, GetLocal, GetTemporary, If, Intermediate, LoadAt, LoadType,
 	MemoryGrow, MemorySize, Return, Select, SetGlobal, SetLocal, SetTemporary, Statement, StoreAt,
-	StoreType, UnOp, UnOpType, Value,
+	StoreType, Terminator, UnOp, UnOpType, Value,
 };
 
 macro_rules! leak_with_predicate {
@@ -112,8 +112,9 @@ impl<'a> TypeInfo<'a> {
 
 #[derive(Default)]
 struct StatList {
-	code: Vec<Statement>,
 	stack: Vec<Expression>,
+	code: Vec<Statement>,
+	last: Option<Terminator>,
 
 	num_result: usize,
 	num_stack: usize,
@@ -286,6 +287,24 @@ impl StatList {
 	}
 }
 
+impl From<StatList> for Forward {
+	fn from(stat: StatList) -> Self {
+		Self {
+			code: stat.code,
+			last: stat.last,
+		}
+	}
+}
+
+impl From<StatList> for Backward {
+	fn from(stat: StatList) -> Self {
+		Self {
+			code: stat.code,
+			last: stat.last,
+		}
+	}
+}
+
 pub struct Builder<'a> {
 	type_info: &'a TypeInfo<'a>,
 
@@ -316,7 +335,7 @@ impl<'a> Builder<'a> {
 			local_data: Vec::new(),
 			num_param: 0,
 			num_stack: data.num_stack,
-			code: Forward { body: data.code },
+			code: data.into(),
 		}
 	}
 
@@ -329,10 +348,11 @@ impl<'a> Builder<'a> {
 			local_data: func.locals().to_vec(),
 			num_param: arity.num_param,
 			num_stack: data.num_stack,
-			code: Forward { body: data.code },
+			code: data.into(),
 		}
 	}
 
+	// FIXME: Sets up temporaries except when used in a weird way
 	fn start_block(&mut self, typ: BlockType) {
 		let mut old = std::mem::take(&mut self.target);
 
@@ -360,10 +380,10 @@ impl<'a> Builder<'a> {
 		self.target.num_stack = now.num_stack;
 
 		match self.target.code.last_mut().unwrap() {
-			Statement::Forward(data) => data.body = now.code,
+			Statement::Forward(data) => *data = now.into(),
-			Statement::Backward(data) => data.body = now.code,
+			Statement::Backward(data) => *data = now.into(),
-			Statement::If(data) if now.is_else => data.falsey = now.code,
+			Statement::If(data) if !now.is_else => data.truthy = now.into(),
-			Statement::If(data) if !now.is_else => data.truthy = now.code,
+			Statement::If(data) if now.is_else => data.falsey = Some(now.into()),
 			_ => unreachable!(),
 		}
 	}
@@ -408,13 +428,13 @@ impl<'a> Builder<'a> {
 		self.target.code.push(data);
 	}
 
-	fn add_return(&mut self) {
+	fn set_return(&mut self) {
-		let data = Statement::Return(Return {
+		let data = Terminator::Return(Return {
 			list: self.target.pop_len(self.num_result),
 		});
 
 		self.target.leak_all();
-		self.target.code.push(data);
+		self.target.last = Some(data);
 	}
 
 	#[cold]
@@ -444,17 +464,17 @@ impl<'a> Builder<'a> {
 		match *inst {
 			Inst::Unreachable => {
 				self.nested_unreachable += 1;
-				self.target.code.push(Statement::Unreachable);
+				self.target.last = Some(Terminator::Unreachable);
 			}
 			Inst::Nop => {}
 			Inst::Block(typ) => {
-				let data = Statement::Forward(Forward { body: Vec::new() });
+				let data = Statement::Forward(Forward::default());
 
 				self.start_block(typ);
 				self.pending.last_mut().unwrap().code.push(data);
 			}
 			Inst::Loop(typ) => {
-				let data = Statement::Backward(Backward { body: Vec::new() });
+				let data = Statement::Backward(Backward::default());
 
 				self.start_block(typ);
 				self.pending.last_mut().unwrap().code.push(data);
@@ -462,8 +482,8 @@ impl<'a> Builder<'a> {
 			Inst::If(typ) => {
 				let data = Statement::If(If {
 					cond: self.target.pop_required(),
-					truthy: Vec::new(),
+					truthy: Forward::default(),
-					falsey: Vec::new(),
+					falsey: None,
 				});
 
 				self.start_block(typ);
@@ -486,12 +506,12 @@ impl<'a> Builder<'a> {
 			Inst::Br(v) => {
 				self.nested_unreachable += 1;
 
-				let data = Statement::Br(Br {
+				let data = Terminator::Br(Br {
 					target: v.try_into().unwrap(),
 				});
 
 				self.target.leak_all();
-				self.target.code.push(data);
+				self.target.last = Some(data);
 			}
 			Inst::BrIf(v) => {
 				let data = Statement::BrIf(BrIf {
@@ -506,17 +526,17 @@ impl<'a> Builder<'a> {
 			Inst::BrTable(ref v) => {
 				self.nested_unreachable += 1;
 
-				let data = Statement::BrTable(BrTable {
+				let data = Terminator::BrTable(BrTable {
 					cond: self.target.pop_required(),
 					data: *v.clone(),
 				});
 
 				self.target.leak_all();
-				self.target.code.push(data);
+				self.target.last = Some(data);
 			}
 			Inst::Return => {
 				self.nested_unreachable += 1;
-				self.add_return();
+				self.set_return();
 			}
 			Inst::Call(i) => {
 				self.add_call(i.try_into().unwrap());
@@ -651,7 +671,7 @@ impl<'a> Builder<'a> {
 		}
 
 		if self.nested_unreachable == 0 && num_result != 0 {
-			self.add_return();
+			self.set_return();
 		}
 
 		std::mem::take(&mut self.target)

wasm-ast/src/node.rs

@@ -660,29 +660,10 @@ impl Expression {
 	}
 }
 
-pub struct Forward {
-	pub body: Vec<Statement>,
-}
-
-pub struct Backward {
-	pub body: Vec<Statement>,
-}
-
-pub struct If {
-	pub cond: Expression,
-	pub truthy: Vec<Statement>,
-	pub falsey: Vec<Statement>,
-}
-
 pub struct Br {
 	pub target: usize,
 }
 
-pub struct BrIf {
-	pub cond: Expression,
-	pub target: usize,
-}
-
 pub struct BrTable {
 	pub cond: Expression,
 	pub data: BrTableData,
@@ -692,6 +673,36 @@ pub struct Return {
 	pub list: Vec<Expression>,
 }
 
+pub enum Terminator {
+	Unreachable,
+	Br(Br),
+	BrTable(BrTable),
+	Return(Return),
+}
+
+#[derive(Default)]
+pub struct Forward {
+	pub code: Vec<Statement>,
+	pub last: Option<Terminator>,
+}
+
+#[derive(Default)]
+pub struct Backward {
+	pub code: Vec<Statement>,
+	pub last: Option<Terminator>,
+}
+
+pub struct If {
+	pub cond: Expression,
+	pub truthy: Forward,
+	pub falsey: Option<Forward>,
+}
+
+pub struct BrIf {
+	pub cond: Expression,
+	pub target: usize,
+}
+
 pub struct Call {
 	pub func: usize,
 	pub result: Range<usize>,
@@ -728,14 +739,10 @@ pub struct StoreAt {
 }
 
 pub enum Statement {
-	Unreachable,
 	Forward(Forward),
 	Backward(Backward),
 	If(If),
-	Br(Br),
 	BrIf(BrIf),
-	BrTable(BrTable),
-	Return(Return),
 	Call(Call),
 	CallIndirect(CallIndirect),
 	SetTemporary(SetTemporary),

wasm-ast/src/visit.rs

@@ -1,7 +1,7 @@
 use crate::node::{
 	Backward, BinOp, Br, BrIf, BrTable, Call, CallIndirect, CmpOp, Expression, Forward, GetGlobal,
 	GetLocal, GetTemporary, If, Intermediate, LoadAt, MemoryGrow, MemorySize, Return, Select,
-	SetGlobal, SetLocal, SetTemporary, Statement, StoreAt, UnOp, Value,
+	SetGlobal, SetLocal, SetTemporary, Statement, StoreAt, Terminator, UnOp, Value,
 };
 
 pub trait Visitor {
@@ -31,20 +31,22 @@ pub trait Visitor {
 
 	fn visit_unreachable(&mut self) {}
 
+	fn visit_br(&mut self, _: &Br) {}
+
+	fn visit_br_table(&mut self, _: &BrTable) {}
+
+	fn visit_return(&mut self, _: &Return) {}
+
+	fn visit_terminator(&mut self, _: &Terminator) {}
+
 	fn visit_forward(&mut self, _: &Forward) {}
 
 	fn visit_backward(&mut self, _: &Backward) {}
 
 	fn visit_if(&mut self, _: &If) {}
 
-	fn visit_br(&mut self, _: &Br) {}
-
 	fn visit_br_if(&mut self, _: &BrIf) {}
 
-	fn visit_br_table(&mut self, _: &BrTable) {}
-
-	fn visit_return(&mut self, _: &Return) {}
-
 	fn visit_call(&mut self, _: &Call) {}
 
 	fn visit_call_indirect(&mut self, _: &CallIndirect) {}
@@ -166,56 +168,12 @@ impl<T: Visitor> Driver<T> for Expression {
 	}
 }
 
-impl<T: Visitor> Driver<T> for Forward {
-	fn accept(&self, visitor: &mut T) {
-		for v in &self.body {
-			v.accept(visitor);
-		}
-
-		visitor.visit_forward(self);
-	}
-}
-
-impl<T: Visitor> Driver<T> for Backward {
-	fn accept(&self, visitor: &mut T) {
-		for v in &self.body {
-			v.accept(visitor);
-		}
-
-		visitor.visit_backward(self);
-	}
-}
-
-impl<T: Visitor> Driver<T> for If {
-	fn accept(&self, visitor: &mut T) {
-		self.cond.accept(visitor);
-
-		for v in &self.truthy {
-			v.accept(visitor);
-		}
-
-		for v in &self.falsey {
-			v.accept(visitor);
-		}
-
-		visitor.visit_if(self);
-	}
-}
-
 impl<T: Visitor> Driver<T> for Br {
 	fn accept(&self, visitor: &mut T) {
 		visitor.visit_br(self);
 	}
 }
 
-impl<T: Visitor> Driver<T> for BrIf {
-	fn accept(&self, visitor: &mut T) {
-		self.cond.accept(visitor);
-
-		visitor.visit_br_if(self);
-	}
-}
-
 impl<T: Visitor> Driver<T> for BrTable {
 	fn accept(&self, visitor: &mut T) {
 		self.cond.accept(visitor);
@@ -234,6 +192,68 @@ impl<T: Visitor> Driver<T> for Return {
 	}
 }
 
+impl<T: Visitor> Driver<T> for Terminator {
+	fn accept(&self, visitor: &mut T) {
+		match self {
+			Self::Unreachable => visitor.visit_unreachable(),
+			Self::Br(v) => v.accept(visitor),
+			Self::BrTable(v) => v.accept(visitor),
+			Self::Return(v) => v.accept(visitor),
+		}
+
+		visitor.visit_terminator(self);
+	}
+}
+
+impl<T: Visitor> Driver<T> for Forward {
+	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_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);
+	}
+}
+
+impl<T: Visitor> Driver<T> for If {
+	fn accept(&self, visitor: &mut T) {
+		self.cond.accept(visitor);
+		self.truthy.accept(visitor);
+
+		if let Some(v) = &self.falsey {
+			v.accept(visitor);
+		}
+
+		visitor.visit_if(self);
+	}
+}
+
+impl<T: Visitor> Driver<T> for BrIf {
+	fn accept(&self, visitor: &mut T) {
+		self.cond.accept(visitor);
+
+		visitor.visit_br_if(self);
+	}
+}
+
 impl<T: Visitor> Driver<T> for Call {
 	fn accept(&self, visitor: &mut T) {
 		for v in &self.param_list {
@@ -292,14 +312,10 @@ impl<T: Visitor> Driver<T> for StoreAt {
 impl<T: Visitor> Driver<T> for Statement {
 	fn accept(&self, visitor: &mut T) {
 		match self {
-			Self::Unreachable => visitor.visit_unreachable(),
 			Self::Forward(v) => v.accept(visitor),
 			Self::Backward(v) => v.accept(visitor),
 			Self::If(v) => v.accept(visitor),
-			Self::Br(v) => v.accept(visitor),
 			Self::BrIf(v) => v.accept(visitor),
-			Self::BrTable(v) => v.accept(visitor),
-			Self::Return(v) => v.accept(visitor),
 			Self::Call(v) => v.accept(visitor),
 			Self::CallIndirect(v) => v.accept(visitor),
 			Self::SetTemporary(v) => v.accept(visitor),