Fix alignment when branching

codegen-luajit/src/backend/statement.rs

@@ -5,7 +5,7 @@ use std::{
 
 use parity_wasm::elements::ValueType;
 use wasm_ast::node::{
-	Backward, Br, BrTable, Call, CallIndirect, Forward, FuncData, If, SetGlobal, SetLocal,
+	Backward, Br, BrIf, BrTable, Call, CallIndirect, Forward, FuncData, If, SetGlobal, SetLocal,
 	SetTemporary, Statement, StoreAt, Terminator,
 };
 
@@ -13,66 +13,40 @@ use super::manager::{
 	write_ascending, write_condition, write_separated, write_variable, Driver, Manager,
 };
 
-fn write_br_at(up: usize, mng: &Manager, w: &mut dyn Write) -> Result<()> {
-	let level = mng.label_list().iter().nth_back(up).unwrap();
-
-	write!(w, "goto continue_at_{level} ")
-}
-
 impl Driver for Br {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-		write_br_at(self.target, mng, w)
-	}
-}
+		let level = *mng.label_list().iter().nth_back(self.target).unwrap();
 
-fn condense_jump_table(list: &[u32]) -> Vec<(usize, usize, u32)> {
-	let mut result = Vec::new();
-	let mut index = 0;
-
-	while index < list.len() {
-		let start = index;
-
-		loop {
-			index += 1;
-
-			// if end of list or next value is not equal, break
-			if index == list.len() || list[index - 1] != list[index] {
-				break;
-			}
+		if !self.align.is_aligned() {
+			write_ascending("reg", self.align.new_range(), w)?;
+			write!(w, " = ")?;
+			write_ascending("reg", self.align.old_range(), w)?;
+			write!(w, " ")?;
 		}
 
-		result.push((start, index - 1, list[start]));
+		write!(w, "goto continue_at_{level} ")
 	}
-
-	result
 }
 
 impl Driver for BrTable {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-		let default = self.data.default.try_into().unwrap();
-
-		// Our condition should be pure so we probably don't need
-		// to emit it in this case.
-		if self.data.table.is_empty() {
-			return write_br_at(default, mng, w);
-		}
-
 		write!(w, "temp = ")?;
 		self.cond.write(mng, w)?;
 
-		for (start, end, dest) in condense_jump_table(&self.data.table) {
-			if start == end {
-				write!(w, "if temp == {start} then ")?;
-			} else {
-				write!(w, "if temp >= {start} and temp <= {end} then ")?;
-			}
+		// Our condition should be pure so we probably don't need
+		// to emit it in this case.
+		if self.data.is_empty() {
+			return self.default.write(mng, w);
+		}
 
-			write_br_at(dest.try_into().unwrap(), mng, w)?;
+		for (case, dest) in self.data.iter().enumerate() {
+			write!(w, "if temp == {case} then ")?;
+			dest.write(mng, w)?;
 			write!(w, "else")?;
 		}
 
 		write!(w, " ")?;
-		write_br_at(default, mng, w)?;
+		self.default.write(mng, w)?;
 		write!(w, "end ")
 	}
 }
@@ -127,6 +101,16 @@ impl Driver for Backward {
 	}
 }
 
+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 ")?;
+		self.target.write(mng, w)?;
+		write!(w, "end ")
+	}
+}
+
 impl Driver for If {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 		write!(w, "if ")?;
@@ -220,6 +204,7 @@ impl Driver for Statement {
 		match self {
 			Self::Forward(s) => s.write(mng, w),
 			Self::Backward(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),
 			Self::CallIndirect(s) => s.write(mng, w),