From ca4f00c79a6cf3cc0f7758f65206c5eb3ece5e33 Mon Sep 17 00:00:00 2001
From: Ian Neal <ian@veridise.com>
Date: Wed, 16 Aug 2023 16:51:22 -0400
Subject: [PATCH] Ensure br is only called with i1 (#45)

---
 circom/tests/controlflow/van544.circom        | 22 +++++++++++++++++++
 .../src/llvm_elements/instructions.rs         |  9 +++++++-
 2 files changed, 30 insertions(+), 1 deletion(-)
 create mode 100644 circom/tests/controlflow/van544.circom

diff --git a/circom/tests/controlflow/van544.circom b/circom/tests/controlflow/van544.circom
new file mode 100644
index 000000000..40d4acc70
--- /dev/null
+++ b/circom/tests/controlflow/van544.circom
@@ -0,0 +1,22 @@
+pragma circom 2.0.0;
+// REQUIRES: circom
+// RUN: rm -rf %t && mkdir %t && %circom --llvm -o %t %s | sed -n 's/.*Written successfully:.* \(.*\)/\1/p' | xargs cat | FileCheck %s
+
+template Conditional() {
+    signal input inp;
+    var q;
+    if (inp) {
+        q = 0;
+    } else {
+        q = 1;
+    }
+}
+
+component main = Conditional();
+
+//CHECK-LABEL: define void @Conditional_{{[0-9]+}}_run
+//CHECK-SAME: ([0 x i256]* %0)
+//CHECK: %[[INP_PTR:.*]] = getelementptr [0 x i256], [0 x i256]* %0, i32 0, i32 0
+//CHECK: %[[INP:.*]] = load i256, i256* %2
+//CHECK: %[[COND:.*]] = icmp ne i256 %[[INP]], 0
+//CHECK: br i1 %[[COND]], label %if.then, label %if.else
diff --git a/code_producers/src/llvm_elements/instructions.rs b/code_producers/src/llvm_elements/instructions.rs
index a2c05c06f..d64e1e378 100644
--- a/code_producers/src/llvm_elements/instructions.rs
+++ b/code_producers/src/llvm_elements/instructions.rs
@@ -548,10 +548,17 @@ pub fn create_conditional_branch<'a>(
     then_block: BasicBlock<'a>,
     else_block: BasicBlock<'a>,
 ) -> AnyValueEnum<'a> {
+    let comparison_type = comparison.get_type();
+    let bool_ty = producer.llvm().module.get_context().bool_type();
+    let bool_comparison = if comparison_type != bool_ty {
+        create_neq(producer, comparison, comparison_type.const_zero()).into_int_value()
+    } else {
+        comparison
+    };
     producer
         .llvm()
         .builder
-        .build_conditional_branch(comparison, then_block, else_block)
+        .build_conditional_branch(bool_comparison, then_block, else_block)
         .as_any_value_enum()
 }