feat: decompose Instruction::Constrain into multiple more basic constraints

compiler/noirc_evaluator/src/ssa/ir/dfg.rs

@@ -166,11 +166,31 @@ impl DataFlowGraph {
                 SimplifiedToMultiple(simplification)
             }
             SimplifyResult::Remove => InstructionRemoved,
-            result @ (SimplifyResult::SimplifiedToInstruction(_) | SimplifyResult::None) => {
-                let instruction = result.instruction().unwrap_or(instruction);
-                let id = self.make_instruction(instruction, ctrl_typevars);
-                self.blocks[block].insert_instruction(id);
-                self.locations.insert(id, call_stack);
+            result @ (SimplifyResult::SimplifiedToInstruction(_)
+            | SimplifyResult::SimplifiedToInstructionMultiple(_)
+            | SimplifyResult::None) => {
+                let instructions = result.instructions().unwrap_or(vec![instruction]);
+
+                if instructions.len() > 1 {
+                    // There's currently no way to pass results from one instruction in `instructions` on to the next.
+                    // We then restrict this to only support multiple instructions if they're all `Instruction::Constrain`
+                    // as this instruction type does not have any results.
+                    assert!(
+                        instructions.iter().all(|instruction| matches!(instruction, Instruction::Constrain(..))),
+                        "`SimplifyResult::SimplifiedToInstructionMultiple` only supports `Constrain` instructions"
+                    );
+                }
+
+                let mut last_id = None;
+
+                for instruction in instructions {
+                    let id = self.make_instruction(instruction, ctrl_typevars.clone());
+                    self.blocks[block].insert_instruction(id);
+                    self.locations.insert(id, call_stack.clone());
+                    last_id = Some(id);
+                }
+
+                let id = last_id.expect("There should be at least 1 simplified instruction");
                 InsertInstructionResult::Results(id, self.instruction_results(id))
             }
         }

compiler/noirc_evaluator/src/ssa/ir/instruction.rs

@@ -432,73 +432,11 @@ impl Instruction {
                 }
             }
             Instruction::Constrain(lhs, rhs, msg) => {
-                if dfg.resolve(*lhs) == dfg.resolve(*rhs) {
-                    // Remove trivial case `assert_eq(x, x)`
-                    SimplifyResult::Remove
+                let constraints = decompose_constrain(*lhs, *rhs, msg.clone(), dfg);
+                if constraints.is_empty() {
+                    Remove
                 } else {
-                    match (&dfg[dfg.resolve(*lhs)], &dfg[dfg.resolve(*rhs)]) {
-                        (
-                            Value::NumericConstant { constant, typ },
-                            Value::Instruction { instruction, .. },
-                        )
-                        | (
-                            Value::Instruction { instruction, .. },
-                            Value::NumericConstant { constant, typ },
-                        ) if *typ == Type::bool() => {
-                            match dfg[*instruction] {
-                                Instruction::Binary(Binary {
-                                    lhs,
-                                    rhs,
-                                    operator: BinaryOp::Eq,
-                                }) if constant.is_one() => {
-                                    // Replace an explicit two step equality assertion
-                                    //
-                                    // v2 = eq v0, u32 v1
-                                    // constrain v2 == u1 1
-                                    //
-                                    // with a direct assertion of equality between the two values
-                                    //
-                                    // v2 = eq v0, u32 v1
-                                    // constrain v0 == v1
-                                    //
-                                    // Note that this doesn't remove the value `v2` as it may be used in other instructions, but it
-                                    // will likely be removed through dead instruction elimination.
-
-                                    SimplifiedToInstruction(Instruction::Constrain(
-                                        lhs,
-                                        rhs,
-                                        msg.clone(),
-                                    ))
-                                }
-                                Instruction::Not(value) => {
-                                    // Replace an assertion that a not instruction is truthy
-                                    //
-                                    // v1 = not v0
-                                    // constrain v1 == u1 1
-                                    //
-                                    // with an assertion that the not instruction input is falsy
-                                    //
-                                    // v1 = not v0
-                                    // constrain v0 == u1 0
-                                    //
-                                    // Note that this doesn't remove the value `v1` as it may be used in other instructions, but it
-                                    // will likely be removed through dead instruction elimination.
-                                    let reversed_constant = FieldElement::from(!constant.is_one());
-                                    let reversed_constant =
-                                        dfg.make_constant(reversed_constant, Type::bool());
-                                    SimplifiedToInstruction(Instruction::Constrain(
-                                        value,
-                                        reversed_constant,
-                                        msg.clone(),
-                                    ))
-                                }
-
-                                _ => None,
-                            }
-                        }
-
-                        _ => None,
-                    }
+                    SimplifiedToInstructionMultiple(constraints)
                 }
             }
             Instruction::ArrayGet { array, index } => {
@@ -646,6 +584,129 @@ fn simplify_cast(value: ValueId, dst_typ: &Type, dfg: &mut DataFlowGraph) -> Sim
     }
 }
 
