From 8792a123f4d6dce4ec21a118aff7abe617de6a3d Mon Sep 17 00:00:00 2001
From: Rerumu <rerumu@protonmail.com>
Date: Thu, 19 May 2022 12:15:42 -0400
Subject: [PATCH] Fix and optimize conditionals

---
 codegen-luau/src/backend/expression.rs | 28 ++++++++-----------------
 codegen-luau/src/backend/manager.rs    | 29 ++++++++++++++++++++++++++
 codegen-luau/src/backend/statement.rs  | 12 ++++++-----
 3 files changed, 45 insertions(+), 24 deletions(-)

diff --git a/codegen-luau/src/backend/expression.rs b/codegen-luau/src/backend/expression.rs
index 14aee64..2ec1e06 100644
--- a/codegen-luau/src/backend/expression.rs
+++ b/codegen-luau/src/backend/expression.rs
@@ -5,9 +5,11 @@ use wasm_ast::node::{
 	UnOp, Value,
 };
 
-use crate::analyzer::operator::{bin_symbol_of, cmp_symbol_of};
+use crate::analyzer::operator::bin_symbol_of;
 
-use super::manager::{write_separated, write_variable, Driver, Manager};
+use super::manager::{
+	write_cmp_op, write_condition, write_separated, write_variable, Driver, Manager,
+};
 
 impl Driver for Recall {
 	fn write(&self, _: &mut Manager, w: &mut dyn Write) -> Result<()> {
@@ -18,8 +20,8 @@ impl Driver for Recall {
 impl Driver for Select {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 		write!(w, "(")?;
-		self.cond.write(mng, w)?;
-		write!(w, "~= 0 and ")?;
+		write_condition(&self.cond, mng, w)?;
+		write!(w, "and ")?;
 		self.a.write(mng, w)?;
 		write!(w, "or ")?;
 		self.b.write(mng, w)?;
@@ -153,21 +155,9 @@ impl Driver for BinOp {
 
 impl Driver for CmpOp {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-		if let Some(symbol) = cmp_symbol_of(self.op) {
-			write!(w, "(")?;
-			self.lhs.write(mng, w)?;
-			write!(w, "{symbol} ")?;
-			self.rhs.write(mng, w)?;
-			write!(w, ")")
-		} else {
-			let (head, tail) = self.op.as_name();
-
-			write!(w, "{head}_{tail}(")?;
-			self.lhs.write(mng, w)?;
-			write!(w, ", ")?;
-			self.rhs.write(mng, w)?;
-			write!(w, ")")
-		}
+		write!(w, "(")?;
+		write_cmp_op(self, mng, w)?;
+		write!(w, "and 1 or 0)")
 	}
 }
 
diff --git a/codegen-luau/src/backend/manager.rs b/codegen-luau/src/backend/manager.rs
index 5cf3990..e184d17 100644
--- a/codegen-luau/src/backend/manager.rs
+++ b/codegen-luau/src/backend/manager.rs
@@ -3,6 +3,10 @@ use std::{
 	ops::Range,
 };
 
+use wasm_ast::node::{CmpOp, Expression};
+
+use crate::analyzer::operator::cmp_symbol_of;
+
 #[derive(PartialEq, Eq)]
 pub enum Label {
 	Forward,
@@ -63,3 +67,28 @@ pub fn write_variable(var: usize, mng: &Manager, w: &mut dyn Write) -> Result<()
 		write!(w, "param_{var} ")
 	}
 }
+
+pub fn write_cmp_op(cmp: &CmpOp, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+	if let Some(symbol) = cmp_symbol_of(cmp.op) {
+		cmp.lhs.write(mng, w)?;
+		write!(w, "{symbol} ")?;
+		cmp.rhs.write(mng, w)
+	} else {
+		let (head, tail) = cmp.op.as_name();
+
+		write!(w, "{head}_{tail}(")?;
+		cmp.lhs.write(mng, w)?;
+		write!(w, ", ")?;
+		cmp.rhs.write(mng, w)?;
+		write!(w, ")")
+	}
+}
+
+pub fn write_condition(data: &Expression, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+	if let Expression::CmpOp(node) = data {
+		write_cmp_op(node, mng, w)
+	} else {
+		data.write(mng, w)?;
+		write!(w, "~= 0 ")
+	}
+}
diff --git a/codegen-luau/src/backend/statement.rs b/codegen-luau/src/backend/statement.rs
index 45f914e..3c83be0 100644
--- a/codegen-luau/src/backend/statement.rs
+++ b/codegen-luau/src/backend/statement.rs
@@ -9,7 +9,9 @@ use wasm_ast::node::{
 	Return, SetGlobal, SetLocal, Statement, StoreAt,
 };
 
-use super::manager::{write_ascending, write_separated, write_variable, Driver, Label, Manager};
+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 ")?;
@@ -86,8 +88,8 @@ impl Driver for If {
 
 		write!(w, "while true do ")?;
 		write!(w, "if ")?;
-		self.cond.write(mng, w)?;
-		write!(w, "~= 0 then ")?;
+		write_condition(&self.cond, mng, w)?;
+		write!(w, "then ")?;
 
 		self.truthy.iter().try_for_each(|s| s.write(mng, w))?;
 
@@ -132,8 +134,8 @@ impl Driver for Br {
 impl Driver for BrIf {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 		write!(w, "if ")?;
-		self.cond.write(mng, w)?;
-		write!(w, "~= 0 then ")?;
+		write_condition(&self.cond, mng, w)?;
+		write!(w, "then ")?;
 		write_br_at(self.target, mng, w)?;
 		write!(w, "end ")
 	}