Ensure br is only called with i1 (#45)

Veridise · Aug 16, 2023 · ca4f00c · ca4f00c
1 parent f70978d
commit ca4f00c
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 1 deletion.
diff --git a/circom/tests/controlflow/van544.circom 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
@@ -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()
 }