+/// Try to decompose this constrain instruction. This constraint will be broken down such that it instead constrains
+/// all the values which are used to compute the values which were being constrained.
+fn decompose_constrain(
+    lhs: ValueId,
+    rhs: ValueId,
+    msg: Option<String>,
+    dfg: &mut DataFlowGraph,
+) -> Vec<Instruction> {
+    let lhs = dfg.resolve(lhs);
+    let rhs = dfg.resolve(rhs);
+
+    if lhs == rhs {
+        // Remove trivial case `assert_eq(x, x)`
+        Vec::new()
+    } else {
+        match (&dfg[lhs], &dfg[rhs]) {
+            (Value::NumericConstant { constant, typ }, Value::Instruction { instruction, .. })
+            | (Value::Instruction { instruction, .. }, Value::NumericConstant { constant, typ })
+                if *typ == Type::bool() =>
+            {
+                match dfg[*instruction] {
+                    Instruction::Binary(Binary { lhs, rhs, operator: BinaryOp::Eq })
+                        if constant.is_one() =>
+                    {
+                        // Replace an explicit two step equality assertion
+                        //
+                        // v2 = eq v0, u32 v1
+                        // constrain v2 == u1 1
+                        //
+                        // with a direct assertion of equality between the two values
+                        //
+                        // v2 = eq v0, u32 v1
+                        // constrain v0 == v1
+                        //
+                        // Note that this doesn't remove the value `v2` as it may be used in other instructions, but it
+                        // will likely be removed through dead instruction elimination.
+
+                        vec![Instruction::Constrain(lhs, rhs, msg)]
+                    }
+
+                    Instruction::Binary(Binary { lhs, rhs, operator: BinaryOp::Mul })
+                        if constant.is_one() && dfg.type_of_value(lhs) == Type::bool() =>
+                    {
+                        // Replace an equality assertion on a boolean multiplication
+                        //
+                        // v2 = mul v0, v1
+                        // constrain v2 == u1 1
+                        //
+                        // with a direct assertion that each value is equal to 1
+                        //
+                        // v2 = mul v0, v1
+                        // constrain v0 == 1
+                        // constrain v1 == 1
+                        //
+                        // This is due to the fact that for `v2` to be 1 then both `v0` and `v1` are 1.
+                        //
+                        // Note that this doesn't remove the value `v2` as it may be used in other instructions, but it
+                        // will likely be removed through dead instruction elimination.
+                        let one = FieldElement::one();
+                        let one = dfg.make_constant(one, Type::bool());
+
+                        [
+                            decompose_constrain(lhs, one, msg.clone(), dfg),
+                            decompose_constrain(rhs, one, msg, dfg),
+                        ]
+                        .concat()
+                    }
+
+                    Instruction::Binary(Binary { lhs, rhs, operator: BinaryOp::Or })
+                        if constant.is_zero() =>
+                    {
+                        // Replace an equality assertion on an OR
+                        //
+                        // v2 = or v0, v1
+                        // constrain v2 == u1 0
+                        //
+                        // with a direct assertion that each value is equal to 0
+                        //
+                        // v2 = or v0, v1
+                        // constrain v0 == 0
+                        // constrain v1 == 0
+                        //
+                        // This is due to the fact that for `v2` to be 0 then both `v0` and `v1` are 0.
+                        //
+                        // Note that this doesn't remove the value `v2` as it may be used in other instructions, but it
+                        // will likely be removed through dead instruction elimination.
+                        let zero = FieldElement::zero();
+                        let zero = dfg.make_constant(zero, dfg.type_of_value(lhs));
+
+                        [
+                            decompose_constrain(lhs, zero, msg.clone(), dfg),
+                            decompose_constrain(rhs, zero, msg, dfg),
+                        ]
+                        .concat()
+                    }
+
+                    Instruction::Not(value) => {
+                        // Replace an assertion that a not instruction is truthy
+                        //
+                        // v1 = not v0
+                        // constrain v1 == u1 1
+                        //
+                        // with an assertion that the not instruction input is falsy
+                        //
+                        // v1 = not v0
+                        // constrain v0 == u1 0
+                        //
+                        // Note that this doesn't remove the value `v1` as it may be used in other instructions, but it
+                        // will likely be removed through dead instruction elimination.
+                        let reversed_constant = FieldElement::from(!constant.is_one());
+                        let reversed_constant = dfg.make_constant(reversed_constant, Type::bool());
+                        decompose_constrain(value, reversed_constant, msg, dfg)
+                    }
+
+                    _ => vec![Instruction::Constrain(lhs, rhs, msg)],
+                }
+            }
+
+            _ => vec![Instruction::Constrain(lhs, rhs, msg)],
+        }
+    }
+}
+
 /// The possible return values for Instruction::return_types
 pub(crate) enum InstructionResultType {
     /// The result type of this instruction matches that of this operand
@@ -1134,6 +1195,10 @@ pub(crate) enum SimplifyResult {
     /// Replace this function with an simpler but equivalent instruction.
     SimplifiedToInstruction(Instruction),
 
+    /// Replace this function with a set of simpler but equivalent instructions.
+    /// This is currently only to be used for [`Instruction::Constrain`].
+    SimplifiedToInstructionMultiple(Vec<Instruction>),
+
     /// Remove the instruction, it is unnecessary
     Remove,
 
@@ -1142,9 +1207,10 @@ pub(crate) enum SimplifyResult {
 }
 
 impl SimplifyResult {
-    pub(crate) fn instruction(self) -> Option<Instruction> {
+    pub(crate) fn instructions(self) -> Option<Vec<Instruction>> {
         match self {
-            SimplifyResult::SimplifiedToInstruction(instruction) => Some(instruction),
+            SimplifyResult::SimplifiedToInstruction(instruction) => Some(vec![instruction]),
+            SimplifyResult::SimplifiedToInstructionMultiple(instructions) => Some(instructions),
             _ => None,
         }
     }

compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs

@@ -445,4 +445,67 @@ mod test {
 
         assert_eq!(instruction, &Instruction::Cast(v0, Type::unsigned(32)));
     }
+
+    #[test]
+    fn constraint_decomposition() {
+        // fn main f0 {
+        //   b0(v0: u1, v1: u1, v2: u1):
+        //     v3 = mul v0 v1
+        //     v4 = not v2
+        //     v5 = mul v3 v4
+        //     constrain v4 u1 1
+        // }
+        //
+        // When constructing this IR, we should automatically decompose the constraint to be in terms of `v0`, `v1` and `v2`.
+        //
+        // The mul instructions are retained and will be removed in the dead instruction elimination pass.
+        let main_id = Id::test_new(0);
+
+        // Compiling main
+        let mut builder = FunctionBuilder::new("main".into(), main_id, RuntimeType::Acir);
+        let v0 = builder.add_parameter(Type::bool());
+        let v1 = builder.add_parameter(Type::bool());
+        let v2 = builder.add_parameter(Type::bool());
+
+        let v3 = builder.insert_binary(v0, BinaryOp::Mul, v1);
+        let v4 = builder.insert_not(v2);
+        let v5 = builder.insert_binary(v3, BinaryOp::Mul, v4);
+
+        // This constraint is automatically decomposed when it is inserted.
+        let v_true = builder.numeric_constant(true, Type::bool());
+        builder.insert_constrain(v5, v_true, None);
+
+        let v_false = builder.numeric_constant(false, Type::bool());
+
+        // Expected output:
+        //
+        // fn main f0 {
+        //   b0(v0: u1, v1: u1, v2: u1):
+        //     v3 = mul v0 v1
+        //     v4 = not v2
+        //     v5 = mul v3 v4
+        //     constrain v0 u1 1
+        //     constrain v1 u1 1
+        //     constrain v2 u1 0
+        // }
+
+        let ssa = builder.finish();
+        let main = ssa.main();
+        let instructions = main.dfg[main.entry_block()].instructions();
+
+        assert_eq!(instructions.len(), 6);
+
+        assert_eq!(
+            main.dfg[instructions[0]],
+            Instruction::Binary(Binary { lhs: v0, operator: BinaryOp::Mul, rhs: v1 })
+        );
+        assert_eq!(main.dfg[instructions[1]], Instruction::Not(v2));
+        assert_eq!(
+            main.dfg[instructions[2]],
+            Instruction::Binary(Binary { lhs: v3, operator: BinaryOp::Mul, rhs: v4 })
+        );
+        assert_eq!(main.dfg[instructions[3]], Instruction::Constrain(v0, v_true, None));
+        assert_eq!(main.dfg[instructions[4]], Instruction::Constrain(v1, v_true, None));
+        assert_eq!(main.dfg[instructions[5]], Instruction::Constrain(v2, v_false, None));
+    }
 }