Improve if statement and top-level jump handling

src/backend/ast/data.rs

@@ -212,8 +212,8 @@ impl Backward {
 
 pub struct If {
 	pub cond: Expression,
-	pub body: Vec<Statement>,
-	pub other: Option<Vec<Statement>>,
+	pub truthy: Vec<Statement>,
+	pub falsey: Option<Vec<Statement>>,
 }
 
 impl If {
@@ -222,11 +222,11 @@ impl If {
 
 		self.cond.accept(visitor);
 
-		for v in &self.body {
+		for v in &self.truthy {
 			v.accept(visitor);
 		}
 
-		if let Some(v) = &self.other {
+		if let Some(v) = &self.falsey {
 			for v in v {
 				v.accept(visitor);
 			}
@@ -403,13 +403,11 @@ pub struct Function {
 	pub num_param: u32,
 	pub num_local: u32,
 	pub num_stack: u32,
-	pub body: Vec<Statement>,
+	pub body: Forward,
 }
 
 impl Function {
 	pub fn accept<V: Visitor>(&self, visitor: &mut V) {
-		for v in &self.body {
-			v.accept(visitor);
-		}
+		self.body.accept(visitor);
 	}
 }

src/backend/ast/transformer.rs

@@ -31,8 +31,8 @@ fn local_sum(list: &[Local]) -> u32 {
 	list.iter().map(Local::count).sum()
 }
 
-fn is_else_stat(list: &[Instruction]) -> bool {
-	list.get(0) == Some(&Instruction::Else)
+fn is_else_stat(inst: &Instruction) -> bool {
+	inst == &Instruction::Else
 }
 
 impl<'a> Transformer<'a> {
@@ -51,7 +51,7 @@ impl<'a> Transformer<'a> {
 		debug_assert!(self.name != usize::MAX, "Not an indexed value");
 
 		let func = &self.wasm.code_section().unwrap().bodies()[self.name];
-		let body = self.new_stored_body(&mut func.code().elements());
+		let body = self.new_forward(&mut func.code().elements());
 
 		Function {
 			num_param: self.arity.in_arity[self.name].num_param,
@@ -205,7 +205,9 @@ impl<'a> Transformer<'a> {
 
 		let mut stat = Vec::new();
 
-		while let Some(inst) = list.get(0) {
+		loop {
+			let inst = &list[0];
+
 			*list = &list[1..];
 
 			if let Ok(op) = UnOp::try_from(inst) {
@@ -384,13 +386,17 @@ impl<'a> Transformer<'a> {
 	}
 
 	fn new_if(&mut self, cond: Expression, list: &mut &[Instruction]) -> If {
-		let body = self.new_stored_body(list);
-		let other = is_else_stat(list).then(|| {
-			*list = &list[1..];
-			self.new_stored_body(list)
-		});
+		let copied = list.clone();
+		let truthy = self.new_stored_body(list);
 
-		If { cond, body, other }
+		let last = copied.len() - list.len() - 1;
+		let falsey = is_else_stat(&copied[last]).then(|| self.new_stored_body(list));
+
+		If {
+			cond,
+			truthy,
+			falsey,
+		}
 	}
 
 	fn new_backward(&mut self, list: &mut &[Instruction]) -> Backward {

src/backend/translator/writer.rs

@@ -271,7 +271,11 @@ impl If {
 
 		d.edition.start_if(&var, w)?;
 
-		self.body.iter().try_for_each(|s| s.output(d, w))?;
+		self.truthy.iter().try_for_each(|s| s.output(d, w))?;
+
+		if let Some(v) = &self.falsey {
+			v.iter().try_for_each(|s| s.output(d, w))?;
+		}
 
 		d.edition.end_if(rem, w)?;
 		d.label_list.pop().unwrap();
@@ -478,7 +482,7 @@ impl Function {
 		self.write_visitor_list(w)?;
 		self.write_variable_list(w)?;
 
-		self.body.iter().try_for_each(|s| s.output(d, w))?;
+		self.body.output(d, w)?;
 
 		write!(w, "end ")
 	}