feat: conditional select

1 files changed, 96 insertions, 0 deletions

diff --git a/src/codegen.c b/src/codegen.c
index 271c5ae..c0e67be 100644
--- a/src/codegen.c
+++ b/src/codegen.c
@@ -1162,6 +1162,55 @@ static void emit_instruction(struct Emitter* emitter, struct InstructionStatemen
 	fprintf(emitter->out, "\n");
 }
 
+// the x86 condition-move suffix for a comparison, true or inverted
+static const char* cmov_cc(enum TokenType comparison, bool when_true)
+{
+	switch (comparison)
+	{
+		case TOKEN_EQUAL_EQUAL:   return when_true ? "e"  : "ne";
+		case TOKEN_BANG_EQUAL:    return when_true ? "ne" : "e";
+		case TOKEN_LESS:          return when_true ? "l"  : "ge";
+		case TOKEN_LESS_EQUAL:    return when_true ? "le" : "g";
+		case TOKEN_GREATER:       return when_true ? "g"  : "le";
+		case TOKEN_GREATER_EQUAL: return when_true ? "ge" : "l";
+		default:                  return NULL;
+	}
+}
+
+// target = if_value if L cmp R else else_value, branchlessly via cmov. The
+// destination defaults to one source and conditionally takes the other, choosing
+// which to default so the target can alias either source safely.
+static void emit_select(struct Emitter* emitter, struct SelectStatement* select)
+{
+	FILE* out = emitter->out;
+	if (select->if_value->kind != EXPR_PRIMARY || select->else_value->kind != EXPR_PRIMARY
+		|| cmov_cc(select->comparison.type, true) == NULL)
+	{
+		fprintf(out, "\t; TODO: unsupported select\n");
+		return;
+	}
+
+	struct Token dst = resolve_register(emitter, select->target);
+	struct Token a = resolve_register(emitter, select->if_value->primary.token);
+	struct Token b = resolve_register(emitter, select->else_value->primary.token);
+	bool dst_is_a = tokens_equal(dst, a);
+
+	if (!dst_is_a && !tokens_equal(dst, b))
+		fprintf(out, "\tmov %.*s, %.*s\n", (int)dst.length, dst.start, (int)b.length, b.start);
+
+	fprintf(out, "\tcmp ");
+	emit_operand(emitter, select->left);
+	fprintf(out, ", ");
+	emit_operand(emitter, select->right);
+	fprintf(out, "\n");
+
+	// if the destination already holds if_value, keep it when the condition is
+	// false and move else_value in otherwise; otherwise the normal true form
+	struct Token source = dst_is_a ? b : a;
+	fprintf(out, "\tcmov%s %.*s, %.*s\n", cmov_cc(select->comparison.type, !dst_is_a),
+		(int)dst.length, dst.start, (int)source.length, source.start);
+}
+
 static void emit_statement(struct Emitter* emitter, struct Statement* statement)
 {
 	FILE* out = emitter->out;
@@ -1170,6 +1219,9 @@ static void emit_statement(struct Emitter* emitter, struct Statement* statement)
 		case STATEMENT_ASSIGN:
 			emit_assign(emitter, &statement->assign);
 			break;
+		case STATEMENT_SELECT:
+			emit_select(emitter, &statement->select);
+			break;
 		case STATEMENT_LABEL:
 			fprintf(out, "%.*s:\n", (int)statement->label.name.length, statement->label.name.start);
 			break;
@@ -1641,6 +1693,47 @@ static void emit_a64_instruction(struct Emitter* emitter, struct InstructionStat
 	fprintf(emitter->out, "\n");
 }
 
+// the AArch64 condition code for a true comparison (used by csel)
+static const char* a64_cond(enum TokenType comparison)
+{
+	switch (comparison)
+	{
+		case TOKEN_EQUAL_EQUAL:   return "eq";
+		case TOKEN_BANG_EQUAL:    return "ne";
+		case TOKEN_LESS:          return "lt";
+		case TOKEN_LESS_EQUAL:    return "le";
+		case TOKEN_GREATER:       return "gt";
+		case TOKEN_GREATER_EQUAL: return "ge";
+		default:                  return NULL;
+	}
+}
+
+// target = if_value if L cmp R else else_value  ->  cmp; csel (one instruction,
+// no aliasing hazard)
+static void emit_a64_select(struct Emitter* emitter, struct SelectStatement* select)
+{
+	FILE* out = emitter->out;
+	const char* cond = a64_cond(select->comparison.type);
+	if (cond == NULL || select->if_value->kind != EXPR_PRIMARY
+		|| select->else_value->kind != EXPR_PRIMARY)
+	{
+		fprintf(out, "\t; TODO: unsupported select\n");
+		return;
+	}
+
+	fprintf(out, "\tcmp ");
+	emit_a64_operand(emitter, select->left);
+	fprintf(out, ", ");
+	emit_a64_operand(emitter, select->right);
+	fprintf(out, "\n\tcsel ");
+	emit_a64_reg(emitter, select->target);
+	fprintf(out, ", ");
+	emit_a64_reg(emitter, select->if_value->primary.token);
+	fprintf(out, ", ");
+	emit_a64_reg(emitter, select->else_value->primary.token);
+	fprintf(out, ", %s\n", cond);
+}
+
 static void emit_a64_statement(struct Emitter* emitter, struct Statement* statement)
 {
 	FILE* out = emitter->out;
@@ -1649,6 +1742,9 @@ static void emit_a64_statement(struct Emitter* emitter, struct Statement* statem
 		case STATEMENT_ASSIGN:
 			emit_a64_assign(emitter, &statement->assign);
 			break;
+		case STATEMENT_SELECT:
+			emit_a64_select(emitter, &statement->select);
+			break;
 		case STATEMENT_LABEL:
 			fprintf(out, "%.*s:\n", (int)statement->label.name.length, statement->label.name.start);
 			break;