feat: use Constrain instructions to replace values with constants

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

@@ -4,6 +4,8 @@
 //! The pass works as follows:
 //! - Re-insert each instruction in order to apply the instruction simplification performed
 //!   by the [`DataFlowGraph`] automatically as new instructions are pushed.
+//! - Check whether any input values have been constrained to be equal to a value of a simpler form
+//!   by a [constrain instruction][Instruction::Constrain]. If so, replace the input value with the simpler form.
 //! - Check whether the instruction is [pure][Instruction::is_pure()]
 //!   and there exists a duplicate instruction earlier in the same block.
 //!   If so, the instruction can be replaced with the results of this previous instruction.
@@ -28,7 +30,7 @@ use crate::ssa::{
         dfg::{DataFlowGraph, InsertInstructionResult},
         function::Function,
         instruction::{Instruction, InstructionId},
-        value::ValueId,
+        value::{Value, ValueId},
     },
     ssa_gen::Ssa,
 };
@@ -75,13 +77,15 @@ impl Context {
 
         // Cache of instructions without any side-effects along with their outputs.
         let mut cached_instruction_results: HashMap<Instruction, Vec<ValueId>> = HashMap::default();
+        let mut constrained_values: HashMap<ValueId, ValueId> = HashMap::default();
 
         for instruction_id in instructions {
             Self::fold_constants_into_instruction(
                 &mut function.dfg,
                 block,
                 instruction_id,
                 &mut cached_instruction_results,
+                &mut constrained_values,
             );
         }
         self.block_queue.extend(function.dfg[block].successors());
@@ -92,8 +96,9 @@ impl Context {
         block: BasicBlockId,
         id: InstructionId,
         instruction_result_cache: &mut HashMap<Instruction, Vec<ValueId>>,
+        constrained_values: &mut HashMap<ValueId, ValueId>,
     ) {
-        let instruction = Self::resolve_instruction(id, dfg);
+        let instruction = Self::resolve_instruction(id, dfg, constrained_values);
         let old_results = dfg.instruction_results(id).to_vec();
 
         // If a copy of this instruction exists earlier in the block, then reuse the previous results.
@@ -105,20 +110,46 @@ impl Context {
         // Otherwise, try inserting the instruction again to apply any optimizations using the newly resolved inputs.
         let new_results = Self::push_instruction(id, instruction.clone(), &old_results, block, dfg);
 
-        // If the instruction is pure then we cache the results so we can reuse them if
-        // the same instruction appears again later in the block.
-        if instruction.is_pure(dfg) {
-            instruction_result_cache.insert(instruction, new_results.clone());
-        }
         Self::replace_result_ids(dfg, &old_results, &new_results);
+
+        Self::cache_instruction(
+            instruction,
+            new_results,
+            dfg,
+            instruction_result_cache,
+            constrained_values,
+        );
     }
 
     /// Fetches an [`Instruction`] by its [`InstructionId`] and fully resolves its inputs.
-    fn resolve_instruction(instruction_id: InstructionId, dfg: &DataFlowGraph) -> Instruction {
+    fn resolve_instruction(
+        instruction_id: InstructionId,
+        dfg: &DataFlowGraph,
+        constrained_values: &mut HashMap<ValueId, ValueId>,
+    ) -> Instruction {
         let instruction = dfg[instruction_id].clone();
 
+        // Alternate between resolving `value_id` in the `dfg` and checking to see if the resolved value
+        // has been constrained to be equal to some simpler value in the current block.
+        //
+        // This allows us to reach a stable final `ValueId` for each instruction input as we add more
+        // constraints to the cache.
+        fn resolve_cache(
+            dfg_resolver: impl Fn(ValueId) -> ValueId,
+            cache: &HashMap<ValueId, ValueId>,
+            value_id: ValueId,
+        ) -> ValueId {
+            let resolved_id = dfg_resolver(value_id);
+            match cache.get(&resolved_id) {
+                Some(cached_value) => resolve_cache(dfg_resolver, cache, *cached_value),
+                None => resolved_id,
+            }
+        }
+
         // Resolve any inputs to ensure that we're comparing like-for-like instructions.
-        instruction.map_values(|value_id| dfg.resolve(value_id))
+        instruction.map_values(|value_id| {
+            resolve_cache(|value_id| dfg.resolve(value_id), constrained_values, value_id)
+        })
     }
 
     /// Pushes a new [`Instruction`] into the [`DataFlowGraph`] which applies any optimizations
@@ -151,6 +182,47 @@ impl Context {
         new_results
     }
 
+    fn cache_instruction(
+        instruction: Instruction,
+        instruction_results: Vec<ValueId>,
+        dfg: &DataFlowGraph,
+        instruction_result_cache: &mut HashMap<Instruction, Vec<ValueId>>,
+        constraint_cache: &mut HashMap<ValueId, ValueId>,
+    ) {
+        // If the instruction was a constraint, then create a link between the two `ValueId`s
+        // to map from the more complex to the simpler value.
+        if let Instruction::Constrain(lhs, rhs, _) = instruction {
+            // These `ValueId`s should be fully resolved now.
+            match (&dfg[lhs], &dfg[rhs]) {
+                // Ignore trivial constraints
+                (Value::NumericConstant { .. }, Value::NumericConstant { .. }) => (),
+
+                // Prefer replacing with constants where possible.
+                (Value::NumericConstant { .. }, _) => {
+                    constraint_cache.insert(rhs, lhs);
+                }
+                (_, Value::NumericConstant { .. }) => {
+                    constraint_cache.insert(lhs, rhs);
+                }
+                // Otherwise prefer block parameters over instruction results.
+                // This is as block parameters are more likely to be a single witness rather than a full expression.
+                (Value::Param { .. }, Value::Instruction { .. }) => {
+                    constraint_cache.insert(rhs, lhs);
+                }
+                (Value::Instruction { .. }, Value::Param { .. }) => {
+                    constraint_cache.insert(lhs, rhs);
+                }
+                (_, _) => (),
+            }
+        }
+
+        // If the instruction doesn't have side-effects, cache the results so we can reuse them if
+        // the same instruction appears again later in the block.
+        if instruction.is_pure(dfg) {
+            instruction_result_cache.insert(instruction, instruction_results);
+        }
+    }
+
     /// Replaces a set of [`ValueId`]s inside the [`DataFlowGraph`] with another.
     fn replace_result_ids(
         dfg: &mut DataFlowGraph,
@@ -321,4 +393,57 @@ mod test {
 
         assert_eq!(instruction, &Instruction::Cast(ValueId::test_new(0), Type::unsigned(32)));
     }
+
+    #[test]
+    fn constrained_value_replacement() {
+        // fn main f0 {
+        //   b0(v0: Field):
+        //     constrain v0 10
+        //     v1 = add v0, Field 1
+        //     constrain v1 11
+        // }
+        //
+        // After constructing this IR, we run constant folding which should replace the second cast
+        // with a reference to the results to the first. This then allows us to optimize away
+        // the constrain instruction as both inputs are known to be equal.
+        //
+        // The first cast instruction is 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::field());
+
+        let field_10 = builder.field_constant(10u128);
+        builder.insert_constrain(v0, field_10, None);
+
+        let field_1 = builder.field_constant(1u128);
+        let v1 = builder.insert_binary(v0, BinaryOp::Add, field_1);
+
+        let field_11 = builder.field_constant(11u128);
+        builder.insert_constrain(v1, field_11, None);
+
+        let mut ssa = builder.finish();
+        let main = ssa.main_mut();
+        let instructions = main.dfg[main.entry_block()].instructions();
+        assert_eq!(instructions.len(), 3);
+
+        // Expected output:
+        //
+        // fn main f0 {
+        //   b0(v0: Field):
+        //     constrain v0 10
+        // }
+        let ssa = ssa.fold_constants();
+        let main = ssa.main();
+        let instructions = main.dfg[main.entry_block()].instructions();
+
+        assert_eq!(instructions.len(), 1);
+        let instruction = &main.dfg[instructions[0]];
+
+        assert_eq!(
+            instruction,
+            &Instruction::Constrain(ValueId::test_new(0), ValueId::test_new(1), None)
+        );
+    }
 }